From 8338d897d535ed1c4f285b50556eec760d233d11 Mon Sep 17 00:00:00 2001 From: Par-star Date: Thu, 23 Jul 2026 15:18:31 +0530 Subject: [PATCH 1/3] Add file_name column option in JSON processing --- src/datasets/packaged_modules/json/json.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/datasets/packaged_modules/json/json.py b/src/datasets/packaged_modules/json/json.py index 9845dd676d4..80a9dc94e92 100644 --- a/src/datasets/packaged_modules/json/json.py +++ b/src/datasets/packaged_modules/json/json.py @@ -56,6 +56,7 @@ class JsonConfig(datasets.BuilderConfig): newlines_in_values: Optional[bool] = None on_mixed_types: Optional[Literal["use_json"]] = "use_json" parse_agent_traces: bool = True + return_file_name: bool = False def __post_init__(self): super().__post_init__() @@ -139,6 +140,14 @@ def _cast_table(self, pa_table: pa.Table, json_field_paths=()) -> pa.Table: def _generate_shards(self, base_files, files_iterables, original_files): yield from base_files + def _add_file_name_column(self, pa_table: pa.Table, file) -> pa.Table: + """Append a 'file_name' column with the source file path, when return_file_name=True.""" + if self.config.return_file_name: + pa_table = pa_table.append_column( + "file_name", pa.array([str(file)] * len(pa_table), type=pa.string()) + ) + return pa_table + def _generate_tables(self, base_files, files_iterables, original_files, allow_full_read=True): json_field_paths = [] is_agent_traces = False @@ -166,6 +175,7 @@ def _generate_tables(self, base_files, files_iterables, original_files, allow_fu if df.columns.tolist() == [0]: df.columns = list(self.config.features) if self.config.features else ["text"] pa_table = pa.Table.from_pandas(df, preserve_index=False) + pa_table = self._add_file_name_column(pa_table, file) yield Key(shard_idx, 0), self._cast_table(pa_table) # If the files are agent traces (one row = one file except for hermes which can have multiple sessions per file) @@ -218,6 +228,7 @@ def _generate_tables(self, base_files, files_iterables, original_files, allow_fu example = json_encode_field(example, json_field_path) examples.append(example) pa_table = pa.Table.from_pylist(examples) + pa_table = self._add_file_name_column(pa_table, file) yield Key(shard_idx, 0), self._cast_table(pa_table) # If the file has one json object per line @@ -336,8 +347,10 @@ def _generate_tables(self, base_files, files_iterables, original_files, allow_fu raise ValueError( f"Failed to convert pandas DataFrame to Arrow Table from file {file}." ) from None + pa_table = self._add_file_name_column(pa_table, file) yield Key(shard_idx, 0), self._cast_table(pa_table) break + pa_table = self._add_file_name_column(pa_table, file) yield ( Key(shard_idx, batch_idx), self._cast_table(pa_table, json_field_paths=json_field_paths), From e6c15b69500951449ef21034b3d087dc3057a8e3 Mon Sep 17 00:00:00 2001 From: Par-star Date: Thu, 23 Jul 2026 15:59:13 +0530 Subject: [PATCH 2/3] Update json.py --- src/datasets/packaged_modules/json/json.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/datasets/packaged_modules/json/json.py b/src/datasets/packaged_modules/json/json.py index 80a9dc94e92..1241ff424e6 100644 --- a/src/datasets/packaged_modules/json/json.py +++ b/src/datasets/packaged_modules/json/json.py @@ -75,7 +75,13 @@ def _info(self): ) if self.config.newlines_in_values is not None: raise ValueError("The JSON loader parameter `newlines_in_values` is no longer supported") - return datasets.DatasetInfo(features=self.config.features) + features = self.config.features + # If the caller passed an explicit schema, make sure it accounts for the + # extra `file_name` column we append when return_file_name=True, otherwise + # `_cast_table` rejects the column since it isn't part of the schema. + if self.config.return_file_name and features is not None and "file_name" not in features: + features = datasets.Features({**features, "file_name": Value("string")}) + return datasets.DatasetInfo(features=features) def _split_generators(self, dl_manager): """We handle string, list and dicts in datafiles""" From c8904da45a999dafe5a179a8f2a715f8e140a428 Mon Sep 17 00:00:00 2001 From: Par-star Date: Fri, 24 Jul 2026 09:35:01 +0530 Subject: [PATCH 3/3] Change list type to Sequence in path_or_paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #5354 ## What Changes the type annotation of `path_or_paths` from `Union[PathLike, list[PathLike]]` to `Union[PathLike, Sequence[PathLike]]` in: - `Dataset.from_csv` - `Dataset.from_json` - `Dataset.from_parquet` - `Dataset.from_text` ## Why `list` is invariant in mypy, so passing a `List[Union[str, bytes, PathLike]]` (or any other list subtype) to these functions raises a mypy error, even though it works correctly at runtime: error: Argument 1 to "from_parquet" has incompatible type "List[str]"; expected "Union[..., List[Union[str, bytes, PathLike[Any]]]]" [arg-type] note: "List" is invariant -- see https://mypy.readthedocs.io/en/stable/common_issues.html#variance note: Consider using "Sequence" instead, which is covariant Since these functions only read/iterate over `path_or_paths` and never mutate it, `Sequence` is the more accurate and mypy-friendly annotation, matching the standard `typing` recommendation for read-only sequence arguments. ## How I tested - Ran `mypy src/datasets/arrow_dataset.py` to confirm no new typing errors. - Ran existing tests: `pytest tests/test_arrow_dataset.py -k "parquet or csv or json or text"` - Confirmed the reported reproduction case from the issue no longer raises a mypy error. ## Backward compatibility No runtime behavior change — `Sequence` still accepts lists, tuples, etc., so this is purely a typing improvement and fully backward compatible. --- src/datasets/arrow_dataset.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py index 59451a640e6..3fc27d84ef6 100644 --- a/src/datasets/arrow_dataset.py +++ b/src/datasets/arrow_dataset.py @@ -1290,7 +1290,7 @@ def from_list( @staticmethod def from_csv( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, cache_dir: str = None, @@ -1430,7 +1430,7 @@ def from_generator( @staticmethod def from_json( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, cache_dir: str = None, @@ -1489,7 +1489,7 @@ def from_json( @staticmethod def from_parquet( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, cache_dir: str = None, @@ -1586,7 +1586,7 @@ def from_parquet( @staticmethod def from_text( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, cache_dir: str = None,