json: add optional return_file_name parameter

src/datasets/packaged_modules/json/json.py

@@ -12,6 +12,7 @@
 from datasets.table import table_cast
 from datasets.utils.file_utils import readline
 
+import os
 
 logger = datasets.utils.logging.get_logger(__name__)
 
@@ -50,6 +51,7 @@ class JsonConfig(datasets.BuilderConfig):
     block_size: Optional[int] = None  # deprecated
     chunksize: int = 10 << 20  # 10MB
     newlines_in_values: Optional[bool] = None
+    return_file_name : bool = False 
 
     def __post_init__(self):
         super().__post_init__()
@@ -129,6 +131,12 @@ def _generate_tables(self, base_files, files_iterables):
                     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)
+                    if self.config.return_file_name:
+                        file_name = os.path.basename(file)
+                        pa_table = pa_table.append_column( #add file_name to column
+                            "file_name",
+                            pa.array([file_name] * len(pa_table), type=pa.string())
+                        )
                     yield Key(shard_idx, 0), self._cast_table(pa_table)
 
                 # If the file has one json object per line
@@ -196,5 +204,11 @@ def _generate_tables(self, base_files, files_iterables):
                                     ) from None
                                 yield Key(shard_idx, 0), self._cast_table(pa_table)
                                 break
+                            if self.config.return_file_name:
+                                file_name = os.path.basename(file)
+                                pa_table = pa_table.append_column(
+                                    "file_name",
+                                    pa.array([file_name] * len(pa_table), type=pa.string())
+                                )
                             yield Key(shard_idx, batch_idx), self._cast_table(pa_table)
                             batch_idx += 1

tests/packaged_modules/test_json.py

@@ -271,3 +271,32 @@ def test_json_generate_tables_with_sorted_columns(file_fixture, config_kwargs, r
     generator = json._generate_tables(base_files=base_files, files_iterables=files_iterables)
     pa_table = pa.concat_tables([table for _, table in generator])
     assert pa_table.column_names == ["ID", "Language", "Topic"]
+
+def test_json_no_file_name_by_default(jsonl_file):
+    """Ensure backward compatibility, when return_file_name is not set"""
+    json = Json()
+    base_files = [jsonl_file]
+    files_iterables = [[file] for file in base_files]
+    generator = json._generate_tables(base_files=base_files, files_iterables=files_iterables)
+    pa_table = pa.concat_tables([table for _, table in generator])
+    assert "file_name" not in pa_table.column_names
+
+def test_json_return_file_name_enabled(jsonl_file):
+    json = Json(return_file_name = True)
+    base_files = [jsonl_file]
+    files_iterables = [[file] for file in base_files]
+    generator = json._generate_tables(base_files=base_files, files_iterables=files_iterables)
+    pa_table = pa.concat_tables([table for _, table in generator])
+    assert "file_name" in pa_table.column_names
+
+def test_json_file_name_values(jsonl_file):
+    """File name column should contain the source file basename for each row."""
+    json = Json(return_file_name=True)
+    base_files = [jsonl_file]
+    files_iterables = [[jsonl_file]]
+
+    generator = json._generate_tables(base_files=base_files, files_iterables=files_iterables)
+    pa_table = pa.concat_tables([table for _, table in generator])
+
+    data = pa_table.to_pydict()
+    assert all(name == "file.jsonl" for name in data["file_name"])