add support for IndexRange and PartitionBlock

Author: nabinchha

src/data_designer/config/seed.py

@@ -19,26 +19,41 @@ class SamplingStrategy(str, Enum):
 
 
 class IndexRange(ConfigBase):
-    start: int = Field(..., ge=0)
-    end: int = Field(..., ge=1)
+    start: int = Field(ge=0, description="The start index of the index range (inclusive)")
+    end: int = Field(ge=0, description="The end index of the index range (inclusive)")
 
     @model_validator(mode="after")
     def _validate_index_range(self) -> Self:
-        if self.start >= self.end:
-            raise ValueError("'start' index must be less than 'end' index")
+        if self.start > self.end:
+            raise ValueError("'start' index must be less than or equal to 'end' index")
         return self
 
+    @property
+    def size(self) -> int:
+        return self.end - self.start + 1
+
 
 class PartitionBlock(ConfigBase):
-    partition_index: int = Field(..., default=0, ge=0)
-    num_partitions: int = Field(..., default=1, ge=1)
+    partition_index: int = Field(default=0, ge=0, description="The index of the partition to sample from")
+    num_partitions: int = Field(default=1, ge=1, description="The total number of partitions in the dataset")
 
     @model_validator(mode="after")
     def _validate_partition_block(self) -> Self:
         if self.partition_index >= self.num_partitions:
             raise ValueError("'partition_index' must be less than 'num_partitions'")
         return self
 
+    def to_index_range(self, dataset_size: int) -> IndexRange:
+        partition_size = dataset_size // self.num_partitions
+        start = self.partition_index * partition_size
+
+        # For the last partition, extend to the end of the dataset to include remainder rows
+        if self.partition_index == self.num_partitions - 1:
+            end = dataset_size - 1
+        else:
+            end = ((self.partition_index + 1) * partition_size) - 1
+        return IndexRange(start=start, end=end)
+
 
 class SeedConfig(ConfigBase):
     dataset: str

src/data_designer/engine/column_generators/generators/seed_dataset.py

@@ -7,12 +7,13 @@
 import duckdb
 import pandas as pd
 
-from data_designer.config.seed import SamplingStrategy
+from data_designer.config.seed import SamplingStrategy, IndexRange, PartitionBlock
 from data_designer.engine.column_generators.generators.base import (
     FromScratchColumnGenerator,
     GenerationStrategy,
     GeneratorMetadata,
 )
+from data_designer.engine.column_generators.utils.errors import SeedDatasetError
 from data_designer.engine.dataset_builders.multi_column_configs import SeedDatasetMultiColumnConfig
 from data_designer.engine.processing.utils import concat_datasets
 from data_designer.engine.resources.resource_provider import ResourceType
@@ -58,18 +59,60 @@ def _initialize(self) -> None:
         self._df_remaining = None
         self._dataset_uri = self.resource_provider.datastore.get_dataset_uri(self.config.dataset)
         self._seed_dataset_size = self.duckdb_conn.execute(f"SELECT COUNT(*) FROM '{self._dataset_uri}'").fetchone()[0]
+        self._index_range = self._resolve_index_range()
+
+    def _validate_selection_strategy(self) -> None:
+        err_msg = None
+        if self.config.selection_strategy is not None:
+            if isinstance(self.config.selection_strategy, IndexRange) and self.config.selection_strategy.end >= self._seed_dataset_size:
+                err_msg = f"Selection strategy 'end' index {self.config.selection_strategy.end} is out of bounds for dataset size {self._seed_dataset_size}"
+            elif isinstance(self.config.selection_strategy, PartitionBlock) and self.config.selection_strategy.num_partitions > self._seed_dataset_size:
+                err_msg = f"Selection strategy 'num_partitions' {self.config.selection_strategy.num_partitions} is out of bounds for dataset size {self._seed_dataset_size}"
+            if err_msg is not None:
+                raise SeedDatasetError(err_msg)
+
+    def _resolve_index_range(self) -> IndexRange | None:
+        self._validate_selection_strategy()
+        index_range = None
+        if self.config.selection_strategy is not None:
+            if isinstance(self.config.selection_strategy, IndexRange):
+                index_range = self.config.selection_strategy
+            elif isinstance(self.config.selection_strategy, PartitionBlock):
+                index_range = self.config.selection_strategy.to_index_range(self._seed_dataset_size)
+        return index_range
 
     def _reset_batch_reader(self, num_records: int) -> None:
         shuffle = self.config.sampling_strategy == SamplingStrategy.SHUFFLE
         shuffle_query = " ORDER BY RANDOM()" if shuffle else ""
