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, diff --git a/src/datasets/packaged_modules/json/json.py b/src/datasets/packaged_modules/json/json.py index 9845dd676d4..1241ff424e6 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__() @@ -74,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""" @@ -139,6 +146,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 +181,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 +234,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 +353,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),