From 197651a5ef0734051c86348ec6ac1db05088e9f0 Mon Sep 17 00:00:00 2001 From: RenhaoSong Date: Fri, 31 Jul 2026 19:00:28 +0800 Subject: [PATCH 1/3] feat: add return_file_name option to json loader --- 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 9845dd676d4..068047a3c19 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__() @@ -338,12 +339,17 @@ 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"], From 121f8e33d3e4d757201387d2a21358ecc7832aba Mon Sep 17 00:00:00 2001 From: RenhaoSong Date: Fri, 31 Jul 2026 19:07:58 +0800 Subject: [PATCH 2/3] test: add tests for return_file_name option --- src/datasets/packaged_modules/json/json.py | 2 +- tests/packaged_modules/test_json.py | 28 ++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/src/datasets/packaged_modules/json/json.py b/src/datasets/packaged_modules/json/json.py index 068047a3c19..b0f62c9e18d 100644 --- a/src/datasets/packaged_modules/json/json.py +++ b/src/datasets/packaged_modules/json/json.py @@ -346,7 +346,7 @@ def _generate_tables(self, base_files, files_iterables, original_files, allow_fu 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)) diff --git a/tests/packaged_modules/test_json.py b/tests/packaged_modules/test_json.py index 2d6fd5b3849..8ca913fbb54 100644 --- a/tests/packaged_modules/test_json.py +++ b/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"] From 679f320772adc15b9a6c23d16853fb2499ef81ef Mon Sep 17 00:00:00 2001 From: RenhaoSong Date: Fri, 31 Jul 2026 21:02:18 +0800 Subject: [PATCH 3/3] feat: add return_file_name option to JSON loader --- src/datasets/packaged_modules/json/json.py | 1 + tests/test_return_file_name.py | 31 ++++++++++++++++++++++ 2 files changed, 32 insertions(+) create mode 100644 tests/test_return_file_name.py diff --git a/src/datasets/packaged_modules/json/json.py b/src/datasets/packaged_modules/json/json.py index b0f62c9e18d..171b0775e52 100644 --- a/src/datasets/packaged_modules/json/json.py +++ b/src/datasets/packaged_modules/json/json.py @@ -351,6 +351,7 @@ 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"], "pi": ["session", "message"], diff --git a/tests/test_return_file_name.py b/tests/test_return_file_name.py new file mode 100644 index 00000000000..ad29c2dd256 --- /dev/null +++ b/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