feat: add return_file_name option to JSON loader

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__()
@@ -338,12 +339,18 @@ def _generate_tables(self, base_files, files_iterables, original_files, allow_fu
                                     ) from None
                                 yield Key(shard_idx, 0), self._cast_table(pa_table)
                                 break
+                            if self.config.return_file_name:
+                                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),
                             )
                             batch_idx += 1
 
+    def _add_file_name_column(self, pa_table, file):
+        """Add a column with the file name of each row."""
+        return pa_table.append_column("file_name", pa.array([str(file)] * pa_table.num_rows))
+
 
 AGENT_TRACES_TYPES_VALUES = {
     "claude_code": ["user", "assistant", "system"],

tests/packaged_modules/test_json.py

@@ -807,3 +807,31 @@ def test_json_load_dataset_without_droid_marker_stays_ordinary_json(tmp_path):
 
     assert dataset.column_names == ["type", "id", "version", "timestamp", "message"]
     assert dataset[0]["type"] == "session_start"
+
+def test_json_load_dataset_with_return_file_name(tmp_path):
+    filename = tmp_path / "file.jsonl"
+    with open(filename, "w") as f:
+        f.write('{"col_1": 1}\n{"col_1": 2}\n')
+    dataset = load_dataset(
+        "json",
+        data_files=str(filename),
+        split="train",
+        cache_dir=str(tmp_path / "cache"),
+        return_file_name=True,
+    )
+    assert dataset.column_names == ["col_1", "file_name"]
+    assert dataset[0] == {"col_1": 1, "file_name": str(filename)}
+    assert dataset[1]["file_name"] == str(filename)
+
+
+def test_json_load_dataset_without_return_file_name(tmp_path):
+    filename = tmp_path / "file.jsonl"
+    with open(filename, "w") as f:
+        f.write('{"col_1": 1}\n{"col_1": 2}\n')
+    dataset = load_dataset(
+        "json",
+        data_files=str(filename),
+        split="train",
+        cache_dir=str(tmp_path / "cache"),
+    )
+    assert dataset.column_names == ["col_1"]

tests/test_return_file_name.py

@@ -0,0 +1,31 @@
+import json
+import os
+import tempfile
+
+from datasets import load_dataset
+
+
+def _make_json_file():
+    tmp_dir = tempfile.mkdtemp()
+    path = os.path.join(tmp_dir, "data.json")
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(
+            [
+                {"question": "什么是RAG？", "answer": "检索增强生成"},
+                {"question": "什么是LoRA？", "answer": "低秩适配"},
+            ],
+            f,
+            ensure_ascii=False,
+        )
+    return path
+
+
+def test_return_file_name_enabled():
+    ds = load_dataset("json", data_files=_make_json_file(), return_file_name=True)
+    assert "file_name" in ds["train"].column_names
+    assert ds["train"][0]["file_name"].endswith("data.json")
+
+
+def test_return_file_name_disabled_by_default():
+    ds = load_dataset("json", data_files=_make_json_file())
+    assert "file_name" not in ds["train"].column_names