Skip to content

Feat/add read nums #144

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

Merged
ChenZiHong-Gavin merged 2 commits into from
Dec 29, 2025
Merged

Feat/add read nums #144

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b5d026e
feat: add param read_nums
ChenZiHong-Gavin Dec 29, 2025
ecb6173
feat: add param read_nums
ChenZiHong-Gavin Dec 29, 2025
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
5 changes: 5 additions & 0 deletions graphgen/operators/read/read.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ def read(
working_dir: Optional[str] = "cache",
parallelism: int = 4,
recursive: bool = True,
read_nums: Optional[int] = None,
**reader_kwargs: Any,
) -> ray.data.Dataset:
"""
Expand All @@ -63,6 +64,7 @@ def read(
:param working_dir: Directory to cache intermediate files (PDF processing)
:param parallelism: Number of parallel workers
:param recursive: Whether to scan directories recursively
:param read_nums: Limit the number of documents to read
:param reader_kwargs: Additional kwargs passed to readers
:return: Ray Dataset containing all documents
"""
Expand Down Expand Up @@ -120,6 +122,9 @@ def read(
}
)

if read_nums is not None:
combined_ds = combined_ds.limit(read_nums)
Comment on lines +125 to +126
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Dec 29, 2025

Choose a reason for hiding this comment

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

high

This implementation has a couple of issues:

  1. Non-determinism (High Severity): The use of limit() here can lead to non-deterministic results. combined_ds is created using ray.data.Dataset.union() on line 116, which does not guarantee the order of records between runs. Consequently, limit(read_nums) will yield different subsets of data each time, which is problematic for reproducibility. To improve determinism, you should ensure the file processing order is stable (e.g., by sorting file lists). For true determinism, a more significant refactor to avoid union() might be necessary.

  2. Missing Input Validation (Medium Severity): The read_nums parameter is not validated. A negative value will cause limit() to raise a ValueError, but it's better to validate inputs explicitly in this function to provide a clearer error message.

Here is a suggestion to address the input validation:

Suggested change
if read_nums is not None:
combined_ds = combined_ds.limit(read_nums)
if read_nums is not None:
if read_nums < 0:
raise ValueError("read_nums cannot be negative")
combined_ds = combined_ds.limit(read_nums)


logger.info("[READ] Successfully read files from %s", input_path)
return combined_ds

Expand Down
Loading