[Data] Auto increasing block size for read_json #42357

omatthew98 · 2024-01-12T00:44:53Z

Why are these changes needed?

This PR adds logic to dynamically increase block_size used in Arrow's JSON loader, if the initial block_size resulting in a more graceful handling of these cases.

Related issue number

Closes #41196

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 :(

Signed-off-by: Matthew Owen <[email protected]>

…lines Signed-off-by: Matthew Owen <[email protected]>

Signed-off-by: Matthew Owen <[email protected]>

scottjlee · 2024-01-19T19:39:14Z

python/ray/data/datasource/json_datasource.py

 from ray.data.datasource.file_based_datasource import FileBasedDatasource
 from ray.util.annotations import PublicAPI

 if TYPE_CHECKING:
    import pyarrow

+logger = logging.getLogger(__name__)


let's use the DatasetLogger instead
https://github.com/ray-project/ray/blob/master/python/ray/data/_internal/dataset_logger.py

scottjlee · 2024-01-19T19:56:25Z

python/ray/data/read_api.py

        When reading large files, the default block size configured in PyArrow can be too small,
        resulting in the following error:
        ``pyarrow.lib.ArrowInvalid: straddling object straddles two block boundaries
-        (try to increase block size?)``.
-
-        To resolve this, use the ``read_options`` parameter to set a larger block size:
+        (try to increase block size?)``. The read will be retried with geometrically
+        increasing block size until the size reaches `DataContext.get_current().target_max_block_size`.
+        The initial block size will start at the PyArrow default block size or it can be
+        manually set through the ``read_options`` parameter as follows:


with this PR, ideally the user shouldn't be seeing this error message any more.
let's instead move this info (related to the "try to increase block size?") into the new code you implemented in _read_stream()

scottjlee · 2024-01-19T20:26:04Z

python/ray/data/datasource/json_datasource.py

+                    read_options=json.ReadOptions(
+                        use_threads=use_threads, block_size=block_size
+                    ),


since ReadOptions can have other attributes/parameters as well, can we pass those or create + modify a copy?

scottjlee · 2024-01-19T20:30:38Z

python/ray/data/datasource/json_datasource.py

+                    and "straddling" not in str(e)
+                    or block_size > max_block_size
+                ):
+                    raise e


let's modify the error message of the exception that is raised to include the largest block size that was tried. and let's also include a link to this GH issue for more details: apache/arrow#25674

Updated the exception to this: pyarrow.lib.ArrowInvalid: straddling object straddles two block boundaries (try to increase block size?) - Auto-increasing block size to 4B failed. More information on this issue can be found here: https://github.com/apache/arrow/issues/25674

Signed-off-by: Matthew Owen <[email protected]>

scottjlee · 2024-01-20T01:02:07Z

python/ray/data/datasource/json_datasource.py

