diff --git a/src/datasets/packaged_modules/json/json.py b/src/datasets/packaged_modules/json/json.py index 9845dd676d4..171b0775e52 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,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"], 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"] 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