wehub-resource-sync
471 lines
16 KiB
Python
471 lines
16 KiB
Python
import gzip
|
|
import json
|
|
import os
|
|
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pyarrow.fs as fs
|
|
import pyarrow.json as pajson
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data import Schema
|
|
from ray.data._internal.datasource.json_datasource import PandasJSONDatasource
|
|
from ray.data._internal.pandas_block import PandasBlockBuilder
|
|
from ray.data._internal.util import rows_same
|
|
from ray.data.block import BlockAccessor
|
|
from ray.data.datasource.file_based_datasource import (
|
|
FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD,
|
|
)
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
# Set the test timeout to 6 minutes
|
|
pytestmark = pytest.mark.timeout(360)
|
|
|
|
|
|
def test_json_read(
|
|
ray_start_regular_shared, target_max_block_size_infinite_or_default, tmp_path
|
|
):
|
|
df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
|
path1 = os.path.join(tmp_path, "test1.json")
|
|
df1.to_json(path1, orient="records", lines=True)
|
|
ds = ray.data.read_json(path1)
|
|
dsdf = ds.to_pandas()
|
|
pd.testing.assert_frame_equal(df1.astype(dsdf.dtypes.to_dict()), dsdf)
|
|
# Metadata ops.
|
|
assert ds.count() == 3
|
|
assert ds.input_files() == [path1]
|
|
assert ds.schema() == Schema(pa.schema([("one", pa.int64()), ("two", pa.string())]))
|
|
|
|
|
|
def test_zipped_json_read(
|
|
ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
|
|
):
|
|
df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
|
path1 = os.path.join(tmp_path, "test1.json.gz")
|
|
df1.to_json(path1, compression="gzip", orient="records", lines=True)
|
|
ds = ray.data.read_json(path1)
|
|
dsdf = ds.to_pandas()
|
|
pd.testing.assert_frame_equal(df1.astype(dsdf.dtypes.to_dict()), dsdf)
|
|
# Metadata ops.
|
|
assert ds.count() == 3
|
|
assert ds.input_files() == [path1]
|
|
|
|
|
|
def test_read_json_fallback_from_pyarrow_failure(
|
|
ray_start_regular_shared, local_path, target_max_block_size_infinite_or_default
|
|
):
|
|
# Try to read this with read_json() to trigger fallback logic
|
|
# to read bytes with json.load().
|
|
data = [{"one": [1]}, {"one": [1, 2]}]
|
|
path1 = os.path.join(local_path, "test1.json")
|
|
with open(path1, "w") as f:
|
|
json.dump(data, f)
|
|
|
|
# pyarrow.json cannot read JSONs containing arrays of different lengths.
|
|
from pyarrow import ArrowInvalid
|
|
|
|
with pytest.raises(ArrowInvalid):
|
|
pajson.read_json(path1)
|
|
|
|
# Ray Data successfully reads this in by
|
|
# falling back to json.load() when pyarrow fails.
|
|
ds = ray.data.read_json(path1)
|
|
assert ds.take_all() == data
|
|
|
|
|
|
def test_json_read_with_read_options(
|
|
ray_start_regular_shared,
|
|
tmp_path,
|
|
target_max_block_size_infinite_or_default,
|
|
):
|
|
# Arrow's JSON ReadOptions isn't serializable in pyarrow < 8.0.0, so this test
|
|
# covers our custom ReadOptions serializer.
|
|
# TODO(Clark): Remove this test and our custom serializer once we require
|
|
# pyarrow >= 8.0.0.
|
|
|
|
df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
|
path1 = os.path.join(tmp_path, "test1.json")
|
|
df1.to_json(path1, orient="records", lines=True)
|
|
ds = ray.data.read_json(
|
|
path1,
|
|
read_options=pajson.ReadOptions(use_threads=False, block_size=2**30),
|
|
)
|
|
dsdf = ds.to_pandas()
|
|
pd.testing.assert_frame_equal(df1.astype(dsdf.dtypes.to_dict()), dsdf)
|
|
# Test metadata ops.
|
|
assert ds.count() == 3
|
|
assert ds.input_files() == [path1]
|
|
assert ds.schema() == Schema(pa.schema([("one", pa.int64()), ("two", pa.string())]))
|
|
|
|
|
|
def test_json_read_with_parse_options(
|
|
ray_start_regular_shared,
|
|
tmp_path,
|
|
target_max_block_size_infinite_or_default,
|
|
):
|
|
# Arrow's JSON ParseOptions isn't serializable in pyarrow < 8.0.0, so this test
|
|
# covers our custom ParseOptions serializer, similar to ReadOptions in above test.
|
|
# TODO(chengsu): Remove this test and our custom serializer once we require
|
|
# pyarrow >= 8.0.0.
|
|
|
|
df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
|
path1 = os.path.join(tmp_path, "test1.json")
|
|
df1.to_json(path1, orient="records", lines=True)
|
|
ds = ray.data.read_json(
|
|
path1,
|
|
parse_options=pajson.ParseOptions(
|
|
explicit_schema=pa.schema([("two", pa.string())]),
|
|
unexpected_field_behavior="ignore",
|
|
),
|
|
)
|
|
dsdf = ds.to_pandas()
|
|
assert len(dsdf.columns) == 1
|
|
pd.testing.assert_series_equal(df1["two"].astype(dsdf["two"].dtype), dsdf["two"])
|
|
# Test metadata ops.
|
|
assert ds.count() == 3
|
|
assert ds.input_files() == [path1]
|
|
assert ds.schema() == Schema(pa.schema([("two", pa.string())]))
|
|
|
|
|
|
@pytest.mark.parametrize("override_num_blocks", [None, 1, 3])
|
|
def test_jsonl_lists(
|
|
ray_start_regular_shared,
|
|
tmp_path,
|
|
override_num_blocks,
|
|
target_max_block_size_infinite_or_default,
|
|
):
|
|
"""Test JSONL with mixed types and schemas."""
|
|
data = [
|
|
["ray", "rocks", "hello"],
|
|
["oh", "no"],
|
|
["rocking", "with", "ray"],
|
|
]
|
|
|
|
path = os.path.join(tmp_path, "test.jsonl")
|
|
with open(path, "w") as f:
|
|
for record in data:
|
|
json.dump(record, f)
|
|
f.write("\n")
|
|
|
|
ds = ray.data.read_json(path, lines=True, override_num_blocks=override_num_blocks)
|
|
result = ds.take_all()
|
|
|
|
assert result[0] == {"0": "ray", "1": "rocks", "2": "hello"}
|
|
assert result[1] == {"0": "oh", "1": "no", "2": None}
|
|
assert result[2] == {"0": "rocking", "1": "with", "2": "ray"}
|
|
|
|
|
|
def test_jsonl_mixed_types(
|
|
ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
|
|
):
|
|
"""Test JSONL with mixed types and schemas."""
|
|
data = [
|
|
{"a": 1, "b": {"c": 2}}, # Nested dict
|
|
{"a": 1, "b": {"c": 3}}, # Nested dict
|
|
{"a": 1, "b": {"c": {"hello": "world"}}}, # Mixed Schema
|
|
]
|
|
|
|
path = os.path.join(tmp_path, "test.jsonl")
|
|
with open(path, "w") as f:
|
|
for record in data:
|
|
json.dump(record, f)
|
|
f.write("\n")
|
|
|
|
ds = ray.data.read_json(path, lines=True)
|
|
result = ds.take_all()
|
|
|
|
assert result[0] == data[0] # Dict stays as is
|
|
assert result[1] == data[1]
|
|
assert result[2] == data[2]
|
|
|
|
|
|
def test_json_write(
|
|
ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
|
|
):
|
|
input_df = pd.DataFrame({"id": [0]})
|
|
ds = ray.data.from_blocks([input_df])
|
|
|
|
ds.write_json(tmp_path)
|
|
|
|
output_df = pd.concat(
|
|
[
|
|
pd.read_json(os.path.join(tmp_path, filename), lines=True)
|
|
for filename in os.listdir(tmp_path)
|
|
]
|
|
)
|
|
|
|
assert rows_same(input_df, output_df)
|
|
|
|
|
|
@pytest.mark.parametrize("override_num_blocks", [None, 2])
|
|
def test_json_roundtrip(
|
|
ray_start_regular_shared,
|
|
tmp_path,
|
|
override_num_blocks,
|
|
target_max_block_size_infinite_or_default,
|
|
):
|
|
df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
|
|
|
ds = ray.data.from_pandas([df], override_num_blocks=override_num_blocks)
|
|
ds.write_json(tmp_path)
|
|
|
|
ds2 = ray.data.read_json(tmp_path)
|
|
ds2df = ds2.to_pandas()
|
|
assert rows_same(ds2df, df)
|
|
for entry in ds2._execute().blocks:
|
|
assert (
|
|
# pyrefly: ignore[no-matching-overload]
|
|
BlockAccessor.for_block(ray.get(entry.ref)).size_bytes()
|
|
== entry.metadata.size_bytes
|
|
)
|
|
|
|
|
|
def test_json_read_small_file_unit_block_size(
|
|
ray_start_regular_shared,
|
|
tmp_path,
|
|
target_max_block_size_infinite_or_default,
|
|
):
|
|
"""Test reading a small JSON file with unit block_size."""
|
|
|
|
df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
|
path1 = os.path.join(tmp_path, "test1.json")
|
|
df1.to_json(path1, orient="records", lines=True)
|
|
ds = ray.data.read_json(path1, read_options=pajson.ReadOptions(block_size=1))
|
|
dsdf = ds.to_pandas()
|
|
pd.testing.assert_frame_equal(df1.astype(dsdf.dtypes.to_dict()), dsdf)
|
|
# Test metadata ops.
|
|
assert ds.count() == 3
|
|
assert ds.input_files() == [path1]
|
|
assert ds.schema() == Schema(pa.schema([("one", pa.int64()), ("two", pa.string())]))
|
|
|
|
|
|
def test_json_read_file_larger_than_block_size(
|
|
ray_start_regular_shared,
|
|
tmp_path,
|
|
target_max_block_size_infinite_or_default,
|
|
):
|
|
"""Test reading a JSON file larger than the block size."""
|
|
block_size = 1024
|
|
num_chars = 2500
|
|
num_rows = 3
|
|
df2 = pd.DataFrame(
|
|
{
|
|
"one": ["a" * num_chars for _ in range(num_rows)],
|
|
"two": ["b" * num_chars for _ in range(num_rows)],
|
|
}
|
|
)
|
|
path2 = os.path.join(tmp_path, "test2.json")
|
|
df2.to_json(path2, orient="records", lines=True)
|
|
ds = ray.data.read_json(
|
|
path2, read_options=pajson.ReadOptions(block_size=block_size)
|
|
)
|
|
dsdf = ds.to_pandas()
|
|
pd.testing.assert_frame_equal(df2.astype(dsdf.dtypes.to_dict()), dsdf)
|
|
# Test metadata ops.
|
|
assert ds.count() == num_rows
|
|
assert ds.input_files() == [path2]
|
|
assert ds.schema() == Schema(
|
|
pa.schema([("one", pa.string()), ("two", pa.string())])
|
|
)
|
|
|
|
|
|
def test_json_read_negative_block_size_fallback(
|
|
ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
|
|
):
|
|
"""Test reading JSON with negative block_size triggers fallback to json.load()."""
|
|
|
|
df3 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
|
path3 = os.path.join(tmp_path, "test3.json")
|
|
df3.to_json(path3, orient="records", lines=True)
|
|
|
|
# Negative Buffer Size, fails with arrow but succeeds in fallback to json.load()
|
|
ds = ray.data.read_json(path3, read_options=pajson.ReadOptions(block_size=-1))
|
|
dsdf = ds.to_pandas()
|
|
pd.testing.assert_frame_equal(df3.astype(dsdf.dtypes.to_dict()), dsdf)
|
|
|
|
|
|
def test_json_read_zero_block_size_failure(
|
|
ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
|
|
):
|
|
"""Test reading JSON with zero block_size fails in both arrow and fallback."""
|
|
|
|
df3 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
|
path3 = os.path.join(tmp_path, "test3.json")
|
|
df3.to_json(path3, orient="records", lines=True)
|
|
|
|
# Zero Buffer Size, fails with arrow and fails in fallback to json.load()
|
|
with pytest.raises(json.decoder.JSONDecodeError, match="Extra data"):
|
|
ds = ray.data.read_json(path3, read_options=pajson.ReadOptions(block_size=0))
|
|
dsdf = ds.to_pandas()
|
|
assert dsdf.equals(df3)
|
|
|
|
|
|
@pytest.mark.parametrize("min_rows_per_file", [5, 10, 50])
|
|
def test_write_min_rows_per_file(
|
|
tmp_path,
|
|
ray_start_regular_shared,
|
|
min_rows_per_file,
|
|
target_max_block_size_infinite_or_default,
|
|
):
|
|
ray.data.range(100, override_num_blocks=20).write_json(
|
|
tmp_path, min_rows_per_file=min_rows_per_file
|
|
)
|
|
|
|
for filename in os.listdir(tmp_path):
|
|
with open(os.path.join(tmp_path, filename), "r") as file:
|
|
num_rows_written = len(file.read().splitlines())
|
|
assert num_rows_written == min_rows_per_file
|
|
|
|
|
|
def test_mixed_gzipped_json_files(
|
|
ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
|
|
):
|
|
# Create a non-empty gzipped JSON file
|
|
non_empty_file_path = os.path.join(tmp_path, "non_empty.json.gz")
|
|
data = [{"col1": "value1", "col2": "value2", "col3": "value3"}]
|
|
with gzip.open(non_empty_file_path, "wt", encoding="utf-8") as f:
|
|
for record in data:
|
|
json.dump(record, f)
|
|
f.write("\n")
|
|
|
|
# Create an empty gzipped JSON file
|
|
empty_file_path = os.path.join(tmp_path, "empty.json.gz")
|
|
with gzip.open(empty_file_path, "wt", encoding="utf-8"):
|
|
pass # Write nothing to create an empty file
|
|
|
|
# Attempt to read both files with Ray
|
|
ds = ray.data.read_json(
|
|
[non_empty_file_path, empty_file_path],
|
|
arrow_open_stream_args={"compression": "gzip"},
|
|
)
|
|
|
|
# The dataset should only contain data from the non-empty file
|
|
assert ds.count() == 1
|
|
# Iterate through each row in the dataset and compare with the expected data
|
|
for row in ds.iter_rows():
|
|
assert row == data[0], f"Row {row} does not match expected {data[0]}"
|
|
|
|
# Verify the data content using take
|
|
retrieved_data = ds.take(1)[0]
|
|
assert (
|
|
retrieved_data == data[0]
|
|
), f"Retrieved data {retrieved_data} does not match expected {data[0]}."
|
|
|
|
|
|
def test_json_with_http_path_parallelization(
|
|
ray_start_regular_shared, httpserver, target_max_block_size_infinite_or_default
|
|
):
|
|
num_files = FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD
|
|
urls = []
|
|
for i in range(num_files):
|
|
httpserver.expect_request(f"/file{i}.json").respond_with_json({"id": i})
|
|
urls.append(httpserver.url_for(f"/file{i}.json"))
|
|
|
|
ds = ray.data.read_json(urls)
|
|
actual_rows = ds.take_all()
|
|
|
|
expected_rows = [{"id": i} for i in range(num_files)]
|
|
assert sorted(actual_rows, key=lambda row: row["id"]) == sorted(
|
|
expected_rows, key=lambda row: row["id"]
|
|
)
|
|
|
|
|
|
class TestPandasJSONDatasource:
|
|
@pytest.mark.parametrize(
|
|
"data",
|
|
[{"a": []}, {"a": [1]}, {"a": [1, 2, 3]}],
|
|
ids=["empty", "single", "multiple"],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"compression,filename",
|
|
[("gzip", "test.json.gz"), ("infer", "test.json")], # infer = default
|
|
)
|
|
def test_read_stream(
|
|
self,
|
|
data,
|
|
tmp_path,
|
|
compression,
|
|
filename,
|
|
target_max_block_size_infinite_or_default,
|
|
):
|
|
# Setup test file.
|
|
df = pd.DataFrame(data)
|
|
path = os.path.join(tmp_path, filename)
|
|
df.to_json(path, orient="records", lines=True, compression=compression)
|
|
|
|
# Setup datasource.
|
|
local_filesystem = fs.LocalFileSystem()
|
|
source = PandasJSONDatasource(
|
|
path, target_output_size_bytes=1, filesystem=local_filesystem
|
|
)
|
|
|
|
# Read stream.
|
|
block_builder = PandasBlockBuilder()
|
|
with source._open_input_source(local_filesystem, path) as f:
|
|
for block in source._read_stream(f, path):
|
|
block_builder.add_block(block)
|
|
block = block_builder.build()
|
|
|
|
# Verify.
|
|
assert rows_same(block, df)
|
|
|
|
def test_read_stream_with_target_output_size_bytes(
|
|
self, tmp_path, target_max_block_size_infinite_or_default
|
|
):
|
|
# Setup test file. It contains 16 lines, each line is 8 MiB.
|
|
df = pd.DataFrame({"data": ["a" * 8 * 1024 * 1024] * 16})
|
|
path = os.path.join(tmp_path, "test.json")
|
|
df.to_json(path, orient="records", lines=True)
|
|
|
|
# Setup datasource. It should read 32 MiB (4 lines) per output.
|
|
local_filesystem = fs.LocalFileSystem()
|
|
source = PandasJSONDatasource(
|
|
path,
|
|
target_output_size_bytes=32 * 1024 * 1024,
|
|
filesystem=local_filesystem,
|
|
)
|
|
|
|
# Read stream.
|
|
block_builder = PandasBlockBuilder()
|
|
with source._open_input_source(local_filesystem, path) as f:
|
|
for block in source._read_stream(f, path):
|
|
assert len(block) == 4
|
|
block_builder.add_block(block)
|
|
block = block_builder.build()
|
|
|
|
# Verify.
|
|
assert rows_same(block, df)
|
|
|
|
def test_read_stream_with_advanced_file_pointer(
|
|
self, tmp_path, target_max_block_size_infinite_or_default
|
|
):
|
|
# Setup test file.
|
|
df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
|
path = os.path.join(tmp_path, "test.json")
|
|
df.to_json(path, orient="records", lines=True)
|
|
|
|
# Setup datasource.
|
|
local_filesystem = fs.LocalFileSystem()
|
|
source = PandasJSONDatasource(
|
|
path, target_output_size_bytes=1, filesystem=local_filesystem
|
|
)
|
|
|
|
# Simulate retrying a stream read on a file handle that was already consumed.
|
|
block_builder = PandasBlockBuilder()
|
|
with source._open_input_source(local_filesystem, path) as f:
|
|
f.read(1)
|
|
for block in source._read_stream(f, path):
|
|
block_builder.add_block(block)
|
|
block = block_builder.build()
|
|
|
|
# Verify.
|
|
assert rows_same(block, df)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|