wehub-resource-sync
294 lines
11 KiB
Python
294 lines
11 KiB
Python
import io
|
|
import logging
|
|
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
|
|
|
|
import pandas as pd
|
|
|
|
from ray.data._internal.pandas_block import PandasBlockAccessor
|
|
from ray.data._internal.tensor_extensions.arrow import pyarrow_table_from_pydict
|
|
from ray.data.context import DataContext
|
|
from ray.data.datasource.file_based_datasource import FileBasedDatasource
|
|
|
|
if TYPE_CHECKING:
|
|
import pyarrow
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
JSON_FILE_EXTENSIONS = [
|
|
"json",
|
|
"jsonl",
|
|
# gzip-compressed files
|
|
"json.gz",
|
|
"jsonl.gz",
|
|
# Brotli-compressed fi;es
|
|
"json.br",
|
|
"jsonl.br",
|
|
# Zstandard-compressed files
|
|
"json.zst",
|
|
"jsonl.zst",
|
|
# lz4-compressed files
|
|
"json.lz4",
|
|
"jsonl.lz4",
|
|
]
|
|
|
|
|
|
class ArrowJSONDatasource(FileBasedDatasource):
|
|
"""JSON datasource, for reading and writing JSON and JSONL files."""
|
|
|
|
def __init__(
|
|
self,
|
|
paths: Union[str, List[str]],
|
|
*,
|
|
arrow_json_args: Optional[Dict[str, Any]] = None,
|
|
**file_based_datasource_kwargs,
|
|
):
|
|
from pyarrow import json
|
|
|
|
super().__init__(paths, **file_based_datasource_kwargs)
|
|
|
|
if arrow_json_args is None:
|
|
arrow_json_args = {}
|
|
|
|
self.read_options = arrow_json_args.pop(
|
|
"read_options", json.ReadOptions(use_threads=False)
|
|
)
|
|
self.arrow_json_args = arrow_json_args
|
|
|
|
def _read_with_pyarrow_read_json(self, buffer: "pyarrow.lib.Buffer"):
|
|
"""Read with PyArrow JSON reader, trying to auto-increase the
|
|
read block size in the case of the read object
|
|
straddling block boundaries."""
|
|
import pyarrow as pa
|
|
import pyarrow.json as pajson
|
|
|
|
# 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?)`. More information on this issue can be found here:
|
|
# https://github.com/apache/arrow/issues/25674
|
|
# 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.
|
|
# >>> import pyarrow.json as pajson
|
|
# >>> block_size = 10 << 20 # Set block size to 10MB
|
|
# >>> ray.data.read_json( # doctest: +SKIP
|
|
# ... "s3://anonymous@ray-example-data/log.json",
|
|
# ... read_options=pajson.ReadOptions(block_size=block_size)
|
|
# ... )
|
|
|
|
init_block_size = self.read_options.block_size
|
|
max_block_size = DataContext.get_current().target_max_block_size
|
|
while True:
|
|
try:
|
|
yield pajson.read_json(
|
|
io.BytesIO(buffer),
|
|
read_options=self.read_options,
|
|
**self.arrow_json_args,
|
|
)
|
|
self.read_options.block_size = init_block_size
|
|
break
|
|
except pa.ArrowInvalid as e:
|
|
if "straddling object straddles two block boundaries" in str(e):
|
|
if (
|
|
max_block_size is None
|
|
or self.read_options.block_size < max_block_size
|
|
):
|
|
# Increase the block size in case it was too small.
|
|
logger.debug(
|
|
f"JSONDatasource read failed with "
|
|
f"block_size={self.read_options.block_size}. Retrying with "
|
|
f"block_size={self.read_options.block_size * 2}."
|
|
)
|
|
self.read_options.block_size *= 2
|
|
else:
|
|
raise pa.ArrowInvalid(
|
|
f"{e} - Auto-increasing block size to "
|
|
f"{self.read_options.block_size} bytes failed. "
|
|
f"Please try manually increasing the block size through "
|
|
f"the `read_options` parameter to a larger size. "
|
|
f"For example: `read_json(..., read_options="
|
|
f"pyarrow.json.ReadOptions(block_size=10 << 25))`"
|
|
f"More information on this issue can be found here: "
|
|
f"https://github.com/apache/arrow/issues/25674"
|
|
)
|
|
else:
|
|
# unrelated error, simply reraise
|
|
raise e
|
|
|
|
def _read_with_python_json(self, buffer: "pyarrow.lib.Buffer"):
|
|
"""Fallback method to read JSON files with Python's native json.load(),
|
|
in case the default pyarrow json reader fails."""
|
|
import json
|
|
|
|
import pyarrow as pa
|
|
|
|
# Check if the buffer is empty
|
|
if buffer.size == 0:
|
|
return
|
|
|
|
parsed_json = json.load(io.BytesIO(buffer))
|
|
try:
|
|
yield pa.Table.from_pylist(parsed_json)
|
|
except AttributeError as e:
|
|
# For PyArrow < 7.0.0, `pa.Table.from_pylist()` is not available.
|
|
# Construct a dict from the list and call
|
|
# `pa.Table.from_pydict()` instead.
|
|
assert "no attribute 'from_pylist'" in str(e), str(e)
|
|
from collections import defaultdict
|
|
|
|
dct = defaultdict(list)
|
|
for row in parsed_json:
|
|
for k, v in row.items():
|
|
dct[k].append(v)
|
|
yield pyarrow_table_from_pydict(dct)
|
|
|
|
# TODO(ekl) The PyArrow JSON reader doesn't support streaming reads.
|
|
def _read_stream(self, f: "pyarrow.NativeFile", path: str):
|
|
import pyarrow as pa
|
|
|
|
buffer: pa.lib.Buffer = f.read_buffer()
|
|
|
|
try:
|
|
yield from self._read_with_pyarrow_read_json(buffer)
|
|
except pa.ArrowInvalid as e:
|
|
# If read with PyArrow fails, try falling back to native json.load().
|
|
logger.warning(
|
|
f"Error reading with pyarrow.json.read_json(). "
|
|
f"Falling back to native json.load(), which may be slower. "
|
|
f"PyArrow error was:\n{e}"
|
|
)
|
|
yield from self._read_with_python_json(buffer)
|
|
|
|
|
|
class PandasJSONDatasource(FileBasedDatasource):
|
|
|
|
# Buffer size in bytes for reading files. Default is 1MB.
|
|
#
|
|
# pandas reads data in small chunks (~8 KiB), which leads to many costly
|
|
# small read requests when accessing cloud storage. To reduce overhead and
|
|
# improve performance, we wrap the file in a larger buffered reader that
|
|
# reads bigger blocks at once.
|
|
_BUFFER_SIZE = 1024**2
|
|
|
|
# In the case of zipped json files, we cannot infer the chunk_size.
|
|
_DEFAULT_CHUNK_SIZE = 10000
|
|
|
|
def __init__(
|
|
self,
|
|
paths: Union[str, List[str]],
|
|
target_output_size_bytes: int,
|
|
**file_based_datasource_kwargs,
|
|
):
|
|
super().__init__(paths, **file_based_datasource_kwargs)
|
|
|
|
self._target_output_size_bytes = target_output_size_bytes
|
|
|
|
def _read_stream(self, f: "pyarrow.NativeFile", path: str):
|
|
chunksize = self._estimate_chunksize(f)
|
|
|
|
with StrictBufferedReader(f, buffer_size=self._BUFFER_SIZE) as stream:
|
|
if chunksize is None:
|
|
# When chunksize=None, pandas returns DataFrame directly
|
|
# (no context manager).
|
|
df = pd.read_json(stream, chunksize=chunksize, lines=True)
|
|
yield _cast_range_index_to_string(df)
|
|
else:
|
|
# When chunksize is a number, pandas returns JsonReader
|
|
# (supports context manager).
|
|
with pd.read_json(stream, chunksize=chunksize, lines=True) as reader:
|
|
for df in reader:
|
|
yield _cast_range_index_to_string(df)
|
|
|
|
def _estimate_chunksize(self, f: "pyarrow.NativeFile") -> Optional[int]:
|
|
"""Estimate the chunksize by sampling the first row.
|
|
|
|
This is necessary to avoid OOMs while reading the file.
|
|
"""
|
|
|
|
if not f.seekable():
|
|
return self._DEFAULT_CHUNK_SIZE
|
|
|
|
# ``_read_stream`` can be recreated on the same file handle when
|
|
# ``FileBasedDatasource`` retries a transient read error.
|
|
f.seek(0)
|
|
|
|
if self._target_output_size_bytes is None:
|
|
return None
|
|
|
|
try:
|
|
with StrictBufferedReader(f, buffer_size=self._BUFFER_SIZE) as stream:
|
|
with pd.read_json(stream, chunksize=1, lines=True) as reader:
|
|
try:
|
|
df = _cast_range_index_to_string(next(reader))
|
|
except StopIteration:
|
|
return 1
|
|
|
|
block_accessor = PandasBlockAccessor.for_block(df)
|
|
if block_accessor.num_rows() == 0:
|
|
chunksize = 1
|
|
else:
|
|
bytes_per_row = block_accessor.size_bytes() / block_accessor.num_rows()
|
|
chunksize = max(
|
|
round(self._target_output_size_bytes / bytes_per_row), 1
|
|
)
|
|
|
|
return chunksize
|
|
finally:
|
|
# Reset file pointer to the beginning for the actual read and for any
|
|
# subsequent retry that reuses the same file handle.
|
|
f.seek(0)
|
|
|
|
def _open_input_source(
|
|
self,
|
|
filesystem: "pyarrow.fs.FileSystem",
|
|
path: str,
|
|
**open_args,
|
|
) -> "pyarrow.NativeFile":
|
|
|
|
compression = self.resolve_compression(path, open_args)
|
|
|
|
if compression is None:
|
|
# We use a seekable file to estimate chunksize.
|
|
return filesystem.open_input_file(path)
|
|
|
|
return super()._open_input_source(filesystem, path, **open_args)
|
|
|
|
|
|
def _cast_range_index_to_string(df: pd.DataFrame):
|
|
# NOTE: PandasBlockAccessor doesn't support RangeIndex, so we need to convert
|
|
# to string.
|
|
if isinstance(df.columns, pd.RangeIndex):
|
|
df.columns = df.columns.astype(str)
|
|
return df
|
|
|
|
|
|
class StrictBufferedReader(io.RawIOBase):
|
|
"""Wrapper that prevents premature file closure and ensures full-buffered reads.
|
|
|
|
This is necessary for two reasons:
|
|
1. The datasource reads the file twice -- first to sample and determine the chunk size,
|
|
and again to load the actual data. Since pandas assumes ownership of the file and
|
|
may close it, we prevent that by explicitly detaching the underlying file before
|
|
closing the buffer.
|
|
|
|
2. pandas wraps the file in a TextIOWrapper to decode bytes into text. TextIOWrapper
|
|
prefers calling read1(), which doesn't prefetch for random-access files
|
|
(e.g., from PyArrow). This wrapper forces all reads through the full buffer to
|
|
avoid inefficient small-range S3 GETs.
|
|
"""
|
|
|
|
def __init__(self, file: io.RawIOBase, buffer_size: int):
|
|
self._file = io.BufferedReader(file, buffer_size=buffer_size)
|
|
|
|
def read(self, size=-1, /):
|
|
return self._file.read(size)
|
|
|
|
def readable(self) -> bool:
|
|
return True
|
|
|
|
def close(self):
|
|
if not self.closed:
|
|
self._file.detach()
|
|
super().close()
|