wehub-resource-sync
541 lines
19 KiB
Python
541 lines
19 KiB
Python
import pickle
|
|
import random
|
|
import sys
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
from packaging.version import parse as parse_version
|
|
|
|
import ray
|
|
import ray.data
|
|
from ray.data._internal.pandas_block import (
|
|
PandasBlockAccessor,
|
|
PandasBlockBuilder,
|
|
PandasBlockColumnAccessor,
|
|
)
|
|
from ray.data._internal.util import is_null
|
|
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
|
from ray.data.context import DataContext
|
|
|
|
# Set seed for the test for size as it related to sampling
|
|
np.random.seed(42)
|
|
|
|
|
|
def simple_series(null):
|
|
return pd.Series([1, 2, null, 6])
|
|
|
|
|
|
@pytest.mark.parametrize("arr", [simple_series(None), simple_series(np.nan)])
|
|
class TestPandasBlockColumnAccessor:
|
|
@pytest.mark.parametrize(
|
|
"ignore_nulls, expected",
|
|
[
|
|
(True, 3),
|
|
(False, 4),
|
|
],
|
|
)
|
|
def test_count(self, arr, ignore_nulls, expected):
|
|
accessor = PandasBlockColumnAccessor(arr)
|
|
result = accessor.count(ignore_nulls=ignore_nulls, as_py=True)
|
|
assert result == expected
|
|
|
|
@pytest.mark.parametrize(
|
|
"ignore_nulls, expected",
|
|
[
|
|
(True, 9),
|
|
(False, np.nan),
|
|
],
|
|
)
|
|
def test_sum(self, arr, ignore_nulls, expected):
|
|
accessor = PandasBlockColumnAccessor(arr)
|
|
result = accessor.sum(ignore_nulls=ignore_nulls, as_py=True)
|
|
assert result == expected or is_null(result) and is_null(expected)
|
|
|
|
@pytest.mark.parametrize(
|
|
"ignore_nulls, expected",
|
|
[
|
|
(True, 1),
|
|
(False, np.nan),
|
|
],
|
|
)
|
|
def test_min(self, arr, ignore_nulls, expected):
|
|
accessor = PandasBlockColumnAccessor(arr)
|
|
result = accessor.min(ignore_nulls=ignore_nulls, as_py=True)
|
|
assert result == expected or is_null(result) and is_null(expected)
|
|
|
|
@pytest.mark.parametrize(
|
|
"ignore_nulls, expected",
|
|
[
|
|
(True, 6),
|
|
(False, np.nan),
|
|
],
|
|
)
|
|
def test_max(self, arr, ignore_nulls, expected):
|
|
accessor = PandasBlockColumnAccessor(arr)
|
|
result = accessor.max(ignore_nulls=ignore_nulls, as_py=True)
|
|
assert result == expected or is_null(result) and is_null(expected)
|
|
|
|
@pytest.mark.parametrize(
|
|
"ignore_nulls, expected",
|
|
[
|
|
(True, 3.0),
|
|
(False, np.nan),
|
|
],
|
|
)
|
|
def test_mean(self, arr, ignore_nulls, expected):
|
|
accessor = PandasBlockColumnAccessor(arr)
|
|
result = accessor.mean(ignore_nulls=ignore_nulls, as_py=True)
|
|
assert result == expected or is_null(result) and is_null(expected)
|
|
|
|
@pytest.mark.parametrize(
|
|
"provided_mean, expected",
|
|
[
|
|
(3.0, 14.0),
|
|
(None, 14.0),
|
|
],
|
|
)
|
|
def test_sum_of_squared_diffs_from_mean(self, arr, provided_mean, expected):
|
|
accessor = PandasBlockColumnAccessor(arr)
|
|
result = accessor.sum_of_squared_diffs_from_mean(
|
|
ignore_nulls=True, mean=provided_mean, as_py=True
|
|
)
|
|
assert result == expected or is_null(result) and is_null(expected)
|
|
|
|
def test_to_pylist(self, arr):
|
|
accessor = PandasBlockColumnAccessor(arr)
|
|
|
|
result = accessor.to_pylist()
|
|
expected = arr.to_list()
|
|
|
|
assert all(
|
|
[a == b or is_null(a) and is_null(b) for a, b in zip(expected, result)]
|
|
)
|
|
|
|
|
|
class TestPandasBlockColumnAccessorAllNullSeries:
|
|
@pytest.fixture
|
|
def all_null_series(self):
|
|
return pd.Series([None] * 3, dtype=np.float64)
|
|
|
|
def test_count_all_null(self, all_null_series):
|
|
accessor = PandasBlockColumnAccessor(all_null_series)
|
|
# When ignoring nulls, count should be 0; otherwise, count returns length.
|
|
assert accessor.count(ignore_nulls=True, as_py=True) == 0
|
|
assert accessor.count(ignore_nulls=False, as_py=True) == len(all_null_series)
|
|
|
|
@pytest.mark.parametrize("ignore_nulls", [True, False])
|
|
def test_sum_all_null(self, all_null_series, ignore_nulls):
|
|
accessor = PandasBlockColumnAccessor(all_null_series)
|
|
result = accessor.sum(ignore_nulls=ignore_nulls)
|
|
assert is_null(result)
|
|
|
|
@pytest.mark.parametrize("ignore_nulls", [True, False])
|
|
def test_min_all_null(self, all_null_series, ignore_nulls):
|
|
accessor = PandasBlockColumnAccessor(all_null_series)
|
|
result = accessor.min(ignore_nulls=ignore_nulls, as_py=True)
|
|
assert is_null(result)
|
|
|
|
@pytest.mark.parametrize("ignore_nulls", [True, False])
|
|
def test_max_all_null(self, all_null_series, ignore_nulls):
|
|
accessor = PandasBlockColumnAccessor(all_null_series)
|
|
result = accessor.max(ignore_nulls=ignore_nulls)
|
|
assert is_null(result)
|
|
|
|
@pytest.mark.parametrize("ignore_nulls", [True, False])
|
|
def test_mean_all_null(self, all_null_series, ignore_nulls):
|
|
accessor = PandasBlockColumnAccessor(all_null_series)
|
|
result = accessor.mean(ignore_nulls=ignore_nulls)
|
|
assert is_null(result)
|
|
|
|
@pytest.mark.parametrize("ignore_nulls", [True, False])
|
|
def test_sum_of_squared_diffs_all_null(self, all_null_series, ignore_nulls):
|
|
accessor = PandasBlockColumnAccessor(all_null_series)
|
|
result = accessor.sum_of_squared_diffs_from_mean(
|
|
ignore_nulls=ignore_nulls, mean=None
|
|
)
|
|
assert is_null(result)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_block, fill_column_name, fill_value, expected_output_block",
|
|
[
|
|
(
|
|
pd.DataFrame({"a": [0, 1]}),
|
|
"b",
|
|
2,
|
|
pd.DataFrame({"a": [0, 1], "b": [2, 2]}),
|
|
),
|
|
],
|
|
)
|
|
def test_fill_column(input_block, fill_column_name, fill_value, expected_output_block):
|
|
block_accessor = PandasBlockAccessor.for_block(input_block)
|
|
|
|
actual_output_block = block_accessor.fill_column(fill_column_name, fill_value)
|
|
|
|
assert actual_output_block.equals(expected_output_block)
|
|
|
|
|
|
def test_pandas_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 PandasBlockBuilder
|
|
pandas_builder = PandasBlockBuilder()
|
|
for row in data_rows:
|
|
pandas_builder.add(row)
|
|
pandas_block = pandas_builder.build()
|
|
|
|
assert pd.api.types.is_datetime64_ns_dtype(pandas_block["col2"])
|
|
|
|
for original_row, result_row in zip(
|
|
data_rows, pandas_block.to_dict(orient="records")
|
|
):
|
|
assert (
|
|
original_row["col2"] == result_row["col2"]
|
|
), "Timestamp mismatch in PandasBlockBuilder output"
|
|
|
|
|
|
def test_dict_and_none_use_arrow_block(ray_start_regular_shared, restore_data_context):
|
|
# Dicts are represented as Arrow struct types, so the block should remain Arrow
|
|
# even if object-extension fallback is disabled.
|
|
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
|
|
|
def fn(batch):
|
|
batch["data_dict"] = [{"data": 0} for _ in range(len(batch["id"]))]
|
|
return batch
|
|
|
|
ds = ray.data.range(10).map_batches(fn)
|
|
ds = ds.materialize()
|
|
block = ray.get(next(ds.iter_internal_ref_bundles()).block_refs[0])
|
|
assert isinstance(block, pa.Table)
|
|
assert block.schema.field("data_dict").type == pa.struct([("data", pa.int64())])
|
|
|
|
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(next(ds2.iter_internal_ref_bundles()).block_refs[0])
|
|
assert isinstance(block, pa.Table)
|
|
|
|
|
|
class TestSizeBytes:
|
|
def test_small(ray_start_regular_shared):
|
|
animals = ["Flamingo", "Centipede"]
|
|
block = pd.DataFrame({"animals": animals})
|
|
|
|
block_accessor = PandasBlockAccessor.for_block(block)
|
|
bytes_size = block_accessor.size_bytes()
|
|
|
|
# check that memory usage is within 10% of the size_bytes
|
|
# For strings, Pandas seems to be fairly accurate, so let's use that.
|
|
memory_usage = block.memory_usage(index=True, deep=True).sum()
|
|
assert bytes_size == pytest.approx(memory_usage, rel=0.1), (
|
|
bytes_size,
|
|
memory_usage,
|
|
)
|
|
|
|
def test_large_str(ray_start_regular_shared):
|
|
animals = [
|
|
random.choice(["alligator", "crocodile", "centipede", "flamingo"])
|
|
for i in range(100_000)
|
|
]
|
|
block = pd.DataFrame({"animals": animals})
|
|
block["animals"] = block["animals"].astype("string")
|
|
|
|
block_accessor = PandasBlockAccessor.for_block(block)
|
|
bytes_size = block_accessor.size_bytes()
|
|
|
|
memory_usage = block.memory_usage(index=True, deep=True).sum()
|
|
assert bytes_size == pytest.approx(memory_usage, rel=0.1), (
|
|
bytes_size,
|
|
memory_usage,
|
|
)
|
|
|
|
def test_large_str_object(ray_start_regular_shared):
|
|
"""Note - this test breaks if you refactor/move the list of animals."""
|
|
num = 100_000
|
|
animals = [
|
|
random.choice(["alligator", "crocodile", "centipede", "flamingo"])
|
|
for i in range(num)
|
|
]
|
|
block = pd.DataFrame({"animals": animals})
|
|
|
|
block_accessor = PandasBlockAccessor.for_block(block)
|
|
bytes_size = block_accessor.size_bytes()
|
|
|
|
# The actual usage should be the index usage + the string data usage.
|
|
memory_usage = block.memory_usage(index=True, deep=False).sum() + sum(
|
|
[sys.getsizeof(animal) for animal in animals]
|
|
)
|
|
|
|
assert bytes_size == pytest.approx(memory_usage, rel=0.1), (
|
|
bytes_size,
|
|
memory_usage,
|
|
)
|
|
|
|
def test_large_floats(ray_start_regular_shared):
|
|
animals = [random.random() for i in range(100_000)]
|
|
block = pd.DataFrame({"animals": animals})
|
|
|
|
block_accessor = PandasBlockAccessor.for_block(block)
|
|
bytes_size = block_accessor.size_bytes()
|
|
|
|
memory_usage = pickle.dumps(block).__sizeof__()
|
|
# check that memory usage is within 10% of the size_bytes
|
|
assert bytes_size == pytest.approx(memory_usage, rel=0.1), (
|
|
bytes_size,
|
|
memory_usage,
|
|
)
|
|
|
|
def test_bytes_object(ray_start_regular_shared):
|
|
def generate_data(batch):
|
|
for _ in range(8):
|
|
yield {"data": [[b"\x00" * 128 * 1024 * 128]]}
|
|
|
|
ds = (
|
|
ray.data.range(1, override_num_blocks=1)
|
|
.map_batches(generate_data, batch_size=1)
|
|
.map_batches(lambda batch: batch, batch_format="pandas")
|
|
)
|
|
|
|
true_value = 128 * 1024 * 128 * 8
|
|
for bundle in ds.iter_internal_ref_bundles():
|
|
size = bundle.size_bytes()
|
|
# assert that true_value is within 10% of bundle.size_bytes()
|
|
assert size == pytest.approx(true_value, rel=0.1), (
|
|
size,
|
|
true_value,
|
|
)
|
|
|
|
def test_nested_numpy(ray_start_regular_shared):
|
|
size = 1024
|
|
rows = 1_000
|
|
data = [
|
|
np.random.randint(size=size, low=0, high=100, dtype=np.int8)
|
|
for _ in range(rows)
|
|
]
|
|
df = pd.DataFrame({"data": data})
|
|
|
|
block_accessor = PandasBlockAccessor.for_block(df)
|
|
block_size = block_accessor.size_bytes()
|
|
true_value = rows * size
|
|
assert block_size == pytest.approx(true_value, rel=0.1), (
|
|
block_size,
|
|
true_value,
|
|
)
|
|
|
|
def test_nested_objects(ray_start_regular_shared):
|
|
size = 10
|
|
rows = 10_000
|
|
lists = [[random.randint(0, 100) for _ in range(size)] for _ in range(rows)]
|
|
data = {"lists": lists}
|
|
block = pd.DataFrame(data)
|
|
|
|
block_accessor = PandasBlockAccessor.for_block(block)
|
|
bytes_size = block_accessor.size_bytes()
|
|
|
|
# List overhead + 10 integers per list
|
|
true_size = rows * (
|
|
sys.getsizeof([random.randint(0, 100) for _ in range(size)]) + size * 28
|
|
)
|
|
|
|
assert bytes_size == pytest.approx(true_size, rel=0.1), (
|
|
bytes_size,
|
|
true_size,
|
|
)
|
|
|
|
def test_mixed_types(ray_start_regular_shared):
|
|
rows = 10_000
|
|
|
|
data = {
|
|
"integers": [random.randint(0, 100) for _ in range(rows)],
|
|
"floats": [random.random() for _ in range(rows)],
|
|
"strings": [
|
|
random.choice(["apple", "banana", "cherry"]) for _ in range(rows)
|
|
],
|
|
"object": [b"\x00" * 128 for _ in range(rows)],
|
|
}
|
|
block = pd.DataFrame(data)
|
|
block_accessor = PandasBlockAccessor.for_block(block)
|
|
bytes_size = block_accessor.size_bytes()
|
|
|
|
# Manually calculate the size
|
|
int_size = rows * 8
|
|
float_size = rows * 8
|
|
str_size = sum(sys.getsizeof(string) for string in data["strings"])
|
|
object_size = rows * sys.getsizeof(b"\x00" * 128)
|
|
|
|
true_size = int_size + float_size + str_size + object_size
|
|
assert bytes_size == pytest.approx(true_size, rel=0.1), (bytes_size, true_size)
|
|
|
|
def test_nested_lists_strings(ray_start_regular_shared):
|
|
rows = 5_000
|
|
nested_lists = ["a"] * 3 + ["bb"] * 4 + ["ccc"] * 3
|
|
data = {
|
|
"nested_lists": [nested_lists for _ in range(rows)],
|
|
}
|
|
block = pd.DataFrame(data)
|
|
block_accessor = PandasBlockAccessor.for_block(block)
|
|
bytes_size = block_accessor.size_bytes()
|
|
|
|
# Manually calculate the size
|
|
list_overhead = sys.getsizeof(block["nested_lists"].iloc[0]) + sum(
|
|
[sys.getsizeof(x) for x in nested_lists]
|
|
)
|
|
true_size = rows * list_overhead
|
|
assert bytes_size == pytest.approx(true_size, rel=0.1), (bytes_size, true_size)
|
|
|
|
@pytest.mark.parametrize("size", [10, 1024])
|
|
def test_multi_level_nesting(ray_start_regular_shared, size):
|
|
rows = 1_000
|
|
data = {
|
|
"complex": [
|
|
{"list": [np.random.rand(size)], "value": {"key": "val"}}
|
|
for _ in range(rows)
|
|
],
|
|
}
|
|
block = pd.DataFrame(data)
|
|
block_accessor = PandasBlockAccessor.for_block(block)
|
|
bytes_size = block_accessor.size_bytes()
|
|
|
|
numpy_size = np.random.rand(size).nbytes
|
|
|
|
values = ["list", "value", "key", "val"]
|
|
str_size = sum([sys.getsizeof(v) for v in values])
|
|
|
|
list_ref_overhead = sys.getsizeof([np.random.rand(size)])
|
|
|
|
dict_overhead1 = sys.getsizeof({"key": "val"})
|
|
|
|
dict_overhead3 = sys.getsizeof(
|
|
{"list": [np.random.rand(size)], "value": {"key": "val"}}
|
|
)
|
|
|
|
true_size = (
|
|
numpy_size + str_size + list_ref_overhead + dict_overhead1 + dict_overhead3
|
|
) * rows
|
|
assert bytes_size == pytest.approx(true_size, rel=0.15), (
|
|
bytes_size,
|
|
true_size,
|
|
)
|
|
|
|
def test_boolean(ray_start_regular_shared):
|
|
data = [random.choice([True, False, None]) for _ in range(100_000)]
|
|
block = pd.DataFrame({"flags": pd.Series(data, dtype="boolean")})
|
|
block_accessor = PandasBlockAccessor.for_block(block)
|
|
bytes_size = block_accessor.size_bytes()
|
|
|
|
# No object case
|
|
true_size = block.memory_usage(index=True, deep=True).sum()
|
|
assert bytes_size == pytest.approx(true_size, rel=0.1), (bytes_size, true_size)
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("10.0.1"),
|
|
reason="ArrowDtype requires pyarrow>=10.0.1",
|
|
)
|
|
def test_arrow(ray_start_regular_shared):
|
|
data = [
|
|
random.choice(["alligator", "crocodile", "flamingo"]) for _ in range(50_000)
|
|
]
|
|
arrow_dtype = pd.ArrowDtype(pa.string())
|
|
block = pd.DataFrame({"animals": pd.Series(data, dtype=arrow_dtype)})
|
|
block_accessor = PandasBlockAccessor.for_block(block)
|
|
bytes_size = block_accessor.size_bytes()
|
|
|
|
true_size = block.memory_usage(index=True, deep=False).sum() + sum(
|
|
sys.getsizeof(x) for x in data
|
|
)
|
|
assert bytes_size == pytest.approx(true_size, rel=0.1), (bytes_size, true_size)
|
|
|
|
def test_deterministic_across_blocks(self):
|
|
"""size_bytes() must return the same value for two blocks holding
|
|
identical data. Non-determinism here can cause a streaming generator
|
|
task to produce different block counts across replay attempts (e.g.
|
|
lineage reconstruction), since each attempt rebuilds the block and
|
|
re-estimates its size. That surfaces as a silent hang or silent data
|
|
loss downstream.
|
|
"""
|
|
# Use enough rows to trigger sampling (sample_size < total_size), and
|
|
# vary the string lengths so a different sample yields a different
|
|
# estimate (this is what makes the non-deterministic case observable).
|
|
data = [f"str_{i}" for i in range(10_000)]
|
|
block1 = pd.DataFrame({"col": pd.Series(data, dtype="string")})
|
|
block2 = pd.DataFrame({"col": pd.Series(data, dtype="string")})
|
|
|
|
first = PandasBlockAccessor.for_block(block1).size_bytes()
|
|
second = PandasBlockAccessor.for_block(block2).size_bytes()
|
|
|
|
assert (
|
|
first == second
|
|
), f"size_bytes() is non-deterministic: first={first}, second={second}"
|
|
|
|
|
|
def test_iter_rows_with_na(ray_start_regular_shared):
|
|
block = pd.DataFrame({"col": [pd.NA]})
|
|
block_accessor = PandasBlockAccessor.for_block(block)
|
|
|
|
rows = block_accessor.iter_rows(public_row_format=True)
|
|
|
|
# We should return None for NaN values.
|
|
assert list(rows) == [{"col": None}]
|
|
|
|
|
|
def test_empty_dataframe_with_object_columns(ray_start_regular_shared):
|
|
"""Test that size_bytes handles empty DataFrames with object/string columns.
|
|
|
|
The warning log:
|
|
"Error calculating size for column 'parent': cannot call `vectorize`
|
|
on size 0 inputs unless `otypes` is set"
|
|
should not be logged in the presence of empty columns.
|
|
"""
|
|
from unittest.mock import patch
|
|
|
|
# Create an empty DataFrame but with defined columns and dtypes
|
|
block = pd.DataFrame(
|
|
{
|
|
"parent": pd.Series([], dtype=object),
|
|
"child": pd.Series([], dtype="string"),
|
|
"data": pd.Series([], dtype=object),
|
|
}
|
|
)
|
|
|
|
block_accessor = PandasBlockAccessor.for_block(block)
|
|
|
|
# Check that NO warning is logged after calling size_bytes
|
|
with patch("ray.data._internal.pandas_block.logger.warning") as mock_warning:
|
|
bytes_size = block_accessor.size_bytes()
|
|
mock_warning.assert_not_called()
|
|
|
|
assert bytes_size >= 0
|
|
|
|
|
|
def test_tensor_column_with_all_nan_preserves_type(ray_start_regular_shared):
|
|
from ray.data.extensions import ArrowTensorType, ArrowTensorTypeV2, TensorArray
|
|
|
|
# Create a DataFrame with all-NaN tensor column
|
|
df = pd.DataFrame(
|
|
{"foo": TensorArray([np.array([np.nan, np.nan]), np.array([np.nan, np.nan])])}
|
|
)
|
|
|
|
block_accessor = PandasBlockAccessor.for_block(df)
|
|
arrow_table = block_accessor.to_arrow()
|
|
|
|
# The column should preserve tensor type
|
|
assert isinstance(
|
|
arrow_table.schema.field("foo").type, (ArrowTensorType, ArrowTensorTypeV2)
|
|
), "TensorDtype column with all-NaN values should preserve tensor type"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|