Support wildcard path pattern for seed dataset

Author: nabinchha

src/data_designer/config/seed.py

@@ -10,7 +10,11 @@
 
 from .base import ConfigBase
 from .datastore import DatastoreSettings
-from .utils.io_helpers import validate_dataset_file_path
+from .utils.io_helpers import (
+    VALID_DATASET_FILE_EXTENSIONS,
+    validate_dataset_file_path,
+    validate_path_contains_files_of_type,
+)
 
 
 class SamplingStrategy(str, Enum):
@@ -130,4 +134,12 @@ def filename(self) -> str:
 class LocalSeedDatasetReference(SeedDatasetReference):
     @field_validator("dataset", mode="after")
     def validate_dataset_is_file(cls, v: str) -> str:
-        return str(validate_dataset_file_path(v))
+        valid_wild_card_versions = {f"*{ext}" for ext in VALID_DATASET_FILE_EXTENSIONS}
+        if any(v.endswith(wildcard) for wildcard in valid_wild_card_versions):
+            parts = v.split("*.")
+            file_path = parts[0]
+            file_extension = parts[-1]
+            validate_path_contains_files_of_type(file_path, file_extension)
+        else:
+            validate_dataset_file_path(v)
+        return v

src/data_designer/config/utils/io_helpers.py

@@ -68,9 +68,11 @@ def validate_dataset_file_path(file_path: str | Path, should_exist: bool = True)
     Args:
         file_path: The path to validate, either as a string or Path object.
         should_exist: If True, verify that the file exists. Defaults to True.
-
     Returns:
         The validated path as a Path object.
+    Raises:
+        InvalidFilePathError: If the path is not a file.
+        InvalidFileFormatError: If the path does not have a valid extension.
     """
     file_path = Path(file_path)
     if should_exist and not Path(file_path).is_file():
@@ -82,6 +84,21 @@ def validate_dataset_file_path(file_path: str | Path, should_exist: bool = True)
     return file_path
 
 
+def validate_path_contains_files_of_type(path: str | Path, file_extension: str) -> None:
+    """Validate that a path contains files of a specific type.
+
+    Args:
+        path: The path to validate. Can contain wildcards like `*.parquet`.
+        file_extension: The extension of the files to validate (without the dot, e.g., "parquet").
+    Returns:
+        None if the path contains files of the specified type, raises an error otherwise.
+    Raises:
+        InvalidFilePathError: If the path does not contain files of the specified type.
+    """
+    if not any(Path(path).glob(f"*.{file_extension}")):
+        raise InvalidFilePathError(f"🛑 Path {path!r} does not contain files of type {file_extension!r}.")
+
+
 def smart_load_dataframe(dataframe: str | Path | pd.DataFrame) -> pd.DataFrame:
     """Load a dataframe from file if a path is given, otherwise return the dataframe.
 

tests/config/test_seed.py

@@ -1,9 +1,30 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from pathlib import Path
+import tempfile
+
+import pandas as pd
 import pytest
 
-from data_designer.config.seed import IndexRange, PartitionBlock
+from data_designer.config.errors import InvalidFilePathError
+from data_designer.config.seed import IndexRange, LocalSeedDatasetReference, PartitionBlock
+
+
+def create_partitions_in_path(temp_dir: Path, extension: str, num_files: int = 2) -> Path:
+    df = pd.DataFrame({"col": [1, 2, 3]})
+
+    for i in range(num_files):
+        file_path = temp_dir / f"partition_{i}.{extension}"
+        if extension == "parquet":
+            df.to_parquet(file_path)
+        elif extension == "csv":
+            df.to_csv(file_path, index=False)
+        elif extension == "json":
+            df.to_json(file_path, orient="records")
+        elif extension == "jsonl":
+            df.to_json(file_path, orient="records", lines=True)
+    return temp_dir
 
 
 def test_index_range_validation():
@@ -54,3 +75,35 @@ def test_partition_block_to_index_range():
     assert index_range.start == 90
     assert index_range.end == 104
     assert index_range.size == 15
+
+
+def test_local_seed_dataset_reference_validation():
+    with pytest.raises(InvalidFilePathError, match="🛑 Path test/dataset.parquet is not a file."):
+        LocalSeedDatasetReference(dataset="test/dataset.parquet")
+
+    # Should not raise an error when referencing supported extensions with wildcard pattern.
+    with tempfile.TemporaryDirectory() as temp_dir:
+        create_partitions_in_path(Path(temp_dir), "parquet")
+        create_partitions_in_path(Path(temp_dir), "csv")
+        create_partitions_in_path(Path(temp_dir), "json")
+        create_partitions_in_path(Path(temp_dir), "jsonl")
+
+        test_cases = [
+            (temp_dir, "parquet"),
+            (temp_dir, "csv"),
+            (temp_dir, "json"),
+            (temp_dir, "jsonl"),
+        ]
+
+        try:
+            for temp_dir, extension in test_cases:
+                reference = LocalSeedDatasetReference(dataset=f"{temp_dir}/*.{extension}")
+                assert reference.dataset == f"{temp_dir}/*.{extension}"
+        except Exception as e:
+            pytest.fail(f"Expected no exception, but got {e}")
+
+    # Should raise an error when referencing a path that does not contain files of the specified type.
+    with tempfile.TemporaryDirectory() as temp_dir:
+        create_partitions_in_path(Path(temp_dir), "parquet")
+        with pytest.raises(InvalidFilePathError, match="does not contain files of type 'csv'"):
+            LocalSeedDatasetReference(dataset=f"{temp_dir}/*.csv")