Feat/add read nums

[ChenZiHong-Gavin]

This pull request introduces a new feature to the read function, enabling users to specify a maximum number of documents to process.

[gemini-code-assist]

Summary of Changes

Hello @ChenZiHong-Gavin, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a new feature to the read function, enabling users to specify a maximum number of documents to process. This enhancement provides greater control over data ingestion, allowing for more efficient handling of large datasets by limiting the scope of operations, which can be particularly useful for testing or processing subsets of data.

Highlights

• New Parameter read_nums: A new optional parameter read_nums has been added to the read function, allowing users to specify a maximum number of documents to read.
• Document Limiting Functionality: The read function now supports limiting the number of documents processed by applying a .limit() operation to the Ray Dataset if read_nums is provided.
• Docstring Update: The docstring for the read function has been updated to include a description for the new read_nums parameter.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a read_nums parameter to limit the number of documents read. While the feature is useful, the current implementation has a significant issue regarding non-deterministic output due to the use of ray.data.Dataset.union() before applying the limit. This can affect reproducibility. Additionally, there is a minor issue of missing input validation for the new parameter. My review comment provides details on both issues and a code suggestion to fix the validation.

[gemini-code-assist]

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)