Skip to content

Use Sequence instead of list for path_or_paths typing #8361

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Par-star wants to merge 3 commits into
base: main
Choose a base branch
from
+24 −5
Open

Use Sequence instead of list for path_or_paths typing #8361

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8338d89
Add file_name column option in JSON processing
Par-star Jul 23, 2026
e6c15b6
Update json.py
Par-star Jul 23, 2026
c8904da
Change list type to Sequence in path_or_paths
Par-star Jul 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions src/datasets/arrow_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
21 changes: 20 additions & 1 deletion src/datasets/packaged_modules/json/json.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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__()
Expand All @@ -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"""
Expand Down Expand Up @@ -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())

@Sanjays2402 Sanjays2402 Jul 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For remote URLs, file here is the DownloadManager cache/extraction path, so return_file_name=True returns a machine-local filename instead of the source URL. The agent-trace branch below uses original_files[shard_idx] for this; should this helper get that value too?

@Par-star Par-star Jul 25, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're right — for remote files, file was the
DownloadManager's resolved local cache path, not the source path/URL.

I've updated _add_file_name_column to use original_files[shard_idx]
consistently across all call sites (the field-specific path, the pandas
fallback path, and the main JSON-lines path). For the agent-traces branch
I reused the file_path variable that was already being computed the
same way for example["file_path"], so both stay consistent.

Pushed the fix — let me know if this looks right to you.

)
return pa_table

def _generate_tables(self, base_files, files_iterables, original_files, allow_full_read=True):
json_field_paths = []
is_agent_traces = False
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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),
Expand Down