[Data] Prefetch data for PandasJSONDatasource

[bveeramani]

Why are these changes needed?

PandasJSONDatasource reads a file twice: once to sample a row and infer the numbers of rows to read per batch, and again to actually load the data. To reset the file after sampling, the datasource opens the file as a random-access file.

The issue is that PyArrow's random-access file doesn't prefetch enough data, which leads to many costly small requests and poor performance.

To mitigate this issue, this PR wraps the file in io.BufferedReader and prefetches more data.

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[gemini-code-assist]

Summary of Changes

Hello @bveeramani, 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 focuses on enhancing the performance of PandasJSONDatasource by addressing inefficiencies in its file I/O. It resolves an issue where repeated small read requests, especially when interacting with PyArrow's random-access files, led to poor performance. The core solution involves integrating a custom buffered reader to prefetch data, thereby minimizing I/O overhead and accelerating the data loading process.

Highlights

• Performance Optimization: Introduced a buffering mechanism within PandasJSONDatasource to significantly reduce the number of small, costly I/O requests when reading JSON files, particularly from cloud storage, by prefetching data.
• New StrictBufferedReader Class: A custom StrictBufferedReader class has been implemented. This wrapper ensures that underlying file objects are read with full buffering and prevents premature closure by external libraries like pandas, which is crucial for the datasource's double-read pattern.
• Configurable Buffer Size: A _BUFFER_SIZE constant (set to 128 KiB) has been added to PandasJSONDatasource, allowing for a configurable buffer size to optimize data prefetching based on typical read patterns.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 is currently in preview and 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 to provide feedback.

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 StrictBufferedReader to improve performance for PandasJSONDatasource by enabling prefetching, which is a solid approach to mitigate the issue of numerous small reads from cloud storage. The implementation is sound. I have a couple of suggestions to enhance robustness and code clarity.

[gemini-code-assist]

This method doesn't handle the case where the input file is empty. If f is an empty file, pd.read_json will return an empty iterator, and next(reader) will raise a StopIteration exception, causing the read to fail.

You can make this more robust by providing a default value to next(). This will handle empty files gracefully by treating them as an empty DataFrame.

Suggested change

	df = _cast_range_index_to_string(next(reader))
	df = _cast_range_index_to_string(next(reader, pd.DataFrame()))

[iamjustinhsu]

how did u notice this?

[iamjustinhsu]

curious, why call it Strict?

[bveeramani]

I observed that the built-in io.BufferedReader implementation doesn't do a good job of actually buffering the data. I think it's because pandas calls BufferedReader.read1, and read1 doesn't prefill the 1 MiB buffer.

The name "Strict" is used to denote that it always performs the buffering.

[iamjustinhsu]

is read1 supposed to be read()?

[bveeramani]

No, TextIOWrapper calls read1 I think

[bveeramani]

how did u notice this?

@iamjustinhsu I ran a batch inference workload, and noticed that the read tasks were unacceptably slow.