update tests + test e2e

Author: nabinchha

src/data_designer/config/config_builder.py

@@ -6,7 +6,7 @@
 import json
 import logging
 from pathlib import Path
-
+from typing import Union, Optional
 from pygments import highlight
 from pygments.formatters import HtmlFormatter
 from pygments.lexers import PythonLexer
@@ -31,6 +31,8 @@
     SamplingStrategy,
     SeedConfig,
     SeedDatasetReference,
+    IndexRange,
+    PartitionBlock,
 )
 from .utils.constants import DEFAULT_REPR_HTML_STYLE, REPR_HTML_TEMPLATE
 from .utils.info import DataDesignerInfo
@@ -113,7 +115,7 @@ def from_config(cls, config: dict | str | Path | BuilderConfig) -> Self:
                     datastore_settings=builder_config.datastore_settings,
                 )
                 builder.set_seed_datastore_settings(builder_config.datastore_settings)
-            builder.with_seed_dataset(seed_dataset_reference, sampling_strategy=config.seed_config.sampling_strategy)
+            builder.with_seed_dataset(seed_dataset_reference, sampling_strategy=config.seed_config.sampling_strategy, selection_strategy=config.seed_config.selection_strategy)
 
         return builder
 
@@ -493,6 +495,7 @@ def with_seed_dataset(
         dataset_reference: SeedDatasetReference,
         *,
         sampling_strategy: SamplingStrategy = SamplingStrategy.ORDERED,
+        selection_strategy: Optional[Union[IndexRange, PartitionBlock]] = None,
     ) -> Self:
         """Add a seed dataset to the current Data Designer configuration.
 
@@ -508,7 +511,7 @@ def with_seed_dataset(
         Returns:
             The current Data Designer config builder instance.
         """
-        self._seed_config = SeedConfig(dataset=dataset_reference.dataset, sampling_strategy=sampling_strategy)
+        self._seed_config = SeedConfig(dataset=dataset_reference.dataset, sampling_strategy=sampling_strategy, selection_strategy=selection_strategy)
         self.set_seed_datastore_settings(
             dataset_reference.datastore_settings if hasattr(dataset_reference, "datastore_settings") else None
         )

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

@@ -60,7 +60,7 @@ def _initialize(self) -> 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:
@@ -115,9 +115,8 @@ def _sample_records(self, num_records: int) -> pd.DataFrame:
         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"  |-- selection strategy: {type(self.config.selection_strategy).__name__}\n{self.config.selection_strategy.model_dump_json(indent=4)}")
             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
         num_zero_record_responses = 0

src/data_designer/engine/dataset_builders/utils/config_compiler.py

@@ -35,6 +35,7 @@ def compile_dataset_builder_column_configs(config: DataDesignerConfig) -> list[D
                 columns=seed_column_configs,
                 dataset=config.seed_config.dataset,
                 sampling_strategy=config.seed_config.sampling_strategy,
+                selection_strategy=config.seed_config.selection_strategy,
             )
         )
 

tests/engine/column_generators/generators/test_seed_dataset.py

@@ -124,6 +124,30 @@ def test_seed_dataset_column_generator_config_structure():
     assert config.columns[0].column_type.value == "seed-dataset"
     assert config.columns[1].name == "col2"
     assert config.columns[1].column_type.value == "seed-dataset"
+    assert config.selection_strategy is None
+
+    # Test PartitionBlock selection strategy
+    config = SeedDatasetMultiColumnConfig(
+        columns=[SeedDatasetColumnConfig(name="col1"), SeedDatasetColumnConfig(name="col2")],
+        dataset="test/dataset",
+        sampling_strategy=SamplingStrategy.SHUFFLE,
+        selection_strategy=PartitionBlock(partition_index=1, num_partitions=3),
+    )
+    assert isinstance(config.selection_strategy, PartitionBlock)
+    assert config.selection_strategy.partition_index == 1
+    assert config.selection_strategy.num_partitions == 3
+
+    # Test IndexRange selection strategy
+    config = SeedDatasetMultiColumnConfig(
+        columns=[SeedDatasetColumnConfig(name="col1"), SeedDatasetColumnConfig(name="col2")],
+        dataset="test/dataset",
+        sampling_strategy=SamplingStrategy.SHUFFLE,
+        selection_strategy=IndexRange(start=0, end=1),
+    )
+    assert isinstance(config.selection_strategy, IndexRange)
+    assert config.selection_strategy.start == 0
+    assert config.selection_strategy.end == 1
+    assert config.selection_strategy.size == 2
 
     # Test constants and enum values
     assert MAX_ZERO_RECORD_RESPONSE_FACTOR == 2