+                if isinstance(e, ArrowInvalid) and "straddling" in str(e):
+                    if self.read_options.block_size < max_block_size:
+                        # Increase the block size in case it was too small.
+                        logger.get_logger(log_to_stdout=False).info(


let's use True for this since it is important, so that the user always sees this message in stdout logs

scottjlee · 2024-01-20T01:03:59Z

python/ray/data/datasource/json_datasource.py

+                self.read_options.block_size = init_block_size
+                break
+            except ArrowInvalid as e:
+                if isinstance(e, ArrowInvalid) and "straddling" in str(e):


we don't need this check inside the except block right? let's also compare to a longer string in case there are other error messages which use the word straddling:

Suggested change

if isinstance(e, ArrowInvalid) and "straddling" in str(e):

if "straddling object straddles two block boundaries" in str(e):

scottjlee · 2024-01-20T01:04:22Z

python/ray/data/datasource/json_datasource.py

+                    else:
+                        raise ArrowInvalid(
+                            f"{e} - Auto-increasing block size to "
+                            f"{self.read_options.block_size}B failed. "


nit

Suggested change

f"{self.read_options.block_size}B failed. "

f"{self.read_options.block_size} bytes failed. "

scottjlee · 2024-01-20T01:07:22Z

python/ray/data/datasource/json_datasource.py

+        # When reading large files, the default block size configured in PyArrow can be
+        # too small, resulting in the following error: `pyarrow.lib.ArrowInvalid:
+        # straddling object straddles two block boundaries (try to increase block
+        # size?)`. The read will be retried with geometrically increasing block size


nit: can we also include the arrow issue link here?

Signed-off-by: Matthew Owen <[email protected]>

c21

Great fix @omatthew98! Just one minor comment.

c21 · 2024-01-22T17:47:22Z

python/ray/data/datasource/json_datasource.py

+                if "straddling object straddles two block boundaries" in str(e):
+                    if self.read_options.block_size < max_block_size:
+                        # Increase the block size in case it was too small.
+                        logger.get_logger(log_to_stdout=True).info(


We don't need to log to stdout, as this would confuse users if there's something wrong.

Signed-off-by: Matthew Owen <[email protected]>

omatthew98 added 10 commits January 11, 2024 16:43

basic structure to auto increase block size for read_json

4857e95

Signed-off-by: Matthew Owen <[email protected]>

use target max block size for the limit we try before failing

f6902fe

Signed-off-by: Matthew Owen <[email protected]>

adding in break to fix behavior

ebe1cfc

Signed-off-by: Matthew Owen <[email protected]>

add in _open_input_source to allow for resetting file pos

afeac34

Signed-off-by: Matthew Owen <[email protected]>

add in test_json_read_across_blocks

360dc40

Signed-off-by: Matthew Owen <[email protected]>

refactoring to use pyarrow buffer

4c34794

Signed-off-by: Matthew Owen <[email protected]>

use smart_open

da74990

Signed-off-by: Matthew Owen <[email protected]>

fix smart_open call, format files

59a3e36

Signed-off-by: Matthew Owen <[email protected]>

simplify test to not rewrite file as it isn't necessary to make long …

f5b387a

…lines Signed-off-by: Matthew Owen <[email protected]>

updating docs to reflect changes

b5a6ffc

Signed-off-by: Matthew Owen <[email protected]>

omatthew98 marked this pull request as ready for review January 18, 2024 20:52

omatthew98 requested review from ericl, scv119, c21, amogkam, scottjlee, bveeramani, raulchen, stephanie-wang and Zandew as code owners January 18, 2024 20:52

scottjlee reviewed Jan 19, 2024

View reviewed changes

pr feedback

838933b

Signed-off-by: Matthew Owen <[email protected]>

omatthew98 requested a review from scottjlee January 19, 2024 21:57

scottjlee approved these changes Jan 20, 2024

View reviewed changes

respond to a few more nits

0177858

Signed-off-by: Matthew Owen <[email protected]>

omatthew98 assigned c21 Jan 20, 2024

c21 changed the title ~~Auto increasing block size for read_json~~ [Data] Auto increasing block size for read_json Jan 22, 2024

c21 approved these changes Jan 22, 2024

View reviewed changes

turn of log to stdout

76dead9

Signed-off-by: Matthew Owen <[email protected]>

c21 merged commit 3d91fb7 into ray-project:master Jan 22, 2024
9 checks passed

garyzhang99 mentioned this pull request Apr 11, 2024

[HOT FIX] Update ray version to support large file json read. modelscope/data-juicer#293

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Data] Auto increasing block size for read_json #42357

[Data] Auto increasing block size for read_json #42357

omatthew98 commented Jan 12, 2024 •

edited

Loading

scottjlee Jan 19, 2024

scottjlee Jan 19, 2024

scottjlee Jan 19, 2024

scottjlee Jan 19, 2024

omatthew98 Jan 19, 2024 •

edited

Loading

scottjlee Jan 20, 2024

scottjlee Jan 20, 2024

scottjlee Jan 20, 2024

scottjlee Jan 20, 2024

c21 left a comment

c21 Jan 22, 2024

	if isinstance(e, ArrowInvalid) and "straddling" in str(e):
	if "straddling object straddles two block boundaries" in str(e):

	f"{self.read_options.block_size}B failed. "
	f"{self.read_options.block_size} bytes failed. "

[Data] Auto increasing block size for read_json #42357

[Data] Auto increasing block size for read_json #42357

Conversation

omatthew98 commented Jan 12, 2024 • edited Loading

Why are these changes needed?

Related issue number

Checks

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

omatthew98 Jan 19, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

c21 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

omatthew98 commented Jan 12, 2024 •

edited

Loading

omatthew98 Jan 19, 2024 •

edited

Loading