-        self._batch_reader = self.duckdb_conn.query(f"SELECT * FROM '{self._dataset_uri}'{shuffle_query}").record_batch(
+
+        if self._index_range is not None:
+            # Use subquery with row_number() window function to filter by index range
+            # IndexRange uses 0-based indexing [start, end] inclusive, row_number() is 1-based
+            # To convert 0-based index i to 1-based row_number: row_number = i + 1
+            # For inclusive range [start, end], we want: row_number > start AND row_number <= end + 1
+            # This gives us 1-based rows [start+1, end+1] which maps to 0-based indices [start, end]
+            read_query = f"""
+                SELECT * EXCLUDE (row_num) FROM (
+                    SELECT *, row_number() OVER () as row_num
+                    FROM '{self._dataset_uri}'
+                ) sub
+                WHERE row_num > {self._index_range.start} AND row_num <= {self._index_range.end + 1}
+                {shuffle_query}
+            """
+        else:
+            read_query = f"SELECT * FROM '{self._dataset_uri}'{shuffle_query}"
+
+        self._batch_reader = self.duckdb_conn.query(read_query).record_batch(
             batch_size=num_records
         )
 
     def _sample_records(self, num_records: int) -> pd.DataFrame:
         logger.info(f"🌱 Sampling {num_records} records from seed dataset")
         logger.info(f"  |-- seed dataset size: {self._seed_dataset_size} records")
         logger.info(f"  |-- sampling strategy: {self.config.sampling_strategy}")
+        if self._index_range is not None:
+            logger.info(f"  |-- selection strategy: {self.config.selection_strategy.model_dump_json()}")
+            logger.info(f"  |-- seed dataset size after selection: {self._index_range.size} records")
 
         df_batch = pd.DataFrame()
         df_sample = pd.DataFrame() if self._df_remaining is None else self._df_remaining

src/data_designer/engine/column_generators/utils/errors.py

@@ -8,3 +8,6 @@ class PromptTemplateRenderError(DataDesignerError): ...
 
 
 class ExpressionTemplateRenderError(DataDesignerError): ...
+
+
+class SeedDatasetError(DataDesignerError): ...

tests/config/test_seed.py

@@ -0,0 +1,52 @@
+import pytest
+
+from data_designer.config.seed import IndexRange, PartitionBlock
+
+def test_index_range_validation():
+    with pytest.raises(ValueError, match="should be greater than or equal to 0"):
+        IndexRange(start=-1, end=10)
+    
+    with pytest.raises(ValueError, match="should be greater than or equal to 0"):
+        IndexRange(start=0, end=-1)
+
+    with pytest.raises(ValueError, match="'start' index must be less than or equal to 'end' index"):
+        IndexRange(start=11, end=10)
+
+
+def test_index_range_size():
+    assert IndexRange(start=0, end=10).size == 11
+    assert IndexRange(start=1, end=10).size == 10
+    assert IndexRange(start=0, end=0).size == 1
+
+
+def test_partition_block_validation():
+    with pytest.raises(ValueError, match="should be greater than or equal to 0"):
+        PartitionBlock(partition_index=-1, num_partitions=10)
+    
+    with pytest.raises(ValueError, match="should be greater than or equal to 1"):
+        PartitionBlock(partition_index=0, num_partitions=0)
+    
+    with pytest.raises(ValueError, match="'partition_index' must be less than 'num_partitions'"):
+        PartitionBlock(partition_index=10, num_partitions=10)
+
+
+def test_partition_block_to_index_range():
+    index_range = PartitionBlock(partition_index=0, num_partitions=10).to_index_range(101)
+    assert index_range.start == 0
+    assert index_range.end == 9
+    assert index_range.size == 10
+
+    index_range = PartitionBlock(partition_index=1, num_partitions=10).to_index_range(105)
+    assert index_range.start == 10
+    assert index_range.end == 19
+    assert index_range.size == 10
+
+    index_range = PartitionBlock(partition_index=2, num_partitions=10).to_index_range(105)
+    assert index_range.start == 20
+    assert index_range.end == 29
+    assert index_range.size == 10
+
+    index_range = PartitionBlock(partition_index=9, num_partitions=10).to_index_range(105)
+    assert index_range.start == 90
+    assert index_range.end == 104
+    assert index_range.size == 15

tests/engine/column_generators/generators/test_seed_dataset.py

@@ -10,14 +10,15 @@
 import pytest
 
 from data_designer.config.columns import SeedDatasetColumnConfig
-from data_designer.config.seed import SamplingStrategy
+from data_designer.config.seed import SamplingStrategy, IndexRange, PartitionBlock
 from data_designer.engine.column_generators.generators.base import GenerationStrategy
 from data_designer.engine.column_generators.generators.seed_dataset import (
     MAX_ZERO_RECORD_RESPONSE_FACTOR,
     SeedDatasetColumnGenerator,
 )
 from data_designer.engine.dataset_builders.multi_column_configs import SeedDatasetMultiColumnConfig
-from data_designer.engine.resources.resource_provider import ResourceType
+from data_designer.engine.resources.resource_provider import ResourceType, ResourceProvider
+from data_designer.engine.column_generators.utils.errors import SeedDatasetError
 
 
 @pytest.fixture
@@ -333,7 +334,11 @@ def test_seed_dataset_column_generator_sample_records_multiple_batches(stub_seed
 # ============================================================================
 
 
-def create_generator_with_real_file(file_path: str, stub_resource_provider) -> SeedDatasetColumnGenerator:
+def create_generator_with_real_file(
+    file_path: str,
+    stub_resource_provider: ResourceProvider,
+    sampling_strategy: SamplingStrategy = SamplingStrategy.ORDERED,
+    selection_strategy: IndexRange | PartitionBlock | None = None) -> SeedDatasetColumnGenerator:
     """Helper function to create a generator with a real file and DuckDB connection."""
     config = SeedDatasetMultiColumnConfig(
         columns=[
@@ -344,7 +349,8 @@ def create_generator_with_real_file(file_path: str, stub_resource_provider) -> S
             SeedDatasetColumnConfig(name="score"),
         ],
         dataset=f"test/{os.path.basename(file_path)}",
-        sampling_strategy=SamplingStrategy.ORDERED,
+        sampling_strategy=sampling_strategy,
+        selection_strategy=selection_strategy,
     )
 
     # Create a real DuckDB connection (in-memory by default)
@@ -605,3 +611,112 @@ def test_seed_dataset_generator_uses_real_duckdb_connection(fixture_name, stub_r
     # Verify the connection can execute count queries
     count_result = generator.duckdb_conn.execute(f"SELECT COUNT(*) FROM '{file_path}'").fetchone()[0]
     assert count_result == 10
+
+
+# ============================================================================
+# Tests for SeedConfig selection strategies
+# ============================================================================
+@pytest.mark.parametrize(
+    "fixture_name",
+    [
+        "seed_dataset_parquet",
+        "seed_dataset_csv",
+        "seed_dataset_json",
+        "seed_dataset_jsonl",
+    ],
+)
+def test_seed_dataset_generator_index_range_selection_strategy(fixture_name, stub_resource_provider, request):
+    """Test that generator correctly applies index range selection strategy."""
+    # Ordered Sampling
+
+    # Range with a subset of items
+    file_path = request.getfixturevalue(fixture_name)
+    generator = create_generator_with_real_file(file_path, stub_resource_provider, sampling_strategy=SamplingStrategy.ORDERED, selection_strategy=IndexRange(start=4, end=8))
+    result = generator.generate_from_scratch(6)
+    assert len(result) == 6
+    assert list(result["name"]) == ["Eve", "Frank", "Grace", "Henry", "Ivy", "Eve"]
+
+    # Range with just one item
+    generator = create_generator_with_real_file(file_path, stub_resource_provider, sampling_strategy=SamplingStrategy.ORDERED, selection_strategy=IndexRange(start=4, end=4))
+    result = generator.generate_from_scratch(1)
+    assert len(result) == 1
+    assert list(result["name"]) == ["Eve"]
+
+    # Range with all items
+    generator = create_generator_with_real_file(file_path, stub_resource_provider, sampling_strategy=SamplingStrategy.ORDERED, selection_strategy=IndexRange(start=0, end=9))
+    result = generator.generate_from_scratch(10)
+    assert len(result) == 10
+    assert list(result["name"]) == ["Alice", "Bob", "Charlie", "David", "Eve", "Frank", "Grace", "Henry", "Ivy", "Jack"]
+
+    # Shuffle Sampling
+
+    # Range with a subset of items
+    generator = create_generator_with_real_file(file_path, stub_resource_provider, sampling_strategy=SamplingStrategy.SHUFFLE, selection_strategy=IndexRange(start=4, end=8))
+    result = generator.generate_from_scratch(10)
+    assert len(result) == 10
+    assert set(result["name"]).issubset({"Eve", "Frank", "Grace", "Henry", "Ivy"})
+
+    # Range with just one item
+    generator = create_generator_with_real_file(file_path, stub_resource_provider, sampling_strategy=SamplingStrategy.SHUFFLE, selection_strategy=IndexRange(start=4, end=4))
+    result = generator.generate_from_scratch(1)
+    assert len(result) == 1
+    assert list(result["name"]) == ["Eve"]
+
+    # Range with all items
+    generator = create_generator_with_real_file(file_path, stub_resource_provider, sampling_strategy=SamplingStrategy.SHUFFLE, selection_strategy=IndexRange(start=0, end=9))
+    result = generator.generate_from_scratch(10)
+    assert len(result) == 10
+    assert set(result["name"]).issubset({"Alice", "Bob", "Charlie", "David", "Eve", "Frank", "Grace", "Henry", "Ivy", "Jack"})
+
+
+@pytest.mark.parametrize(
+    "fixture_name",
+    [
+        "seed_dataset_parquet",
+        "seed_dataset_csv",
+        "seed_dataset_json",
+        "seed_dataset_jsonl",
+    ],
+)
+def test_seed_dataset_generator_partition_block_selection_strategy(fixture_name, stub_resource_provider, request):
+    """Test that generator correctly applies partition block selection strategy."""
+    file_path = request.getfixturevalue(fixture_name)
+    generator = create_generator_with_real_file(
+        file_path,
+        stub_resource_provider,
+        sampling_strategy=SamplingStrategy.ORDERED,
+        selection_strategy=PartitionBlock(partition_index=1, num_partitions=3)
+    )
+    result = generator.generate_from_scratch(5)
+    assert len(result) == 5
+    # Requesting 5 items from a 3-item partition should cycle:
+    assert list(result["name"]) == ["David", "Eve", "Frank", "David", "Eve"]
+
+    generator = create_generator_with_real_file(
+        file_path,
+        stub_resource_provider,
+        sampling_strategy=SamplingStrategy.SHUFFLE,
+        selection_strategy=PartitionBlock(partition_index=4, num_partitions=5))
+    result = generator.generate_from_scratch(10)
+    assert len(result) == 10
+    assert set(result["name"]).issubset({"Jack", "Ivy"})
+
+
+@pytest.mark.parametrize(
+    "fixture_name",
+    [
+        "seed_dataset_parquet",
+        "seed_dataset_csv",
+        "seed_dataset_json",
+        "seed_dataset_jsonl",
+    ],
+)
+def test_seed_dataset_generator_invalid_selection_strategies(fixture_name, stub_resource_provider, request):
+    """Test that generator raises an error for invalid selection strategies."""
+    file_path = request.getfixturevalue(fixture_name)
+    with pytest.raises(SeedDatasetError, match="Selection strategy 'end' index 10 is out of bounds for dataset size 10"):
+        generator = create_generator_with_real_file(file_path, stub_resource_provider, selection_strategy=IndexRange(start=1, end=10))
+        generator.generate_from_scratch(1)
+    with pytest.raises(SeedDatasetError, match="Selection strategy 'num_partitions' 11 is out of bounds for dataset size 10"):
+        generator = create_generator_with_real_file(file_path, stub_resource_provider, selection_strategy=PartitionBlock(partition_index=0, num_partitions=11))
+        generator.generate_from_scratch(1)