wehub-resource-sync
151 lines
5.1 KiB
Python
151 lines
5.1 KiB
Python
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
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
version.parse(pa.__version__) < version.parse("19.0.0"),
|
|
reason="Namespace expressions tests require PyArrow >= 19.0",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def map_dataset():
|
|
"""Fixture that creates a dataset backed by an Arrow MapArray column."""
|
|
map_items = [
|
|
{"attrs": {"color": "red", "size": "M"}},
|
|
{"attrs": {"brand": "Ray"}},
|
|
]
|
|
map_type = pa.map_(pa.string(), pa.string())
|
|
arrow_table = pa.table(
|
|
{"attrs": pa.array([row["attrs"] for row in map_items], type=map_type)}
|
|
)
|
|
return ray.data.from_arrow(arrow_table)
|
|
|
|
|
|
def _assert_result(result_df: pd.DataFrame, expected_df: pd.DataFrame, drop_cols: list):
|
|
"""Helper to drop columns and assert equality."""
|
|
result_df = result_df.drop(columns=drop_cols)
|
|
assert rows_same(result_df, expected_df)
|
|
|
|
|
|
class TestMapNamespace:
|
|
"""Tests for map namespace operations using the shared map_dataset fixture."""
|
|
|
|
def test_map_keys(self, map_dataset):
|
|
result = map_dataset.with_column("keys", col("attrs").map.keys()).to_pandas()
|
|
expected = pd.DataFrame({"keys": [["color", "size"], ["brand"]]})
|
|
_assert_result(result, expected, drop_cols=["attrs"])
|
|
|
|
def test_map_values(self, map_dataset):
|
|
result = map_dataset.with_column(
|
|
"values", col("attrs").map.values()
|
|
).to_pandas()
|
|
expected = pd.DataFrame({"values": [["red", "M"], ["Ray"]]})
|
|
_assert_result(result, expected, drop_cols=["attrs"])
|
|
|
|
def test_map_chaining(self, map_dataset):
|
|
# map.keys() returns a list, so .list.len() should apply
|
|
result = map_dataset.with_column(
|
|
"num_keys", col("attrs").map.keys().list.len()
|
|
).to_pandas()
|
|
expected = pd.DataFrame({"num_keys": [2, 1]})
|
|
_assert_result(result, expected, drop_cols=["attrs"])
|
|
|
|
|
|
def test_physical_map_extraction():
|
|
"""Test extraction works on List<Struct> (Physical Maps)."""
|
|
# Construct List<Struct<k, v>>
|
|
struct_type = pa.struct([pa.field("k", pa.string()), pa.field("v", pa.int64())])
|
|
list_type = pa.list_(struct_type)
|
|
|
|
data_py = [[{"k": "a", "v": 1}], [{"k": "b", "v": 2}]]
|
|
arrow_table = pa.Table.from_arrays(
|
|
[pa.array(data_py, type=list_type)], names=["data"]
|
|
)
|
|
ds = ray.data.from_arrow(arrow_table)
|
|
|
|
result = (
|
|
ds.with_column("keys", col("data").map.keys())
|
|
.with_column("values", col("data").map.values())
|
|
.to_pandas()
|
|
)
|
|
|
|
expected = pd.DataFrame(
|
|
{
|
|
"data": data_py,
|
|
"keys": [["a"], ["b"]],
|
|
"values": [[1], [2]],
|
|
}
|
|
)
|
|
assert rows_same(result, expected)
|
|
|
|
|
|
def test_map_sliced_offsets():
|
|
"""Test extraction works correctly on sliced Arrow arrays (offset > 0)."""
|
|
items = [{"m": {"id": i}} for i in range(10)]
|
|
map_type = pa.map_(pa.string(), pa.int64())
|
|
arrays = pa.array([row["m"] for row in items], type=map_type)
|
|
table = pa.Table.from_arrays([arrays], names=["m"])
|
|
|
|
# Force offsets by slicing the table before ingestion
|
|
sliced_table = table.slice(offset=7, length=3)
|
|
ds = ray.data.from_arrow(sliced_table)
|
|
|
|
result = ds.with_column("vals", col("m").map.values()).to_pandas()
|
|
expected = pd.DataFrame({"vals": [[7], [8], [9]]})
|
|
_assert_result(result, expected, drop_cols=["m"])
|
|
|
|
|
|
def test_map_nulls_and_empty():
|
|
"""Test handling of null maps and empty maps."""
|
|
items_data = [{"m": {"a": 1}}, {"m": {}}, {"m": None}]
|
|
|
|
map_type = pa.map_(pa.string(), pa.int64())
|
|
arrays = pa.array([row["m"] for row in items_data], type=map_type)
|
|
arrow_table = pa.Table.from_arrays([arrays], names=["m"])
|
|
ds = ray.data.from_arrow(arrow_table)
|
|
|
|
rows = (
|
|
ds.with_column("keys", col("m").map.keys())
|
|
.with_column("values", col("m").map.values())
|
|
.take_all()
|
|
)
|
|
|
|
assert list(rows[0]["keys"]) == ["a"] and list(rows[0]["values"]) == [1]
|
|
assert len(rows[1]["keys"]) == 0 and len(rows[1]["values"]) == 0
|
|
assert rows[2]["keys"] is None and rows[2]["values"] is None
|
|
|
|
|
|
def test_empty_chunked_array():
|
|
"""Test extraction works on empty ChunkedArray (zero chunks)."""
|
|
from ray.data.namespace_expressions.map_namespace import (
|
|
MapComponent,
|
|
_extract_map_component,
|
|
)
|
|
|
|
# Create empty ChunkedArray with map type
|
|
map_type = pa.map_(pa.string(), pa.int64())
|
|
empty_chunked = pa.chunked_array([], type=map_type)
|
|
assert empty_chunked.num_chunks == 0
|
|
|
|
# Extract keys - should return empty ChunkedArray with list<string> type
|
|
keys_result = _extract_map_component(empty_chunked, MapComponent.KEYS)
|
|
assert len(keys_result) == 0
|
|
assert keys_result.type == pa.list_(pa.string())
|
|
|
|
# Extract values - should return empty ChunkedArray with list<int64> type
|
|
values_result = _extract_map_component(empty_chunked, MapComponent.VALUES)
|
|
assert len(values_result) == 0
|
|
assert values_result.type == pa.list_(pa.int64())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|