In PyArrow 20.0.0+, when using open_json to read data in batches, an …

[HunterLine]

In PyArrow 20.0.0+, when using open_json to read data in batches, an error may occur if a single record is too large:

Underlying PyArrow Error: straddling object straddles two block boundaries (try to increase block size?)

This PR fixes an issue in DataJuicer related to this behavior.

[gemini-code-assist]

Summary of Changes

Hello, 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 addresses a critical data loading issue in PyArrow versions 20.0.0 and above, specifically when processing large JSON records that exceed the default block size. By introducing configurable read_options to increase the PyArrow block_size, the system can now robustly handle very large single JSON objects, commonly found in LLM datasets, preventing "straddling object" errors and ensuring seamless data ingestion. The change also improves error diagnostics and includes a new test case to verify the solution.

Highlights

• PyArrow JSON Reading Fix: Resolved an issue where PyArrow's open_json failed to read large JSON records (in versions 20.0.0 and above) by allowing configurable read_options to increase the block_size.
• Configurable Read Options: Introduced the ability to pass read_options (e.g., for block_size) through the configuration (self.cfg) to the JSON data loading process.
• Improved Error Reporting: Enhanced the error message for pyarrow.lib.ArrowInvalid during JSON file reading to include the underlying PyArrow error for better debugging.
• New Test Case: Added a dedicated test to validate the fix for reading large JSON objects using read_json_stream with an increased block_size.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• data_juicer/core/data/load_strategy.py
  • Modified load_data to accept read_options as a keyword argument and pass it to read_json_stream when loading JSON data.
• data_juicer/core/data/ray_dataset.py
  • Updated the exception handling for pyarrow.lib.ArrowInvalid in _read_stream to include the original PyArrow error message, providing more context for debugging.
• data_juicer/core/executor/ray_executor.py
  • Added logic to retrieve read_options from the executor's configuration (self.cfg) and pass them to datasetbuilder.load_dataset.
  • Included a comment explaining the rationale for increasing block_size for large JSON objects in LLM data to prevent ArrowInvalid errors.
• tests/core/data/test_ray_dataset.py
  • Added a new test method test_read_json_stream to verify that read_json_stream can successfully read large JSON objects by setting an appropriate block_size in read_options.
  • Applied minor formatting adjustments to existing test data and TEST_TAG decorators.

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 addresses an issue with reading large JSON records in PyArrow by allowing read_options to be passed to increase the block size. The changes are logical and include a new test case to verify the fix. My review includes a few suggestions to improve code style and maintainability, such as using .get() for dictionary access, replacing os.system with a safer alternative, and using named constants for magic numbers in tests.

[Dludora]

Could we expose read_options so it can be set via YAML files? I see it's not currently included in config.py.

[HunterLine]

feat(config): add read_options config item

Add read_options parameter to the base argument parser in config.py.
This allows passing custom read options (e.g. block_size) to PyArrow
reading functions, taking advantage of improved performance when
loading large JSON files.