diff --git a/nemo_curator/stages/file_partitioning.py b/nemo_curator/stages/file_partitioning.py index a5c1a5c526..2829927481 100644 --- a/nemo_curator/stages/file_partitioning.py +++ b/nemo_curator/stages/file_partitioning.py @@ -64,6 +64,12 @@ class FilePartitioningStage(ProcessingStage[_EmptyTask, FileGroupTask]): def __post_init__(self): """Initialize default values.""" + if self.files_per_partition is not None and self.blocksize is not None: + logger.warning( + "Both 'files_per_partition' and 'blocksize' were specified. " + "'files_per_partition' will take precedence and 'blocksize' will be ignored." + ) + self.blocksize = None if self.file_extensions is None: self.file_extensions = [".jsonl", ".json", ".parquet"] if self.storage_options is None: diff --git a/tests/stages/common/test_file_partitioning.py b/tests/stages/common/test_file_partitioning.py index 838289a957..866c3beb8e 100644 --- a/tests/stages/common/test_file_partitioning.py +++ b/tests/stages/common/test_file_partitioning.py @@ -69,12 +69,11 @@ def test_initialization_default_values(self): assert stage.limit is None assert stage.name == "file_partitioning" - def test_initialization_custom_values(self): - """Test initialization with custom parameter values.""" + def test_initialization_custom_values_with_files_per_partition(self): + """Test initialization with custom parameter values using files_per_partition.""" stage = FilePartitioningStage( file_paths="/custom/path", files_per_partition=5, - blocksize="128MB", file_extensions=[".txt", ".json"], storage_options={"key": "value"}, limit=3, @@ -82,6 +81,23 @@ def test_initialization_custom_values(self): assert stage.file_paths == "/custom/path" assert stage.files_per_partition == 5 + assert stage.blocksize is None + assert stage.file_extensions == [".txt", ".json"] + assert stage.storage_options == {"key": "value"} + assert stage.limit == 3 + + def test_initialization_custom_values_with_blocksize(self): + """Test initialization with custom parameter values using blocksize.""" + stage = FilePartitioningStage( + file_paths="/custom/path", + blocksize="128MB", + file_extensions=[".txt", ".json"], + storage_options={"key": "value"}, + limit=3, + ) + + assert stage.file_paths == "/custom/path" + assert stage.files_per_partition is None assert stage.blocksize == "128MB" assert stage.file_extensions == [".txt", ".json"] assert stage.storage_options == {"key": "value"} @@ -194,6 +210,19 @@ def test_process_with_blocksize(self, empty_task: _EmptyTask, tmp_path: Path): assert len(task.data) == 1 assert task.data[0] == test_files[i] + def test_both_blocksize_and_files_per_partition_warns(self, caplog: pytest.LogCaptureFixture): + """Test that specifying both blocksize and files_per_partition logs a warning and ignores blocksize.""" + with caplog.at_level("WARNING"): + stage = FilePartitioningStage( + file_paths="/test/path", + files_per_partition=2, + blocksize="128MB", + ) + assert stage.files_per_partition == 2 + assert stage.blocksize is None + assert "files_per_partition" in caplog.text + assert "blocksize" in caplog.text + def test_process_empty_file_list(self, empty_task: _EmptyTask): """Test processing with empty file list.""" stage = FilePartitioningStage(file_paths=[])