wehub-resource-sync
314 lines
11 KiB
Python
314 lines
11 KiB
Python
"""Integration tests for struct namespace expressions.
|
|
|
|
These tests require Ray and test end-to-end struct namespace expression evaluation.
|
|
"""
|
|
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
from packaging import version
|
|
|
|
import ray
|
|
from ray.data._internal.util import rows_same
|
|
from ray.data.expressions import col
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
version.parse(pa.__version__) < version.parse("19.0.0"),
|
|
reason="Namespace expressions tests require PyArrow >= 19.0",
|
|
)
|
|
|
|
|
|
def _create_dataset(items_data, dataset_format, arrow_table=None):
|
|
if dataset_format == "arrow":
|
|
if arrow_table is not None:
|
|
ds = ray.data.from_arrow(arrow_table)
|
|
else:
|
|
table = pa.Table.from_pylist(items_data)
|
|
ds = ray.data.from_arrow(table)
|
|
elif dataset_format == "pandas":
|
|
if arrow_table is not None:
|
|
df = arrow_table.to_pandas()
|
|
else:
|
|
df = pd.DataFrame(items_data)
|
|
ds = ray.data.from_blocks([df])
|
|
return ds
|
|
|
|
|
|
DATASET_FORMATS = ["pandas", "arrow"]
|
|
|
|
|
|
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
|
class TestStructNamespace:
|
|
"""Tests for struct namespace operations."""
|
|
|
|
def test_struct_bracket_bool_index_raises(self, dataset_format):
|
|
"""Test struct[bool] raises TypeError instead of being treated as int."""
|
|
del dataset_format # Unused, required by class-level parametrization.
|
|
|
|
with pytest.raises(
|
|
TypeError, match="Struct indices must be strings or integers"
|
|
):
|
|
col("user").struct[True]
|
|
|
|
@pytest.mark.parametrize("bad_index", ["1", True])
|
|
def test_struct_field_by_index_non_integer_raises(self, dataset_format, bad_index):
|
|
"""Test struct.field_by_index() rejects non-integer indices."""
|
|
del dataset_format # Unused, required by class-level parametrization.
|
|
|
|
with pytest.raises(TypeError, match="Struct field index must be an integer"):
|
|
col("user").struct.field_by_index(bad_index)
|
|
|
|
def test_struct_field_by_index_negative_raises(self, dataset_format):
|
|
"""Test struct.field_by_index() rejects negative indices."""
|
|
del dataset_format # Unused, required by class-level parametrization.
|
|
|
|
with pytest.raises(
|
|
ValueError, match="Struct field index must be non-negative, got -1"
|
|
):
|
|
col("user").struct.field_by_index(-1)
|
|
|
|
with pytest.raises(
|
|
ValueError, match="Struct field index must be non-negative, got -1"
|
|
):
|
|
col("user").struct[-1]
|
|
|
|
def test_struct_field(self, ray_start_regular_shared, dataset_format):
|
|
"""Test struct.field() extracts field."""
|
|
arrow_table = pa.table(
|
|
{
|
|
"user": pa.array(
|
|
[
|
|
{"name": "Alice", "age": 30},
|
|
{"name": "Bob", "age": 25},
|
|
],
|
|
type=pa.struct(
|
|
[
|
|
pa.field("name", pa.string()),
|
|
pa.field("age", pa.int32()),
|
|
]
|
|
),
|
|
)
|
|
}
|
|
)
|
|
items_data = [
|
|
{"user": {"name": "Alice", "age": 30}},
|
|
{"user": {"name": "Bob", "age": 25}},
|
|
]
|
|
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
|
|
|
result = ds.with_column("age", col("user").struct.field("age")).to_pandas()
|
|
expected = pd.DataFrame(
|
|
{
|
|
"user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
|
|
"age": [30, 25],
|
|
}
|
|
)
|
|
assert rows_same(result, expected)
|
|
|
|
def test_struct_bracket(self, ray_start_regular_shared, dataset_format):
|
|
"""Test struct['field'] bracket notation."""
|
|
arrow_table = pa.table(
|
|
{
|
|
"user": pa.array(
|
|
[
|
|
{"name": "Alice", "age": 30},
|
|
{"name": "Bob", "age": 25},
|
|
],
|
|
type=pa.struct(
|
|
[
|
|
pa.field("name", pa.string()),
|
|
pa.field("age", pa.int32()),
|
|
]
|
|
),
|
|
)
|
|
}
|
|
)
|
|
items_data = [
|
|
{"user": {"name": "Alice", "age": 30}},
|
|
{"user": {"name": "Bob", "age": 25}},
|
|
]
|
|
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
|
|
|
result = ds.with_column("name", col("user").struct["name"]).to_pandas()
|
|
expected = pd.DataFrame(
|
|
{
|
|
"user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
|
|
"name": ["Alice", "Bob"],
|
|
}
|
|
)
|
|
assert rows_same(result, expected)
|
|
|
|
def test_struct_field_by_index(self, ray_start_regular_shared, dataset_format):
|
|
"""Test struct.field_by_index() extracts field by position."""
|
|
if dataset_format == "pandas":
|
|
pytest.skip(
|
|
"Index-based struct access requires stable Arrow struct field ordering."
|
|
)
|
|
arrow_table = pa.table(
|
|
{
|
|
"user": pa.array(
|
|
[
|
|
{"name": "Alice", "age": 30},
|
|
{"name": "Bob", "age": 25},
|
|
],
|
|
type=pa.struct(
|
|
[
|
|
pa.field("name", pa.string()),
|
|
pa.field("age", pa.int32()),
|
|
]
|
|
),
|
|
)
|
|
}
|
|
)
|
|
items_data = [
|
|
{"user": {"name": "Alice", "age": 30}},
|
|
{"user": {"name": "Bob", "age": 25}},
|
|
]
|
|
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
|
|
|
result = ds.with_column("age", col("user").struct.field_by_index(1)).to_pandas()
|
|
expected = pd.DataFrame(
|
|
{
|
|
"user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
|
|
"age": [30, 25],
|
|
}
|
|
)
|
|
assert rows_same(result, expected)
|
|
|
|
def test_struct_bracket_with_index(self, ray_start_regular_shared, dataset_format):
|
|
"""Test struct[index] bracket notation."""
|
|
if dataset_format == "pandas":
|
|
pytest.skip(
|
|
"Index-based struct access requires stable Arrow struct field ordering."
|
|
)
|
|
arrow_table = pa.table(
|
|
{
|
|
"user": pa.array(
|
|
[
|
|
{"name": "Alice", "age": 30},
|
|
{"name": "Bob", "age": 25},
|
|
],
|
|
type=pa.struct(
|
|
[
|
|
pa.field("name", pa.string()),
|
|
pa.field("age", pa.int32()),
|
|
]
|
|
),
|
|
)
|
|
}
|
|
)
|
|
items_data = [
|
|
{"user": {"name": "Alice", "age": 30}},
|
|
{"user": {"name": "Bob", "age": 25}},
|
|
]
|
|
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
|
|
|
result = ds.with_column("name", col("user").struct[0]).to_pandas()
|
|
expected = pd.DataFrame(
|
|
{
|
|
"user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
|
|
"name": ["Alice", "Bob"],
|
|
}
|
|
)
|
|
assert rows_same(result, expected)
|
|
|
|
def test_struct_nested_field(self, ray_start_regular_shared, dataset_format):
|
|
"""Test nested struct field access with .field()."""
|
|
arrow_table = pa.table(
|
|
{
|
|
"user": pa.array(
|
|
[
|
|
{"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
|
|
{"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
|
|
],
|
|
type=pa.struct(
|
|
[
|
|
pa.field("name", pa.string()),
|
|
pa.field(
|
|
"address",
|
|
pa.struct(
|
|
[
|
|
pa.field("city", pa.string()),
|
|
pa.field("zip", pa.string()),
|
|
]
|
|
),
|
|
),
|
|
]
|
|
),
|
|
)
|
|
}
|
|
)
|
|
items_data = [
|
|
{"user": {"name": "Alice", "address": {"city": "NYC", "zip": "10001"}}},
|
|
{"user": {"name": "Bob", "address": {"city": "LA", "zip": "90001"}}},
|
|
]
|
|
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
|
|
|
result = ds.with_column(
|
|
"city", col("user").struct.field("address").struct.field("city")
|
|
).to_pandas()
|
|
expected = pd.DataFrame(
|
|
{
|
|
"user": [
|
|
{"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
|
|
{"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
|
|
],
|
|
"city": ["NYC", "LA"],
|
|
}
|
|
)
|
|
assert rows_same(result, expected)
|
|
|
|
def test_struct_nested_bracket(self, ray_start_regular_shared, dataset_format):
|
|
"""Test nested struct field access with brackets."""
|
|
arrow_table = pa.table(
|
|
{
|
|
"user": pa.array(
|
|
[
|
|
{"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
|
|
{"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
|
|
],
|
|
type=pa.struct(
|
|
[
|
|
pa.field("name", pa.string()),
|
|
pa.field(
|
|
"address",
|
|
pa.struct(
|
|
[
|
|
pa.field("city", pa.string()),
|
|
pa.field("zip", pa.string()),
|
|
]
|
|
),
|
|
),
|
|
]
|
|
),
|
|
)
|
|
}
|
|
)
|
|
items_data = [
|
|
{"user": {"name": "Alice", "address": {"city": "NYC", "zip": "10001"}}},
|
|
{"user": {"name": "Bob", "address": {"city": "LA", "zip": "90001"}}},
|
|
]
|
|
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
|
|
|
result = ds.with_column(
|
|
"zip", col("user").struct["address"].struct["zip"]
|
|
).to_pandas()
|
|
expected = pd.DataFrame(
|
|
{
|
|
"user": [
|
|
{"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
|
|
{"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
|
|
],
|
|
"zip": ["10001", "90001"],
|
|
}
|
|
)
|
|
assert rows_same(result, expected)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|