wehub-resource-sync
147 lines
4.0 KiB
Python
147 lines
4.0 KiB
Python
from __future__ import annotations
|
|
|
|
import pickle
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.compute as pc
|
|
import pytest
|
|
from rerun.experimental import Selector
|
|
|
|
|
|
def test_parse_and_execute_field() -> None:
|
|
arr = pa.StructArray.from_arrays(
|
|
[pa.array([1.0, 2.0, 3.0])],
|
|
names=["x"],
|
|
)
|
|
result = Selector(".x").execute(arr)
|
|
assert result is not None
|
|
assert result == pa.array([1.0, 2.0, 3.0])
|
|
|
|
|
|
def test_nested_field() -> None:
|
|
inner = pa.StructArray.from_arrays(
|
|
[pa.array([10, 20])],
|
|
names=["value"],
|
|
)
|
|
outer = pa.StructArray.from_arrays(
|
|
[inner],
|
|
names=["nested"],
|
|
)
|
|
result = Selector(".nested.value").execute(outer)
|
|
assert result is not None
|
|
assert result == pa.array([10, 20])
|
|
|
|
|
|
def test_invalid_query() -> None:
|
|
with pytest.raises(ValueError, match="Failed to parse selector"):
|
|
Selector("invalid query without dot")
|
|
|
|
|
|
def test_repr_and_str() -> None:
|
|
assert "Selector" in repr(Selector(".field"))
|
|
assert ".field" in str(Selector(".field"))
|
|
|
|
|
|
def test_pipe_with_callable() -> None:
|
|
arr = pa.StructArray.from_arrays(
|
|
[pa.array([1.0, 2.0, 3.0])],
|
|
names=["x"],
|
|
)
|
|
result = Selector(".x").pipe(lambda a: pc.multiply(a, 2)).execute(arr)
|
|
assert result is not None
|
|
assert result == pa.array([2.0, 4.0, 6.0])
|
|
|
|
|
|
def test_pipe_with_plain_callable() -> None:
|
|
arr = pa.StructArray.from_arrays(
|
|
[pa.array([1.0, 2.0, 3.0])],
|
|
names=["x"],
|
|
)
|
|
result = Selector(".x").pipe(lambda a: pa.array([v.as_py() ** 2 for v in a])).execute(arr)
|
|
assert result is not None
|
|
assert result == pa.array([1.0, 4.0, 9.0])
|
|
|
|
|
|
def test_pipe_with_selector() -> None:
|
|
inner = pa.StructArray.from_arrays(
|
|
[pa.array([10, 20])],
|
|
names=["value"],
|
|
)
|
|
outer = pa.StructArray.from_arrays(
|
|
[inner],
|
|
names=["nested"],
|
|
)
|
|
result = Selector(".nested").pipe(Selector(".value")).execute(outer)
|
|
assert result is not None
|
|
assert result == pa.array([10, 20])
|
|
|
|
|
|
def test_pipe_callable_returning_none() -> None:
|
|
arr = pa.StructArray.from_arrays(
|
|
[pa.array([1, 2, 3])],
|
|
names=["x"],
|
|
)
|
|
result = Selector(".x").pipe(lambda _a: None).execute(arr)
|
|
assert result is None
|
|
|
|
|
|
def test_execute_per_row() -> None:
|
|
inner_type = pa.struct([pa.field("value", pa.int64())])
|
|
arr = pa.array(
|
|
[
|
|
[{"value": 1}, {"value": 2}],
|
|
[{"value": 3}],
|
|
],
|
|
type=pa.list_(inner_type),
|
|
)
|
|
result = Selector(".value").execute_per_row(arr)
|
|
assert result is not None
|
|
assert result == pa.array([[1, 2], [3]], type=pa.list_(pa.int64()))
|
|
|
|
|
|
def test_pipe_chaining() -> None:
|
|
inner = pa.StructArray.from_arrays(
|
|
[pa.array([5.0, 10.0])],
|
|
names=["val"],
|
|
)
|
|
outer = pa.StructArray.from_arrays(
|
|
[inner],
|
|
names=["data"],
|
|
)
|
|
result = Selector(".data").pipe(Selector(".val")).pipe(lambda a: pc.add(a, 1)).execute(outer)
|
|
assert result is not None
|
|
assert result == pa.array([6.0, 11.0])
|
|
|
|
|
|
def test_pickle_roundtrip_simple() -> None:
|
|
arr = pa.StructArray.from_arrays(
|
|
[pa.array([1.0, 2.0, 3.0])],
|
|
names=["x"],
|
|
)
|
|
original = Selector(".x")
|
|
restored = pickle.loads(pickle.dumps(original))
|
|
assert isinstance(restored, Selector)
|
|
assert str(restored) == str(original)
|
|
assert restored.execute(arr) == pa.array([1.0, 2.0, 3.0])
|
|
|
|
|
|
def test_pickle_roundtrip_piped_selectors() -> None:
|
|
inner = pa.StructArray.from_arrays(
|
|
[pa.array([10, 20])],
|
|
names=["value"],
|
|
)
|
|
outer = pa.StructArray.from_arrays(
|
|
[inner],
|
|
names=["nested"],
|
|
)
|
|
original = Selector(".nested").pipe(Selector(".value"))
|
|
restored = pickle.loads(pickle.dumps(original))
|
|
assert isinstance(restored, Selector)
|
|
assert restored.execute(outer) == pa.array([10, 20])
|
|
|
|
|
|
def test_pickle_rejects_callable_pipe() -> None:
|
|
sel = Selector(".x").pipe(lambda a: pc.multiply(a, 2))
|
|
with pytest.raises(TypeError, match="Cannot pickle Selector"):
|
|
pickle.dumps(sel)
|