wehub-resource-sync
352 lines
11 KiB
Python
352 lines
11 KiB
Python
import base64
|
|
import os
|
|
import sys
|
|
import types
|
|
from decimal import Decimal
|
|
from tempfile import TemporaryDirectory
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
from pyarrow import ArrowInvalid
|
|
|
|
import ray
|
|
from ray._common.test_utils import run_string_as_driver
|
|
from ray.data._internal.arrow_block import (
|
|
ArrowBlockAccessor,
|
|
ArrowBlockBuilder,
|
|
)
|
|
from ray.data._internal.arrow_ops.transform_pyarrow import combine_chunked_array
|
|
from ray.data._internal.util import GiB, MiB
|
|
from ray.data.block import BlockAccessor
|
|
from ray.data.context import DataContext
|
|
|
|
|
|
def test_combine_chunked_fixed_width_array_large():
|
|
"""Verifies `combine_chunked_array` on fixed-width arrays > 2 GiB, produces
|
|
single contiguous PA Array"""
|
|
|
|
# 144 MiB
|
|
ones_1gb = np.ones(shape=(550, 128, 128, 4), dtype=np.int32()).ravel()
|
|
|
|
# Total ~2.15 GiB
|
|
input_ = pa.chunked_array(
|
|
[
|
|
pa.array(ones_1gb),
|
|
]
|
|
* 16
|
|
)
|
|
|
|
assert round(input_.nbytes / GiB, 2) == 2.15
|
|
|
|
result = combine_chunked_array(input_)
|
|
|
|
assert isinstance(result, pa.Int32Array)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"array_type,input_factory",
|
|
[
|
|
(
|
|
pa.binary(),
|
|
lambda num_bytes: np.arange(num_bytes, dtype=np.uint8).tobytes(),
|
|
),
|
|
(
|
|
pa.string(),
|
|
lambda num_bytes: base64.encodebytes(
|
|
np.arange(num_bytes, dtype=np.int8).tobytes()
|
|
).decode("ascii"),
|
|
),
|
|
(pa.list_(pa.uint8()), lambda num_bytes: np.arange(num_bytes, dtype=np.uint8)),
|
|
],
|
|
)
|
|
def test_combine_chunked_variable_width_array_large(array_type, input_factory):
|
|
"""Verifies `combine_chunked_array` on variable-width arrays > 2 GiB,
|
|
safely produces new ChunkedArray with provided chunks recombined into
|
|
larger ones up to INT32_MAX in size"""
|
|
|
|
one_half_gb_arr = pa.array([input_factory(GiB / 2)], type=array_type)
|
|
chunked_arr = pa.chunked_array(
|
|
[one_half_gb_arr, one_half_gb_arr, one_half_gb_arr, one_half_gb_arr]
|
|
)
|
|
|
|
# 2 GiB + offsets (4 x int32)
|
|
num_bytes = chunked_arr.nbytes
|
|
expected_num_bytes = 4 * one_half_gb_arr.nbytes
|
|
|
|
num_chunks = len(chunked_arr.chunks)
|
|
assert num_chunks == 4
|
|
assert num_bytes == expected_num_bytes
|
|
|
|
# Assert attempt to combine directly fails
|
|
with pytest.raises(ArrowInvalid):
|
|
chunked_arr.combine_chunks()
|
|
|
|
# Safe combination succeeds by avoiding overflowing combination
|
|
combined = combine_chunked_array(chunked_arr)
|
|
|
|
num_bytes = combined.nbytes
|
|
|
|
num_chunks = len(combined.chunks)
|
|
assert num_chunks == 2
|
|
assert num_bytes == expected_num_bytes
|
|
|
|
|
|
def test_add_rows_with_different_column_names(ray_start_regular_shared):
|
|
builder = ArrowBlockBuilder()
|
|
|
|
builder.add({"col1": "spam"})
|
|
builder.add({"col2": "foo"})
|
|
block = builder.build()
|
|
|
|
expected_table = pa.Table.from_pydict(
|
|
{"col1": ["spam", None], "col2": [None, "foo"]}
|
|
)
|
|
assert block.equals(expected_table)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def binary_dataset_single_file_gt_2gb():
|
|
total_size = int(2.1 * GiB)
|
|
chunk_size = 256 * MiB
|
|
num_chunks = total_size // chunk_size
|
|
remainder = total_size % chunk_size
|
|
|
|
with TemporaryDirectory() as tmp_dir:
|
|
dataset_path = f"{tmp_dir}/binary_dataset_gt_2gb_single_file"
|
|
|
|
# Create directory
|
|
os.mkdir(dataset_path)
|
|
|
|
with open(f"{dataset_path}/chunk.bin", "wb") as f:
|
|
for i in range(num_chunks):
|
|
f.write(b"a" * chunk_size)
|
|
|
|
print(f">>> Written chunk #{i}")
|
|
|
|
if remainder:
|
|
f.write(b"a" * remainder)
|
|
|
|
print(f">>> Wrote chunked dataset at: {dataset_path}")
|
|
|
|
yield dataset_path, total_size
|
|
|
|
print(f">>> Cleaning up dataset: {dataset_path}")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"col_name",
|
|
[
|
|
"bytes",
|
|
# TODO fix numpy conversion
|
|
# "text",
|
|
],
|
|
)
|
|
def test_single_row_gt_2gb(
|
|
ray_start_regular_shared,
|
|
restore_data_context,
|
|
binary_dataset_single_file_gt_2gb,
|
|
col_name,
|
|
):
|
|
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
dataset_path, target_binary_size = binary_dataset_single_file_gt_2gb
|
|
|
|
def _id(row):
|
|
bs = row[col_name]
|
|
assert round(len(bs) / GiB, 1) == round(target_binary_size / GiB, 1)
|
|
return row
|
|
|
|
if col_name == "text":
|
|
ds = ray.data.read_text(dataset_path)
|
|
elif col_name == "bytes":
|
|
ds = ray.data.read_binary_files(dataset_path)
|
|
|
|
total = ds.map(_id).count()
|
|
|
|
assert total == 1
|
|
|
|
|
|
def test_random_shuffle(ray_start_regular_shared):
|
|
TOTAL_ROWS = 10000
|
|
table = pa.table({"id": pa.array(range(TOTAL_ROWS))})
|
|
block_accessor = ArrowBlockAccessor(table)
|
|
|
|
# Perform the random shuffle
|
|
shuffled_table = block_accessor.random_shuffle(random_seed=None)
|
|
assert shuffled_table.num_rows == TOTAL_ROWS
|
|
|
|
# Access the shuffled data
|
|
block_accessor = ArrowBlockAccessor(shuffled_table)
|
|
shuffled_data = block_accessor.to_pandas()["id"].tolist()
|
|
original_data = list(range(TOTAL_ROWS))
|
|
|
|
# Ensure the shuffled data is not identical to the original
|
|
assert (
|
|
shuffled_data != original_data
|
|
), "Shuffling should result in a different order"
|
|
|
|
# Ensure the entire set of original values is still in the shuffled dataset
|
|
assert (
|
|
sorted(shuffled_data) == original_data
|
|
), "The shuffled data should contain all the original values"
|
|
|
|
|
|
def test_register_arrow_types(ray_start_regular_shared, tmp_path):
|
|
# Test that our custom arrow extension types are registered on initialization.
|
|
ds = ray.data.from_items(np.zeros((8, 8, 8), dtype=np.int64))
|
|
tmp_file = f"{tmp_path}/test.parquet"
|
|
ds.write_parquet(tmp_file)
|
|
|
|
ds = ray.data.read_parquet(tmp_file)
|
|
schema = "Column Type\n------ ----\nitem ArrowTensorTypeV2(shape=(8, 8), dtype=int64)"
|
|
assert str(ds.schema()) == schema
|
|
|
|
# Also run in driver script to eliminate existing imports.
|
|
driver_script = """import ray
|
|
ds = ray.data.read_parquet("{0}")
|
|
schema = ds.schema()
|
|
assert str(schema) == \"\"\"{1}\"\"\"
|
|
""".format(
|
|
tmp_file, schema
|
|
)
|
|
run_string_as_driver(driver_script)
|
|
|
|
|
|
def test_dict_doesnt_fallback_to_pandas_block(ray_start_regular_shared):
|
|
# If the UDF returns a column with dict, previously, we would
|
|
# fall back to pandas, because we couldn't convert it to
|
|
# an Arrow block. This test checks that the block
|
|
# construction now correctly goes to Arrow.
|
|
def fn(batch):
|
|
batch["data_dict"] = [{"data": 0} for _ in range(len(batch["id"]))]
|
|
batch["data_objects"] = [
|
|
types.SimpleNamespace(a=1, b="test") for _ in range(len(batch["id"]))
|
|
]
|
|
return batch
|
|
|
|
ds = ray.data.range(10).map_batches(fn)
|
|
ds = ds.materialize()
|
|
block = ray.get(ds.get_internal_block_refs()[0])
|
|
assert isinstance(block, pa.Table), type(block)
|
|
df_from_block = block.to_pandas()
|
|
assert df_from_block["data_dict"].iloc[0] == {"data": 0}
|
|
assert df_from_block["data_objects"].iloc[0] == types.SimpleNamespace(a=1, b="test")
|
|
|
|
def fn2(batch):
|
|
batch["data_none"] = [None for _ in range(len(batch["id"]))]
|
|
return batch
|
|
|
|
ds2 = ray.data.range(10).map_batches(fn2)
|
|
ds2 = ds2.materialize()
|
|
block = ray.get(ds2.get_internal_block_refs()[0])
|
|
assert isinstance(block, pa.Table), type(block)
|
|
df_from_block = block.to_pandas()
|
|
assert df_from_block["data_none"].iloc[0] is None
|
|
|
|
|
|
# Test for https://github.com/ray-project/ray/issues/49338.
|
|
def test_build_block_with_null_column(ray_start_regular_shared, restore_data_context):
|
|
ctx = DataContext.get_current()
|
|
ctx.execution_options.preserve_order = True
|
|
|
|
# The blocks need to contain a tensor column to trigger the bug.
|
|
block1 = BlockAccessor.batch_to_block(
|
|
{"string": [None], "array": np.zeros((1, 2, 2))}
|
|
)
|
|
block2 = BlockAccessor.batch_to_block(
|
|
{"string": ["spam"], "array": np.zeros((1, 2, 2))}
|
|
)
|
|
|
|
builder = ArrowBlockBuilder()
|
|
builder.add_block(block1)
|
|
builder.add_block(block2)
|
|
block = builder.build()
|
|
|
|
rows = list(BlockAccessor.for_block(block).iter_rows(True))
|
|
assert len(rows) == 2
|
|
|
|
assert rows[0]["string"] is None
|
|
assert rows[1]["string"] == "spam"
|
|
assert np.array_equal(rows[0]["array"], np.zeros((2, 2)))
|
|
assert np.array_equal(rows[1]["array"], np.zeros((2, 2)))
|
|
|
|
|
|
def test_arrow_block_timestamp_ns(ray_start_regular_shared):
|
|
# Input data with nanosecond precision timestamps
|
|
data_rows = [
|
|
{"col1": 1, "col2": pd.Timestamp("2023-01-01T00:00:00.123456789")},
|
|
{"col1": 2, "col2": pd.Timestamp("2023-01-01T01:15:30.987654321")},
|
|
{"col1": 3, "col2": pd.Timestamp("2023-01-01T02:30:15.111111111")},
|
|
{"col1": 4, "col2": pd.Timestamp("2023-01-01T03:45:45.222222222")},
|
|
{"col1": 5, "col2": pd.Timestamp("2023-01-01T05:00:00.333333333")},
|
|
]
|
|
|
|
# Initialize ArrowBlockBuilder
|
|
arrow_builder = ArrowBlockBuilder()
|
|
for row in data_rows:
|
|
arrow_builder.add(row)
|
|
arrow_block = arrow_builder.build()
|
|
|
|
assert arrow_block.schema.field("col2").type == pa.timestamp("ns")
|
|
for i, row in enumerate(data_rows):
|
|
result_timestamp = arrow_block["col2"][i].as_py()
|
|
# Convert both values to pandas Timestamp to preserve nanosecond precision for
|
|
# comparison.
|
|
assert pd.Timestamp(row["col2"]) == pd.Timestamp(
|
|
result_timestamp
|
|
), f"Timestamp mismatch at row {i} in ArrowBlockBuilder output"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_array,transform,expected_type,expected_values",
|
|
[
|
|
(
|
|
pa.array([None, None], type=pa.string()),
|
|
None,
|
|
pa.string(),
|
|
[None, None],
|
|
),
|
|
(
|
|
pa.array([None, None], type=pa.list_(pa.string())),
|
|
None,
|
|
pa.list_(pa.string()),
|
|
[None, None],
|
|
),
|
|
(
|
|
pa.array([None, None], type=pa.decimal128(10, 2)),
|
|
lambda df: df.fillna({"x": 0}),
|
|
pa.decimal128(10, 2),
|
|
[Decimal("0.00"), Decimal("0.00")],
|
|
),
|
|
(
|
|
pa.array([["a", "b"], None], type=pa.list_(pa.string())),
|
|
None,
|
|
pa.list_(pa.string()),
|
|
[["a", "b"], None],
|
|
),
|
|
],
|
|
)
|
|
def test_arrow_block_to_pandas_preserves_arrow_types_through_roundtrip(
|
|
input_array, transform, expected_type, expected_values
|
|
):
|
|
table = pa.table({"x": input_array})
|
|
|
|
df = ArrowBlockAccessor(table).to_pandas()
|
|
assert isinstance(df.dtypes["x"], pd.ArrowDtype)
|
|
assert df.dtypes["x"].pyarrow_dtype == expected_type
|
|
|
|
if transform is not None:
|
|
df = transform(df)
|
|
|
|
roundtripped = BlockAccessor.for_block(df).to_arrow()
|
|
|
|
assert roundtripped.schema.field("x").type == expected_type
|
|
assert roundtripped.to_pydict() == {"x": expected_values}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|