partition_index -> index

Author: nabinchha

src/data_designer/config/seed.py

@@ -34,24 +34,24 @@ def size(self) -> int:
 
 
 class PartitionBlock(ConfigBase):
-    partition_index: int = Field(default=0, ge=0, description="The index of the partition to sample from")
+    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'")
+        if self.index >= self.num_partitions:
+            raise ValueError("'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
+        start = self.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:
+        if self.index == self.num_partitions - 1:
             end = dataset_size - 1
         else:
-            end = ((self.partition_index + 1) * partition_size) - 1
+            end = ((self.index + 1) * partition_size) - 1
         return IndexRange(start=start, end=end)
 
 

tests/config/test_seed.py

@@ -25,32 +25,32 @@ def test_index_range_size():
 
 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)
+        PartitionBlock(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)
+        PartitionBlock(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)
+    with pytest.raises(ValueError, match="'index' must be less than 'num_partitions'"):
+        PartitionBlock(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)
+    index_range = PartitionBlock(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)
+    index_range = PartitionBlock(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)
+    index_range = PartitionBlock(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)
+    index_range = PartitionBlock(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

@@ -131,10 +131,10 @@ def test_seed_dataset_column_generator_config_structure():
         columns=[SeedDatasetColumnConfig(name="col1"), SeedDatasetColumnConfig(name="col2")],
         dataset="test/dataset",
         sampling_strategy=SamplingStrategy.SHUFFLE,
-        selection_strategy=PartitionBlock(partition_index=1, num_partitions=3),
+        selection_strategy=PartitionBlock(index=1, num_partitions=3),
     )
     assert isinstance(config.selection_strategy, PartitionBlock)
-    assert config.selection_strategy.partition_index == 1
+    assert config.selection_strategy.index == 1
     assert config.selection_strategy.num_partitions == 3
 
     # Test IndexRange selection strategy
@@ -742,7 +742,7 @@ def test_seed_dataset_generator_partition_block_selection_strategy(fixture_name,
         file_path,
         stub_resource_provider,
         sampling_strategy=SamplingStrategy.ORDERED,
-        selection_strategy=PartitionBlock(partition_index=1, num_partitions=3),
+        selection_strategy=PartitionBlock(index=1, num_partitions=3),
     )
     result = generator.generate_from_scratch(5)
     assert len(result) == 5
@@ -753,7 +753,7 @@ def test_seed_dataset_generator_partition_block_selection_strategy(fixture_name,
         file_path,
         stub_resource_provider,
         sampling_strategy=SamplingStrategy.SHUFFLE,
-        selection_strategy=PartitionBlock(partition_index=4, num_partitions=5),
+        selection_strategy=PartitionBlock(index=4, num_partitions=5),
     )
     result = generator.generate_from_scratch(10)
     assert len(result) == 10
@@ -783,6 +783,6 @@ def test_seed_dataset_generator_invalid_selection_strategies(fixture_name, stub_
         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)
+            file_path, stub_resource_provider, selection_strategy=PartitionBlock(index=0, num_partitions=11)
         )
         generator.generate_from_scratch(1)