wehub-resource-sync
3257 lines
106 KiB
Python
3257 lines
106 KiB
Python
import re
|
|
import types
|
|
from typing import Iterable
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
from packaging.version import parse as parse_version
|
|
|
|
from ray.data._internal.arrow_ops.transform_pyarrow import (
|
|
MIN_PYARROW_VERSION_TYPE_PROMOTION,
|
|
_align_struct_fields,
|
|
_has_unhashable_pandas_types,
|
|
concat,
|
|
hash_partition,
|
|
shuffle,
|
|
try_combine_chunked_columns,
|
|
unify_schemas,
|
|
)
|
|
from ray.data._internal.tensor_extensions.arrow import (
|
|
ArrowTensorTypeV2,
|
|
_extension_array_concat_supported,
|
|
create_arrow_fixed_shape_tensor_type,
|
|
)
|
|
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
|
from ray.data.block import BlockAccessor
|
|
from ray.data.extensions import (
|
|
ArrowConversionError,
|
|
ArrowPythonObjectArray,
|
|
ArrowPythonObjectType,
|
|
ArrowTensorArray,
|
|
ArrowTensorType,
|
|
ArrowVariableShapedTensorArray,
|
|
ArrowVariableShapedTensorType,
|
|
)
|
|
|
|
|
|
def test_try_defragment_table():
|
|
chunks = np.array_split(np.arange(1000), 10)
|
|
|
|
t = pa.Table.from_pydict(
|
|
{
|
|
"id": pa.chunked_array([pa.array(c) for c in chunks]),
|
|
}
|
|
)
|
|
|
|
assert len(t["id"].chunks) == 10
|
|
|
|
dt = try_combine_chunked_columns(t)
|
|
|
|
assert len(dt["id"].chunks) == 1
|
|
assert dt == t
|
|
|
|
|
|
def test_try_combine_chunked_columns_min_chunks_to_combine():
|
|
"""Test that the min_chunks_to_combine parameter controls the combining
|
|
threshold."""
|
|
# Create a table with 3 chunks per column.
|
|
t = pa.Table.from_pydict(
|
|
{
|
|
"a": pa.chunked_array(
|
|
[pa.array([1, 2]), pa.array([3, 4]), pa.array([5, 6])]
|
|
),
|
|
"b": pa.chunked_array(
|
|
[pa.array([7, 8]), pa.array([9, 10]), pa.array([11, 12])]
|
|
),
|
|
}
|
|
)
|
|
assert t["a"].num_chunks == 3
|
|
assert t["b"].num_chunks == 3
|
|
|
|
# Default threshold (10) should NOT combine since 3 < 10.
|
|
result = try_combine_chunked_columns(t)
|
|
assert result["a"].num_chunks == 3
|
|
assert result["b"].num_chunks == 3
|
|
|
|
# min_chunks_to_combine=1 should always combine.
|
|
result = try_combine_chunked_columns(t, min_chunks_to_combine=1)
|
|
assert result["a"].num_chunks == 1
|
|
assert result["b"].num_chunks == 1
|
|
assert result == t
|
|
|
|
# min_chunks_to_combine=3 should combine (3 >= 3).
|
|
result = try_combine_chunked_columns(t, min_chunks_to_combine=3)
|
|
assert result["a"].num_chunks == 1
|
|
assert result["b"].num_chunks == 1
|
|
|
|
# min_chunks_to_combine=4 should NOT combine (3 < 4).
|
|
result = try_combine_chunked_columns(t, min_chunks_to_combine=4)
|
|
assert result["a"].num_chunks == 3
|
|
assert result["b"].num_chunks == 3
|
|
|
|
|
|
def test_hash_partitioning():
|
|
# Test hash-partitioning of the empty table
|
|
empty_table = pa.Table.from_pydict({"idx": []})
|
|
|
|
assert {} == hash_partition(empty_table, hash_cols=["idx"], num_partitions=5)
|
|
|
|
# Test hash-partitioning of table into 1 partition (returns table itself)
|
|
t = pa.Table.from_pydict({"idx": list(range(10))})
|
|
|
|
assert {0: t} == hash_partition(t, hash_cols=["idx"], num_partitions=1)
|
|
|
|
# Test hash-partitioning of proper table
|
|
idx = list(range(100))
|
|
|
|
t = pa.Table.from_pydict(
|
|
{
|
|
"idx": pa.array(idx),
|
|
"ints": pa.array(idx),
|
|
"floats": pa.array([float(i) for i in idx]),
|
|
"strings": pa.array([str(i) for i in idx]),
|
|
"structs": pa.array(
|
|
[
|
|
{
|
|
"value": i,
|
|
}
|
|
for i in idx
|
|
]
|
|
),
|
|
}
|
|
)
|
|
|
|
single_partition_dict = hash_partition(t, hash_cols=["idx"], num_partitions=1)
|
|
|
|
# There's just 1 partition
|
|
assert len(single_partition_dict) == 1
|
|
assert t == single_partition_dict.get(0)
|
|
|
|
def _concat_and_sort_partitions(parts: Iterable[pa.Table]) -> pa.Table:
|
|
return pa.concat_tables(parts).sort_by("idx")
|
|
|
|
_5_partition_dict = hash_partition(t, hash_cols=["strings"], num_partitions=5)
|
|
|
|
assert len(_5_partition_dict) == 5
|
|
assert t == _concat_and_sort_partitions(_5_partition_dict.values())
|
|
|
|
# There could be no more partitions than elements
|
|
_structs_partition_dict = hash_partition(
|
|
t, hash_cols=["structs"], num_partitions=101
|
|
)
|
|
|
|
assert len(_structs_partition_dict) <= 101
|
|
assert t == _concat_and_sort_partitions(_structs_partition_dict.values())
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"pa_type,expected",
|
|
[
|
|
# Nested types -> unhashable in pandas (convert to dict/list)
|
|
(pa.struct([("a", pa.int32())]), True),
|
|
(pa.list_(pa.int32()), True),
|
|
(pa.large_list(pa.int32()), True),
|
|
(pa.list_(pa.int32(), 3), True), # fixed_size_list
|
|
(pa.map_(pa.string(), pa.int32()), True),
|
|
(pa.dense_union([pa.field("x", pa.int32())]), True),
|
|
# Ray extension types -> numpy arrays / arbitrary objects in pandas
|
|
(ArrowTensorTypeV2((2, 2), pa.int64()), True),
|
|
(ArrowPythonObjectType(), True),
|
|
# Hashable primitives -> must stay False so we keep the fast path
|
|
(pa.int32(), False),
|
|
(pa.float64(), False),
|
|
(pa.bool_(), False),
|
|
(pa.string(), False),
|
|
(pa.large_string(), False),
|
|
(pa.binary(), False),
|
|
(pa.decimal128(10, 2), False),
|
|
(pa.date32(), False),
|
|
(pa.timestamp("ns"), False),
|
|
(pa.dictionary(pa.int32(), pa.string()), False),
|
|
],
|
|
)
|
|
def test_has_unhashable_pandas_types(pa_type, expected):
|
|
schema = pa.schema([("c", pa_type)])
|
|
assert _has_unhashable_pandas_types(schema) is expected
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < parse_version("16.0.0"),
|
|
reason="list_view / large_list_view require pyarrow 16+",
|
|
)
|
|
def test_has_unhashable_pandas_types_list_views():
|
|
# Regression: list_view/large_list_view also convert to Python lists in
|
|
# pandas, so they must be flagged as unhashable like list/large_list.
|
|
for view_type in (pa.list_view(pa.int32()), pa.large_list_view(pa.int32())):
|
|
schema = pa.schema([("c", view_type)])
|
|
assert _has_unhashable_pandas_types(schema) is True
|
|
|
|
|
|
def test_hash_partition_null_struct_consistent_across_blocks():
|
|
struct_t = pa.struct([("v", pa.int32())])
|
|
num_partitions = 8
|
|
|
|
all_null = pa.Table.from_pydict(
|
|
{"k": pa.array([None, None, None], type=struct_t), "idx": [0, 1, 2]}
|
|
)
|
|
mixed = pa.Table.from_pydict(
|
|
{
|
|
"k": pa.array([None, {"v": 1}, None], type=struct_t),
|
|
"idx": [10, 11, 12],
|
|
}
|
|
)
|
|
|
|
p1 = hash_partition(all_null, hash_cols=["k"], num_partitions=num_partitions)
|
|
p2 = hash_partition(mixed, hash_cols=["k"], num_partitions=num_partitions)
|
|
|
|
def null_partition_id(parts):
|
|
# Return the partition id holding null-key rows (there should be
|
|
# exactly one — identical null keys must co-locate).
|
|
null_pids = {
|
|
pid for pid, tbl in parts.items() if any(tbl["k"].is_null().to_pylist())
|
|
}
|
|
assert len(null_pids) == 1, null_pids
|
|
return next(iter(null_pids))
|
|
|
|
assert null_partition_id(p1) == null_partition_id(p2)
|
|
|
|
|
|
def test_shuffle():
|
|
t = pa.Table.from_pydict(
|
|
{
|
|
"index": pa.array(list(range(10))),
|
|
}
|
|
)
|
|
|
|
shuffled = shuffle(t, seed=0xDEED)
|
|
|
|
assert shuffled == pa.Table.from_pydict(
|
|
{"index": pa.array([4, 3, 6, 8, 7, 1, 5, 2, 9, 0])}
|
|
)
|
|
|
|
|
|
def test_arrow_concat_empty(simple_concat_data):
|
|
# Test empty.
|
|
assert concat(simple_concat_data["empty"]) == pa.table([])
|
|
|
|
|
|
def test_arrow_concat_single_block(simple_concat_data):
|
|
# Test single block:
|
|
out = concat([simple_concat_data["single_block"]])
|
|
assert len(out) == 2
|
|
assert out == simple_concat_data["single_block"]
|
|
|
|
|
|
def test_arrow_concat_basic(basic_concat_blocks, basic_concat_expected):
|
|
# Test two basic tables.
|
|
ts = basic_concat_blocks
|
|
out = concat(ts)
|
|
# Check length.
|
|
assert len(out) == basic_concat_expected["length"]
|
|
# Check schema.
|
|
assert out.column_names == basic_concat_expected["column_names"]
|
|
assert out.schema.types == basic_concat_expected["schema_types"]
|
|
# Confirm that concatenation is zero-copy (i.e. it didn't trigger chunk
|
|
# consolidation).
|
|
assert out["a"].num_chunks == basic_concat_expected["chunks"]
|
|
assert out["b"].num_chunks == basic_concat_expected["chunks"]
|
|
# Check content.
|
|
assert out["a"].to_pylist() == basic_concat_expected["content"]["a"]
|
|
assert out["b"].to_pylist() == basic_concat_expected["content"]["b"]
|
|
# Check equivalence.
|
|
expected = pa.concat_tables(ts)
|
|
assert out == expected
|
|
|
|
|
|
def test_arrow_concat_null_promotion(null_promotion_blocks, null_promotion_expected):
|
|
# Test null column --> well-typed column promotion.
|
|
ts = null_promotion_blocks
|
|
out = concat(ts)
|
|
# Check length.
|
|
assert len(out) == null_promotion_expected["length"]
|
|
# Check schema.
|
|
assert out.column_names == null_promotion_expected["column_names"]
|
|
assert out.schema.types == null_promotion_expected["schema_types"]
|
|
# Confirm that concatenation is zero-copy (i.e. it didn't trigger chunk
|
|
# consolidation).
|
|
assert out["a"].num_chunks == null_promotion_expected["chunks"]
|
|
assert out["b"].num_chunks == null_promotion_expected["chunks"]
|
|
# Check content.
|
|
assert out["a"].to_pylist() == null_promotion_expected["content"]["a"]
|
|
assert out["b"].to_pylist() == null_promotion_expected["content"]["b"]
|
|
# Check equivalence.
|
|
expected = pa.concat_tables(ts, promote=True)
|
|
assert out == expected
|
|
|
|
|
|
def test_arrow_concat_tensor_extension_uniform(
|
|
uniform_tensor_blocks, uniform_tensor_expected
|
|
):
|
|
# Test tensor column concatenation.
|
|
t1, t2 = uniform_tensor_blocks
|
|
ts = [t1, t2]
|
|
out = concat(ts)
|
|
|
|
# Check length.
|
|
assert len(out) == uniform_tensor_expected["length"]
|
|
|
|
# Check schema.
|
|
assert out.column_names == ["a"]
|
|
assert out.schema == uniform_tensor_expected["schema"]
|
|
|
|
# Confirm that concatenation is zero-copy (i.e. it didn't trigger chunk
|
|
# consolidation).
|
|
assert out["a"].num_chunks == uniform_tensor_expected["chunks"]
|
|
|
|
# Check content.
|
|
content = uniform_tensor_expected["content"]
|
|
np.testing.assert_array_equal(out["a"].chunk(0).to_numpy_ndarray(), content[0])
|
|
np.testing.assert_array_equal(out["a"].chunk(1).to_numpy_ndarray(), content[1])
|
|
|
|
# Check equivalence.
|
|
expected = pa.concat_tables(ts, promote=True)
|
|
assert out == expected
|
|
|
|
|
|
def test_arrow_concat_tensor_extension_variable_shaped(
|
|
variable_shaped_tensor_blocks, variable_shaped_tensor_expected
|
|
):
|
|
# Test variable_shaped tensor column concatenation.
|
|
t1, t2 = variable_shaped_tensor_blocks
|
|
ts = [t1, t2]
|
|
out = concat(ts)
|
|
# Check length.
|
|
assert len(out) == variable_shaped_tensor_expected["length"]
|
|
# Check schema.
|
|
assert out.column_names == ["a"]
|
|
assert out.schema == variable_shaped_tensor_expected["schema"]
|
|
# Confirm that concatenation is zero-copy (i.e. it didn't trigger chunk
|
|
# consolidation).
|
|
assert out["a"].num_chunks == variable_shaped_tensor_expected["chunks"]
|
|
# Check content.
|
|
content = variable_shaped_tensor_expected["content"]
|
|
for o, e in zip(out["a"].chunk(0).to_numpy(), content[0]):
|
|
np.testing.assert_array_equal(o, e)
|
|
for o, e in zip(out["a"].chunk(1).to_numpy(), content[1]):
|
|
np.testing.assert_array_equal(o, e)
|
|
# NOTE: We don't check equivalence with pyarrow.concat_tables since it currently
|
|
# fails for this case.
|
|
|
|
|
|
@pytest.mark.parametrize("preserve_order", [True, False])
|
|
def test_arrow_concat_tensor_extension_uniform_and_variable_shaped(
|
|
mixed_tensor_blocks, mixed_tensor_expected, preserve_order
|
|
):
|
|
# Test concatenating a homogeneous-shaped tensor column with a variable-shaped
|
|
# tensor column.
|
|
t1, t2 = mixed_tensor_blocks
|
|
ts = [t1, t2]
|
|
out = concat(ts, preserve_order=preserve_order)
|
|
# Check length.
|
|
assert len(out) == mixed_tensor_expected["length"]
|
|
# Check schema.
|
|
assert out.column_names == ["a"]
|
|
assert out.schema == mixed_tensor_expected["schema"]
|
|
|
|
# Confirm that concatenation is zero-copy (i.e. it didn't trigger chunk
|
|
# consolidation).
|
|
assert out["a"].num_chunks == mixed_tensor_expected["chunks"]
|
|
|
|
# Collect all arrays from output and expected.
|
|
actual = [
|
|
arr
|
|
for chunk_idx in range(out["a"].num_chunks)
|
|
for arr in out["a"].chunk(chunk_idx).to_numpy()
|
|
]
|
|
expected = [arr for chunk in mixed_tensor_expected["content"] for arr in chunk]
|
|
assert len(actual) == len(expected)
|
|
|
|
if not preserve_order:
|
|
actual = sorted(actual, key=lambda arr: arr.tobytes())
|
|
expected = sorted(expected, key=lambda arr: arr.tobytes())
|
|
|
|
for a, e in zip(actual, expected):
|
|
np.testing.assert_array_equal(a, e)
|
|
# NOTE: We don't check equivalence with pyarrow.concat_tables since it currently
|
|
# fails for this case.
|
|
|
|
|
|
def test_arrow_concat_tensor_extension_uniform_but_different(
|
|
different_shape_tensor_blocks, different_shape_tensor_expected
|
|
):
|
|
# Test concatenating two homogeneous-shaped tensor columns with differing shapes
|
|
# between them.
|
|
t1, t2 = different_shape_tensor_blocks
|
|
ts = [t1, t2]
|
|
out = concat(ts)
|
|
# Check length.
|
|
assert len(out) == different_shape_tensor_expected["length"]
|
|
# Check schema.
|
|
assert out.column_names == ["a"]
|
|
assert out.schema == different_shape_tensor_expected["schema"]
|
|
# Confirm that concatenation is zero-copy (i.e. it didn't trigger chunk
|
|
# consolidation).
|
|
assert out["a"].num_chunks == different_shape_tensor_expected["chunks"]
|
|
# Check content.
|
|
content = different_shape_tensor_expected["content"]
|
|
for o, e in zip(out["a"].chunk(0).to_numpy(), content[0]):
|
|
np.testing.assert_array_equal(o, e)
|
|
for o, e in zip(out["a"].chunk(1).to_numpy(), content[1]):
|
|
np.testing.assert_array_equal(o, e)
|
|
# NOTE: We don't check equivalence with pyarrow.concat_tables since it currently
|
|
# fails for this case.
|
|
|
|
|
|
@pytest.mark.parametrize("preserve_order", [True, False])
|
|
def test_arrow_concat_with_objects(
|
|
object_concat_blocks, object_concat_expected, preserve_order
|
|
):
|
|
t3 = concat(object_concat_blocks, preserve_order=preserve_order)
|
|
assert isinstance(t3, pa.Table)
|
|
assert len(t3) == object_concat_expected["length"]
|
|
assert isinstance(t3.schema.field("a").type, object_concat_expected["a_type"])
|
|
assert object_concat_expected["b_type"](t3.schema.field("b").type)
|
|
|
|
actual_a = t3.column("a").to_pylist()
|
|
actual_b = t3.column("b").to_pylist()
|
|
expected_a = object_concat_expected["content"]["a"]
|
|
expected_b = object_concat_expected["content"]["b"]
|
|
|
|
if preserve_order:
|
|
assert actual_a == expected_a
|
|
assert actual_b == expected_b
|
|
else:
|
|
assert sorted(actual_a, key=str) == sorted(expected_a, key=str)
|
|
assert sorted(actual_b, key=str) == sorted(expected_b, key=str)
|
|
|
|
|
|
def test_struct_with_different_field_names(
|
|
struct_different_field_names_blocks, struct_different_field_names_expected
|
|
):
|
|
# Ensures that when concatenating tables with struct columns having different
|
|
# field names, missing fields in each struct are filled with None in the
|
|
# resulting table.
|
|
|
|
# Concatenate tables with different field names in struct
|
|
t3 = concat(struct_different_field_names_blocks)
|
|
|
|
assert isinstance(t3, pa.Table)
|
|
assert len(t3) == struct_different_field_names_expected["length"]
|
|
|
|
# Check the entire schema
|
|
assert t3.schema == struct_different_field_names_expected["schema"]
|
|
|
|
# Check that missing fields are filled with None
|
|
assert (
|
|
t3.column("a").to_pylist()
|
|
== struct_different_field_names_expected["content"]["a"]
|
|
)
|
|
assert (
|
|
t3.column("d").to_pylist()
|
|
== struct_different_field_names_expected["content"]["d"]
|
|
)
|
|
|
|
|
|
def test_nested_structs(nested_structs_blocks, nested_structs_expected):
|
|
# Checks that deeply nested structs (3 levels of nesting) are handled properly
|
|
# during concatenation and the resulting table preserves the correct nesting
|
|
# structure.
|
|
|
|
# Concatenate tables with nested structs and missing fields
|
|
t3 = concat(nested_structs_blocks)
|
|
assert isinstance(t3, pa.Table)
|
|
assert len(t3) == nested_structs_expected["length"]
|
|
|
|
# Validate the schema of the resulting table
|
|
assert t3.schema == nested_structs_expected["schema"]
|
|
|
|
# Validate the data in the concatenated table
|
|
assert t3.column("a").to_pylist() == nested_structs_expected["content"]["a"]
|
|
assert t3.column("d").to_pylist() == nested_structs_expected["content"]["d"]
|
|
|
|
|
|
def test_struct_with_null_values(
|
|
struct_null_values_blocks, struct_null_values_expected
|
|
):
|
|
# Ensures that when concatenating tables with struct columns containing null
|
|
# values, the null values are properly handled, and the result reflects the
|
|
# expected structure.
|
|
|
|
# Concatenate tables with struct columns containing null values
|
|
t3 = concat(struct_null_values_blocks)
|
|
assert isinstance(t3, pa.Table)
|
|
assert len(t3) == struct_null_values_expected["length"]
|
|
|
|
# Validate the schema of the resulting table
|
|
assert (
|
|
t3.schema == struct_null_values_expected["schema"]
|
|
), f"Expected schema: {struct_null_values_expected['schema']}, but got {t3.schema}"
|
|
|
|
# Verify the PyArrow table content
|
|
assert t3.column("a").to_pylist() == struct_null_values_expected["content"]["a"]
|
|
|
|
result = t3.column("d").to_pylist()
|
|
expected = struct_null_values_expected["content"]["d"]
|
|
assert result == expected, f"Expected {expected}, but got {result}"
|
|
|
|
|
|
def test_struct_with_mismatched_lengths(
|
|
struct_mismatched_lengths_blocks, struct_mismatched_lengths_expected
|
|
):
|
|
# Verifies that when concatenating tables with struct columns of different lengths,
|
|
# the missing values are properly padded with None in the resulting table.
|
|
|
|
# Concatenate tables with struct columns of different lengths
|
|
t3 = concat(struct_mismatched_lengths_blocks)
|
|
assert isinstance(t3, pa.Table)
|
|
assert (
|
|
len(t3) == struct_mismatched_lengths_expected["length"]
|
|
) # Check that the resulting table has the correct number of rows
|
|
|
|
# Validate the schema of the resulting table
|
|
assert (
|
|
t3.schema == struct_mismatched_lengths_expected["schema"]
|
|
), f"Expected schema: {struct_mismatched_lengths_expected['schema']}, but got {t3.schema}"
|
|
|
|
# Verify the content of the resulting table
|
|
assert (
|
|
t3.column("a").to_pylist() == struct_mismatched_lengths_expected["content"]["a"]
|
|
)
|
|
result = t3.column("d").to_pylist()
|
|
expected = struct_mismatched_lengths_expected["content"]["d"]
|
|
|
|
assert result == expected, f"Expected {expected}, but got {result}"
|
|
|
|
|
|
def test_struct_with_empty_arrays(
|
|
struct_empty_arrays_blocks, struct_empty_arrays_expected
|
|
):
|
|
# Checks the behavior when concatenating tables with structs containing empty
|
|
# arrays, verifying that null structs are correctly handled.
|
|
|
|
# Concatenate tables with struct columns containing null values
|
|
t3 = concat(struct_empty_arrays_blocks)
|
|
|
|
# Verify that the concatenated result is a valid PyArrow Table
|
|
assert isinstance(t3, pa.Table)
|
|
assert (
|
|
len(t3) == struct_empty_arrays_expected["length"]
|
|
) # Check that the concatenated table has 3 rows
|
|
|
|
# Validate the schema of the resulting concatenated table
|
|
assert (
|
|
t3.schema == struct_empty_arrays_expected["schema"]
|
|
), f"Expected schema: {struct_empty_arrays_expected['schema']}, but got {t3.schema}"
|
|
|
|
# Verify the content of the concatenated table
|
|
assert t3.column("a").to_pylist() == struct_empty_arrays_expected["content"]["a"]
|
|
result = t3.column("d").to_pylist()
|
|
expected = struct_empty_arrays_expected["content"]["d"]
|
|
|
|
assert result == expected, f"Expected {expected}, but got {result}"
|
|
|
|
|
|
def test_struct_with_arrow_variable_shaped_tensor_type(
|
|
struct_variable_shaped_tensor_blocks, struct_variable_shaped_tensor_expected
|
|
):
|
|
# Test concatenating tables with struct columns containing ArrowVariableShapedTensorType
|
|
# fields, ensuring proper handling of variable-shaped tensors within structs.
|
|
|
|
# Concatenate tables with struct columns containing variable-shaped tensors
|
|
t3 = concat(struct_variable_shaped_tensor_blocks)
|
|
assert isinstance(t3, pa.Table)
|
|
assert len(t3) == struct_variable_shaped_tensor_expected["length"]
|
|
|
|
# Validate the schema of the resulting table
|
|
assert (
|
|
t3.schema == struct_variable_shaped_tensor_expected["schema"]
|
|
), f"Expected schema: {struct_variable_shaped_tensor_expected['schema']}, but got {t3.schema}"
|
|
|
|
# Verify the content of the resulting table
|
|
assert (
|
|
t3.column("id").to_pylist()
|
|
== struct_variable_shaped_tensor_expected["content"]["id"]
|
|
)
|
|
|
|
# Check that the struct column contains the expected data
|
|
result_structs = t3.column("struct_with_tensor").to_pylist()
|
|
assert len(result_structs) == 4
|
|
|
|
# Verify each struct contains the correct metadata and tensor data
|
|
expected_metadata = ["row1", "row2", "row3", "row4"]
|
|
for i, (struct, expected_meta) in enumerate(zip(result_structs, expected_metadata)):
|
|
assert struct["metadata"] == expected_meta
|
|
assert isinstance(struct["tensor"], np.ndarray)
|
|
|
|
# Verify tensor shapes match expectations
|
|
if i == 0:
|
|
assert struct["tensor"].shape == (2, 2)
|
|
np.testing.assert_array_equal(
|
|
struct["tensor"], np.ones((2, 2), dtype=np.float32)
|
|
)
|
|
elif i == 1:
|
|
assert struct["tensor"].shape == (3, 3)
|
|
np.testing.assert_array_equal(
|
|
struct["tensor"], np.zeros((3, 3), dtype=np.float32)
|
|
)
|
|
elif i == 2:
|
|
assert struct["tensor"].shape == (1, 4)
|
|
np.testing.assert_array_equal(
|
|
struct["tensor"], np.ones((1, 4), dtype=np.float32)
|
|
)
|
|
elif i == 3:
|
|
assert struct["tensor"].shape == (2, 1)
|
|
np.testing.assert_array_equal(
|
|
struct["tensor"], np.zeros((2, 1), dtype=np.float32)
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < MIN_PYARROW_VERSION_TYPE_PROMOTION,
|
|
reason="Requires PyArrow >= 14.0.0 for type promotion in nested struct fields",
|
|
)
|
|
def test_struct_with_diverging_primitive_types():
|
|
"""Test concatenating tables with struct fields that have diverging primitive types.
|
|
|
|
This tests the scenario where struct fields have the same name but different
|
|
primitive types (e.g., int64 vs float64), which requires type promotion.
|
|
"""
|
|
import pyarrow as pa
|
|
|
|
# Table 1: struct with (a: int64, b: string)
|
|
t1 = pa.table(
|
|
{
|
|
"data": pa.array(
|
|
[{"a": 1, "b": "hello"}, {"a": 2, "b": "world"}],
|
|
type=pa.struct([pa.field("a", pa.int64()), pa.field("b", pa.string())]),
|
|
)
|
|
}
|
|
)
|
|
|
|
# Table 2: struct with (a: float64, c: int32)
|
|
# Field 'a' has different type, field 'b' missing, field 'c' new
|
|
t2 = pa.table(
|
|
{
|
|
"data": pa.array(
|
|
[{"a": 1.5, "c": 100}, {"a": 2.5, "c": 200}],
|
|
type=pa.struct(
|
|
[pa.field("a", pa.float64()), pa.field("c", pa.int32())]
|
|
),
|
|
)
|
|
}
|
|
)
|
|
|
|
# Concatenate with type promotion
|
|
result = concat([t1, t2], promote_types=True)
|
|
|
|
# Verify schema: field 'a' should be promoted to float64
|
|
expected_struct_type = pa.struct(
|
|
[
|
|
pa.field("a", pa.float64()),
|
|
pa.field("b", pa.string()),
|
|
pa.field("c", pa.int32()),
|
|
]
|
|
)
|
|
assert result.schema == pa.schema([pa.field("data", expected_struct_type)])
|
|
|
|
# Verify data: int64 values should be cast to float64, missing fields filled with None
|
|
expected_data = [
|
|
{"a": 1.0, "b": "hello", "c": None},
|
|
{"a": 2.0, "b": "world", "c": None},
|
|
{"a": 1.5, "b": None, "c": 100},
|
|
{"a": 2.5, "b": None, "c": 200},
|
|
]
|
|
assert result.column("data").to_pylist() == expected_data
|
|
|
|
|
|
def test_arrow_concat_object_with_tensor_fails(object_with_tensor_fails_blocks):
|
|
with pytest.raises(ArrowConversionError) as exc_info:
|
|
concat(object_with_tensor_fails_blocks)
|
|
assert "objects and tensors" in str(exc_info.value.__cause__)
|
|
|
|
|
|
def test_unify_schemas(unify_schemas_basic_schemas, unify_schemas_multicol_schemas):
|
|
# Unifying a schema with the same schema as itself
|
|
schemas = unify_schemas_basic_schemas
|
|
assert (
|
|
unify_schemas([schemas["tensor_arr_1"], schemas["tensor_arr_1"]])
|
|
== schemas["tensor_arr_1"]
|
|
)
|
|
|
|
# Single columns with different shapes
|
|
contains_diff_shaped = [schemas["tensor_arr_1"], schemas["tensor_arr_2"]]
|
|
assert unify_schemas(contains_diff_shaped) == pa.schema(
|
|
[
|
|
("tensor_arr", ArrowVariableShapedTensorType(pa.int32(), 2)),
|
|
]
|
|
)
|
|
|
|
# Single columns with same shapes
|
|
contains_diff_types = [schemas["tensor_arr_1"], schemas["tensor_arr_3"]]
|
|
assert unify_schemas(contains_diff_types) == pa.schema(
|
|
[
|
|
("tensor_arr", ArrowTensorType((3, 5), pa.int32())),
|
|
]
|
|
)
|
|
|
|
# Single columns with a variable shaped tensor, same ndim
|
|
contains_var_shaped = [schemas["tensor_arr_1"], schemas["var_tensor_arr"]]
|
|
assert unify_schemas(contains_var_shaped) == pa.schema(
|
|
[
|
|
("tensor_arr", ArrowVariableShapedTensorType(pa.int32(), 2)),
|
|
]
|
|
)
|
|
|
|
# Single columns with a variable shaped tensor, different ndim
|
|
contains_1d2d = [schemas["tensor_arr_1"], schemas["var_tensor_arr_1d"]]
|
|
assert unify_schemas(contains_1d2d) == pa.schema(
|
|
[
|
|
("tensor_arr", ArrowVariableShapedTensorType(pa.int32(), 2)),
|
|
]
|
|
)
|
|
contains_2d3d = [schemas["tensor_arr_1"], schemas["var_tensor_arr_3d"]]
|
|
assert unify_schemas(contains_2d3d) == pa.schema(
|
|
[
|
|
("tensor_arr", ArrowVariableShapedTensorType(pa.int32(), 3)),
|
|
]
|
|
)
|
|
|
|
# Multi-column schemas
|
|
multicol = unify_schemas_multicol_schemas
|
|
assert unify_schemas(
|
|
[multicol["multicol_schema_1"], multicol["multicol_schema_2"]]
|
|
) == pa.schema(
|
|
[
|
|
("col_int", pa.int32()),
|
|
("col_fixed_tensor", ArrowTensorType((4, 2), pa.int32())),
|
|
("col_var_tensor", ArrowVariableShapedTensorType(pa.int16(), 5)),
|
|
]
|
|
)
|
|
|
|
assert unify_schemas(
|
|
[multicol["multicol_schema_1"], multicol["multicol_schema_3"]]
|
|
) == pa.schema(
|
|
[
|
|
("col_int", pa.int32()),
|
|
("col_fixed_tensor", ArrowVariableShapedTensorType(pa.int32(), 3)),
|
|
("col_var_tensor", ArrowVariableShapedTensorType(pa.int16(), 5)),
|
|
]
|
|
)
|
|
|
|
# Unifying >2 schemas together
|
|
assert unify_schemas(
|
|
[
|
|
multicol["multicol_schema_1"],
|
|
multicol["multicol_schema_2"],
|
|
multicol["multicol_schema_3"],
|
|
]
|
|
) == pa.schema(
|
|
[
|
|
("col_int", pa.int32()),
|
|
("col_fixed_tensor", ArrowVariableShapedTensorType(pa.int32(), 3)),
|
|
("col_var_tensor", ArrowVariableShapedTensorType(pa.int16(), 5)),
|
|
]
|
|
)
|
|
|
|
|
|
def test_unify_schemas_object_types(unify_schemas_object_types_schemas):
|
|
"""Test handling of object types (columns_with_objects functionality)."""
|
|
schemas = unify_schemas_object_types_schemas
|
|
|
|
# Should convert to ArrowPythonObjectType
|
|
result = unify_schemas([schemas["object_schema"], schemas["int_schema"]])
|
|
assert result == schemas["expected"]
|
|
|
|
# Test multiple object types
|
|
result = unify_schemas(
|
|
[schemas["object_schema"], schemas["int_schema"], schemas["float_schema"]]
|
|
)
|
|
assert result == schemas["expected"]
|
|
|
|
|
|
def test_unify_schemas_incompatible_tensor_dtypes(
|
|
unify_schemas_incompatible_tensor_schemas,
|
|
):
|
|
"""Test error handling for incompatible tensor dtypes."""
|
|
import pyarrow as pa
|
|
|
|
with pytest.raises(
|
|
pa.lib.ArrowTypeError,
|
|
match=re.escape(
|
|
"Can't unify tensor types with divergent scalar types: [ArrowTensorType(shape=(2, 2), dtype=int32), ArrowTensorType(shape=(2, 2), dtype=float)]"
|
|
),
|
|
):
|
|
unify_schemas(unify_schemas_incompatible_tensor_schemas)
|
|
|
|
|
|
def test_unify_schemas_objects_and_tensors(unify_schemas_objects_and_tensors_schemas):
|
|
"""Test error handling for intersection of objects and tensors."""
|
|
with pytest.raises(ValueError, match="Found columns with both objects and tensors"):
|
|
unify_schemas(unify_schemas_objects_and_tensors_schemas)
|
|
|
|
|
|
def test_unify_schemas_missing_tensor_fields(
|
|
unify_schemas_missing_tensor_fields_schemas,
|
|
):
|
|
"""Test handling of missing tensor fields in structs (has_missing_fields logic)."""
|
|
schemas = unify_schemas_missing_tensor_fields_schemas
|
|
|
|
# Should convert tensor to variable-shaped to accommodate missing field
|
|
result = unify_schemas([schemas["with_tensor"], schemas["without_tensor"]])
|
|
assert result == schemas["expected"]
|
|
|
|
|
|
def test_unify_schemas_nested_struct_tensors(
|
|
unify_schemas_nested_struct_tensors_schemas,
|
|
):
|
|
"""Test handling of nested structs with tensor fields."""
|
|
schemas = unify_schemas_nested_struct_tensors_schemas
|
|
|
|
# Should convert nested tensor to variable-shaped
|
|
result = unify_schemas([schemas["with_tensor"], schemas["without_tensor"]])
|
|
assert result == schemas["expected"]
|
|
|
|
|
|
def test_unify_schemas_edge_cases(unify_schemas_edge_cases_data):
|
|
"""Test edge cases and robustness."""
|
|
data = unify_schemas_edge_cases_data
|
|
|
|
# Empty schema list
|
|
with pytest.raises(Exception): # Should handle gracefully
|
|
unify_schemas(data["empty_schemas"])
|
|
|
|
# Single schema
|
|
assert unify_schemas([data["single_schema"]]) == data["single_schema"]
|
|
|
|
# Schemas with no common columns
|
|
result = unify_schemas(
|
|
[data["no_common_columns"]["schema1"], data["no_common_columns"]["schema2"]]
|
|
)
|
|
assert result == data["no_common_columns"]["expected"]
|
|
|
|
# All null schemas
|
|
result = unify_schemas(
|
|
[data["all_null_schemas"]["schema1"], data["all_null_schemas"]["schema2"]]
|
|
)
|
|
assert result == data["all_null_schemas"]["schema1"]
|
|
|
|
|
|
def test_unify_schemas_mixed_tensor_types(unify_schemas_mixed_tensor_data):
|
|
"""Test handling of mixed tensor types (fixed and variable shaped)."""
|
|
data = unify_schemas_mixed_tensor_data
|
|
|
|
# Should result in variable-shaped tensor
|
|
result = unify_schemas([data["fixed_shape"], data["variable_shaped"]])
|
|
assert result == data["expected_variable"]
|
|
|
|
# Test with different shapes but same dtype
|
|
result = unify_schemas([data["fixed_shape"], data["different_shape"]])
|
|
assert result == data["expected_variable"]
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
get_pyarrow_version() < MIN_PYARROW_VERSION_TYPE_PROMOTION,
|
|
reason="Requires Arrow version of at least 14.0.0",
|
|
)
|
|
def test_unify_schemas_type_promotion(unify_schemas_type_promotion_data):
|
|
data = unify_schemas_type_promotion_data
|
|
|
|
# No type promotion
|
|
assert (
|
|
unify_schemas(
|
|
[data["non_null"], data["nullable"]],
|
|
promote_types=False,
|
|
)
|
|
== data["nullable"]
|
|
)
|
|
|
|
# No type promotion
|
|
with pytest.raises(pa.lib.ArrowTypeError) as exc_info:
|
|
unify_schemas(
|
|
[data["int64"], data["float64"]],
|
|
promote_types=False,
|
|
)
|
|
|
|
assert "Unable to merge: Field A has incompatible types: int64 vs double" == str(
|
|
exc_info.value
|
|
)
|
|
|
|
# Type promoted
|
|
assert (
|
|
unify_schemas(
|
|
[data["int64"], data["float64"]],
|
|
promote_types=True,
|
|
)
|
|
== data["float64"]
|
|
)
|
|
|
|
|
|
def test_arrow_block_select(block_select_data):
|
|
data = block_select_data
|
|
block_accessor = BlockAccessor.for_block(data["table"])
|
|
|
|
block = block_accessor.select(data["single_column"]["columns"])
|
|
assert block.schema == data["single_column"]["expected_schema"]
|
|
assert block.to_pandas().equals(data["df"][data["single_column"]["columns"]])
|
|
|
|
block = block_accessor.select(data["multiple_columns"]["columns"])
|
|
assert block.schema == data["multiple_columns"]["expected_schema"]
|
|
assert block.to_pandas().equals(data["df"][data["multiple_columns"]["columns"]])
|
|
|
|
with pytest.raises(ValueError):
|
|
block = block_accessor.select([lambda x: x % 3, "two"])
|
|
|
|
|
|
def test_arrow_block_slice_copy(block_slice_data):
|
|
# Test that ArrowBlock slicing properly copies the underlying Arrow
|
|
# table.
|
|
def check_for_copy(table1, table2, a, b, is_copy):
|
|
expected_slice = table1.slice(a, b - a)
|
|
assert table2.equals(expected_slice)
|
|
assert table2.schema == table1.schema
|
|
assert table1.num_columns == table2.num_columns
|
|
for col1, col2 in zip(table1.columns, table2.columns):
|
|
assert col1.num_chunks == col2.num_chunks
|
|
for chunk1, chunk2 in zip(col1.chunks, col2.chunks):
|
|
bufs1 = chunk1.buffers()
|
|
bufs2 = chunk2.buffers()
|
|
expected_offset = 0 if is_copy else a
|
|
assert chunk2.offset == expected_offset
|
|
assert len(chunk2) == b - a
|
|
if is_copy:
|
|
assert bufs2[1].address != bufs1[1].address
|
|
else:
|
|
assert bufs2[1].address == bufs1[1].address
|
|
|
|
data = block_slice_data["normal"]
|
|
table = data["table"]
|
|
a, b = data["slice_params"]["a"], data["slice_params"]["b"]
|
|
block_accessor = BlockAccessor.for_block(table)
|
|
|
|
# Test with copy.
|
|
table2 = block_accessor.slice(a, b, True)
|
|
check_for_copy(table, table2, a, b, is_copy=True)
|
|
|
|
# Test without copy.
|
|
table2 = block_accessor.slice(a, b, False)
|
|
check_for_copy(table, table2, a, b, is_copy=False)
|
|
|
|
|
|
def test_arrow_block_slice_copy_empty(block_slice_data):
|
|
# Test that ArrowBlock slicing properly copies the underlying Arrow
|
|
# table when the table is empty.
|
|
data = block_slice_data["empty"]
|
|
table = data["table"]
|
|
a, b = data["slice_params"]["a"], data["slice_params"]["b"]
|
|
expected_slice = table.slice(a, b - a)
|
|
block_accessor = BlockAccessor.for_block(table)
|
|
|
|
# Test with copy.
|
|
table2 = block_accessor.slice(a, b, True)
|
|
assert table2.equals(expected_slice)
|
|
assert table2.schema == table.schema
|
|
assert table2.num_rows == 0
|
|
|
|
# Test without copy.
|
|
table2 = block_accessor.slice(a, b, False)
|
|
assert table2.equals(expected_slice)
|
|
assert table2.schema == table.schema
|
|
assert table2.num_rows == 0
|
|
|
|
|
|
@pytest.mark.parametrize("preserve_order", [True, False])
|
|
def test_mixed_tensor_types_same_dtype(
|
|
mixed_tensor_types_same_dtype_blocks,
|
|
mixed_tensor_types_same_dtype_expected,
|
|
preserve_order,
|
|
):
|
|
"""Test mixed tensor types with same data type but different shapes."""
|
|
|
|
t1, t2 = mixed_tensor_types_same_dtype_blocks
|
|
|
|
t3 = concat([t1, t2], preserve_order=preserve_order)
|
|
assert isinstance(t3, pa.Table)
|
|
assert len(t3) == mixed_tensor_types_same_dtype_expected["length"]
|
|
|
|
# Verify schema - should have tensor field as variable-shaped
|
|
assert t3.schema == mixed_tensor_types_same_dtype_expected["schema"]
|
|
tensor_field = t3.schema.field("tensor")
|
|
assert isinstance(tensor_field.type, ArrowVariableShapedTensorType)
|
|
|
|
# Verify content
|
|
result_tensors = t3.column("tensor").to_pylist()
|
|
assert len(result_tensors) == mixed_tensor_types_same_dtype_expected["length"]
|
|
|
|
expected_tensors = mixed_tensor_types_same_dtype_expected["tensor_values"]
|
|
|
|
if not preserve_order:
|
|
result_tensors = sorted(result_tensors, key=lambda arr: arr.tobytes())
|
|
expected_tensors = sorted(expected_tensors, key=lambda arr: arr.tobytes())
|
|
|
|
for result_tensor, expected_tensor in zip(result_tensors, expected_tensors):
|
|
assert isinstance(result_tensor, np.ndarray)
|
|
assert result_tensor.shape == expected_tensor.shape
|
|
assert result_tensor.dtype == expected_tensor.dtype
|
|
np.testing.assert_array_equal(result_tensor, expected_tensor)
|
|
|
|
|
|
def test_mixed_tensor_types_fixed_shape_different(
|
|
mixed_tensor_types_fixed_shape_blocks, mixed_tensor_types_fixed_shape_expected
|
|
):
|
|
"""Test mixed tensor types with different fixed shapes."""
|
|
|
|
t1, t2 = mixed_tensor_types_fixed_shape_blocks
|
|
|
|
t3 = concat([t1, t2])
|
|
assert isinstance(t3, pa.Table)
|
|
assert len(t3) == mixed_tensor_types_fixed_shape_expected["length"]
|
|
|
|
# Verify schema - should have tensor field as variable-shaped
|
|
assert t3.schema == mixed_tensor_types_fixed_shape_expected["schema"]
|
|
tensor_field = t3.schema.field("tensor")
|
|
assert isinstance(tensor_field.type, ArrowVariableShapedTensorType)
|
|
|
|
# Verify content
|
|
result_tensors = t3.column("tensor").to_pylist()
|
|
assert len(result_tensors) == mixed_tensor_types_fixed_shape_expected["length"]
|
|
|
|
expected_tensors = mixed_tensor_types_fixed_shape_expected["tensor_values"]
|
|
|
|
# Verify each tensor
|
|
for i, (result_tensor, expected_tensor) in enumerate(
|
|
zip(result_tensors, expected_tensors)
|
|
):
|
|
assert isinstance(result_tensor, np.ndarray)
|
|
assert result_tensor.shape == expected_tensor.shape
|
|
assert result_tensor.dtype == expected_tensor.dtype
|
|
np.testing.assert_array_equal(result_tensor, expected_tensor)
|
|
|
|
|
|
def test_mixed_tensor_types_variable_shaped(
|
|
mixed_tensor_types_variable_shaped_blocks,
|
|
mixed_tensor_types_variable_shaped_expected,
|
|
):
|
|
"""Test mixed tensor types with variable-shaped tensors."""
|
|
|
|
t1, t2 = mixed_tensor_types_variable_shaped_blocks
|
|
|
|
t3 = concat([t1, t2])
|
|
assert isinstance(t3, pa.Table)
|
|
assert len(t3) == mixed_tensor_types_variable_shaped_expected["length"]
|
|
|
|
# Verify schema - should have tensor field as variable-shaped
|
|
assert t3.schema == mixed_tensor_types_variable_shaped_expected["schema"]
|
|
tensor_field = t3.schema.field("tensor")
|
|
assert isinstance(tensor_field.type, ArrowVariableShapedTensorType)
|
|
|
|
# Verify content
|
|
result_tensors = t3.column("tensor").to_pylist()
|
|
assert len(result_tensors) == mixed_tensor_types_variable_shaped_expected["length"]
|
|
|
|
expected_tensors = mixed_tensor_types_variable_shaped_expected["tensor_values"]
|
|
|
|
# Verify each tensor
|
|
for i, (result_tensor, expected_tensor) in enumerate(
|
|
zip(result_tensors, expected_tensors)
|
|
):
|
|
assert isinstance(result_tensor, np.ndarray)
|
|
assert result_tensor.shape == expected_tensor.shape
|
|
assert result_tensor.dtype == expected_tensor.dtype
|
|
np.testing.assert_array_equal(result_tensor, expected_tensor)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not _extension_array_concat_supported(),
|
|
reason="ExtensionArrays support concatenation only in Pyarrow >= 12.0",
|
|
)
|
|
@pytest.mark.parametrize("preserve_order", [True, False])
|
|
def test_mixed_tensor_types_in_struct(
|
|
struct_with_mixed_tensor_types_blocks,
|
|
struct_with_mixed_tensor_types_expected,
|
|
preserve_order,
|
|
):
|
|
"""Test that the fix works for mixed tensor types in structs."""
|
|
|
|
t1, t2 = struct_with_mixed_tensor_types_blocks
|
|
|
|
t3 = concat([t1, t2], preserve_order=preserve_order)
|
|
assert isinstance(t3, pa.Table)
|
|
assert len(t3) == struct_with_mixed_tensor_types_expected["length"]
|
|
|
|
# Verify the result has the expected structure
|
|
assert t3.schema == struct_with_mixed_tensor_types_expected["schema"]
|
|
assert "id" in t3.column_names
|
|
assert "struct" in t3.column_names
|
|
|
|
# Verify struct field contains both types of tensors
|
|
struct_data = t3.column("struct").to_pylist()
|
|
assert len(struct_data) == struct_with_mixed_tensor_types_expected["length"]
|
|
|
|
expected_struct_values = struct_with_mixed_tensor_types_expected["struct_values"]
|
|
|
|
if not preserve_order:
|
|
# Sort both by the "id" column so we can compare element-by-element.
|
|
ids = t3.column("id").to_pylist()
|
|
order = sorted(range(len(ids)), key=lambda i: ids[i])
|
|
struct_data = [struct_data[i] for i in order]
|
|
|
|
# Verify struct values
|
|
for i, (struct_row, expected_values) in enumerate(
|
|
zip(struct_data, expected_struct_values)
|
|
):
|
|
for key, expected_value in expected_values.items():
|
|
assert struct_row[key] == expected_value
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not _extension_array_concat_supported(),
|
|
reason="ExtensionArrays support concatenation only in Pyarrow >= 12.0",
|
|
)
|
|
def test_nested_struct_with_mixed_tensor_types(
|
|
nested_struct_with_mixed_tensor_types_blocks,
|
|
nested_struct_with_mixed_tensor_types_expected,
|
|
):
|
|
"""Test nested structs with mixed tensor types at different levels."""
|
|
|
|
t1, t2 = nested_struct_with_mixed_tensor_types_blocks
|
|
|
|
t3 = concat([t1, t2])
|
|
assert isinstance(t3, pa.Table)
|
|
assert len(t3) == nested_struct_with_mixed_tensor_types_expected["length"]
|
|
|
|
# Verify the result has the expected structure
|
|
assert t3.schema == nested_struct_with_mixed_tensor_types_expected["schema"]
|
|
assert "id" in t3.column_names
|
|
assert "complex_struct" in t3.column_names
|
|
|
|
# Verify nested struct field contains both types of tensors
|
|
struct_data = t3.column("complex_struct").to_pylist()
|
|
assert len(struct_data) == nested_struct_with_mixed_tensor_types_expected["length"]
|
|
|
|
expected_fields = nested_struct_with_mixed_tensor_types_expected["expected_fields"]
|
|
|
|
# Check that nested structures are preserved
|
|
for field in expected_fields:
|
|
if field in ["nested", "outer_tensor", "outer_value"]:
|
|
assert field in struct_data[0]
|
|
elif field in ["inner_tensor", "inner_value"]:
|
|
assert field in struct_data[0]["nested"]
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not _extension_array_concat_supported(),
|
|
reason="ExtensionArrays support concatenation only in Pyarrow >= 12.0",
|
|
)
|
|
def test_multiple_tensor_fields_in_struct(
|
|
multiple_tensor_fields_struct_blocks, multiple_tensor_fields_struct_expected
|
|
):
|
|
"""Test structs with multiple tensor fields of different types."""
|
|
|
|
t1, t2 = multiple_tensor_fields_struct_blocks
|
|
|
|
t3 = concat([t1, t2])
|
|
assert isinstance(t3, pa.Table)
|
|
assert len(t3) == multiple_tensor_fields_struct_expected["length"]
|
|
|
|
# Verify the result has the expected structure
|
|
assert t3.schema == multiple_tensor_fields_struct_expected["schema"]
|
|
assert "id" in t3.column_names
|
|
assert "multi_tensor_struct" in t3.column_names
|
|
|
|
# Verify struct field contains both types of tensors
|
|
struct_data = t3.column("multi_tensor_struct").to_pylist()
|
|
assert len(struct_data) == multiple_tensor_fields_struct_expected["length"]
|
|
|
|
expected_fields = multiple_tensor_fields_struct_expected["expected_fields"]
|
|
|
|
# Check that all tensor fields are present
|
|
for row in struct_data:
|
|
for field in expected_fields:
|
|
assert field in row
|
|
|
|
|
|
def test_struct_with_incompatible_tensor_dtypes_fails():
|
|
"""Test that concatenating structs with incompatible tensor dtypes fails gracefully."""
|
|
|
|
# Block 1: Struct with float32 fixed-shape tensor
|
|
tensor_data1 = np.ones((2, 2), dtype=np.float32)
|
|
|
|
# Block 2: Struct with int64 variable-shaped tensor (different dtype)
|
|
tensor_data2 = np.array(
|
|
[
|
|
np.ones((3, 3), dtype=np.int64),
|
|
np.zeros((1, 4), dtype=np.int64),
|
|
],
|
|
dtype=object,
|
|
)
|
|
|
|
t1, t2 = _create_struct_tensor_blocks(
|
|
tensor_data1, tensor_data2, "fixed", "variable"
|
|
)
|
|
|
|
# This should fail because of incompatible tensor dtypes
|
|
with pytest.raises(
|
|
ArrowConversionError,
|
|
match=re.escape(
|
|
"Can't unify tensor types with divergent scalar types: [ArrowTensorTypeV2(shape=(2,), dtype=float), ArrowVariableShapedTensorType(ndim=2, dtype=int64)]"
|
|
),
|
|
):
|
|
concat([t1, t2])
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not _extension_array_concat_supported(),
|
|
reason="ExtensionArrays support concatenation only in Pyarrow >= 12.0",
|
|
)
|
|
@pytest.mark.parametrize("preserve_order", [True, False])
|
|
def test_struct_with_additional_fields(
|
|
struct_with_additional_fields_blocks,
|
|
struct_with_additional_fields_expected,
|
|
preserve_order,
|
|
):
|
|
"""Test structs where some blocks have additional fields."""
|
|
|
|
t1, t2 = struct_with_additional_fields_blocks
|
|
|
|
t3 = concat([t1, t2], preserve_order=preserve_order)
|
|
assert isinstance(t3, pa.Table)
|
|
assert len(t3) == struct_with_additional_fields_expected["length"]
|
|
|
|
# Verify the result has the expected structure
|
|
assert t3.schema == struct_with_additional_fields_expected["schema"]
|
|
assert "id" in t3.column_names
|
|
assert "struct" in t3.column_names
|
|
|
|
# Verify struct field contains both types of tensors
|
|
ids = t3.column("id").to_pylist()
|
|
struct_data = t3.column("struct").to_pylist()
|
|
assert len(struct_data) == struct_with_additional_fields_expected["length"]
|
|
|
|
field_presence = struct_with_additional_fields_expected["field_presence"]
|
|
extra_values = struct_with_additional_fields_expected["extra_values"]
|
|
|
|
if not preserve_order:
|
|
order = sorted(range(len(ids)), key=lambda i: ids[i])
|
|
struct_data = [struct_data[i] for i in order]
|
|
|
|
# Check field presence and values
|
|
for i, row in enumerate(struct_data):
|
|
for field, should_be_present in field_presence.items():
|
|
assert (field in row) == should_be_present
|
|
|
|
# Check extra field values
|
|
if "extra" in row:
|
|
assert row["extra"] == extra_values[i]
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not _extension_array_concat_supported(),
|
|
reason="ExtensionArrays support concatenation only in Pyarrow >= 12.0",
|
|
)
|
|
def test_struct_with_null_tensor_values(
|
|
struct_with_null_tensor_values_blocks, struct_with_null_tensor_values_expected
|
|
):
|
|
"""Test structs where some fields are missing and get filled with nulls."""
|
|
|
|
t1, t2 = struct_with_null_tensor_values_blocks
|
|
|
|
t3 = concat([t1, t2])
|
|
assert isinstance(t3, pa.Table)
|
|
assert len(t3) == struct_with_null_tensor_values_expected["length"]
|
|
|
|
# Validate schema - should have both fields
|
|
assert t3.schema == struct_with_null_tensor_values_expected["schema"]
|
|
|
|
# Validate result
|
|
assert t3.column("id").to_pylist() == struct_with_null_tensor_values_expected["ids"]
|
|
|
|
# Check the struct column directly to avoid the Arrow tensor extension null bug
|
|
struct_column = t3.column("struct")
|
|
expected_values = struct_with_null_tensor_values_expected["values"]
|
|
expected_tensor_validity = struct_with_null_tensor_values_expected[
|
|
"tensor_validity"
|
|
]
|
|
|
|
# Check each row
|
|
for i, (expected_value, expected_valid) in enumerate(
|
|
zip(expected_values, expected_tensor_validity)
|
|
):
|
|
assert struct_column[i]["value"].as_py() == expected_value
|
|
|
|
if expected_valid:
|
|
assert struct_column[i]["tensor"] is not None
|
|
else:
|
|
# Check that the tensor field is null by checking its validity
|
|
tensor_field = struct_column[i]["tensor"]
|
|
assert tensor_field.is_valid is False
|
|
|
|
|
|
# Test fixtures for _align_struct_fields tests
|
|
@pytest.fixture
|
|
def simple_struct_blocks():
|
|
"""Fixture for simple struct blocks with missing fields."""
|
|
# Block 1: Struct with fields 'a' and 'b'
|
|
struct_data1 = [{"a": 1, "b": "x"}, {"a": 2, "b": "y"}]
|
|
|
|
# Block 2: Struct with fields 'a' and 'c' (missing 'b', has 'c')
|
|
struct_data2 = [{"a": 3, "c": True}, {"a": 4, "c": False}]
|
|
|
|
return _create_basic_struct_blocks(
|
|
struct_data1, struct_data2, id_data1=None, id_data2=None
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def simple_struct_schema():
|
|
"""Fixture for simple struct schema with all fields."""
|
|
struct_fields = [("a", pa.int64()), ("b", pa.string()), ("c", pa.bool_())]
|
|
return _create_struct_schema(struct_fields, include_id=False)
|
|
|
|
|
|
@pytest.fixture
|
|
def nested_struct_blocks():
|
|
"""Fixture for nested struct blocks with missing fields."""
|
|
# Block 1: Nested struct with inner fields 'x' and 'y'
|
|
struct_data1 = [{"inner": {"x": 1, "y": "a"}}, {"inner": {"x": 2, "y": "b"}}]
|
|
|
|
# Block 2: Nested struct with inner fields 'x' and 'z' (missing 'y', has 'z')
|
|
struct_data2 = [{"inner": {"x": 3, "z": 1.5}}, {"inner": {"x": 4, "z": 2.5}}]
|
|
|
|
return _create_basic_struct_blocks(
|
|
struct_data1, struct_data2, column_name="outer", id_data1=None, id_data2=None
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def nested_struct_schema():
|
|
"""Fixture for nested struct schema with all fields."""
|
|
inner_fields = [("x", pa.int64()), ("y", pa.string()), ("z", pa.float64())]
|
|
struct_fields = [("inner", pa.struct(inner_fields))]
|
|
return _create_struct_schema(
|
|
struct_fields,
|
|
include_id=False,
|
|
other_fields=[("outer", pa.struct(struct_fields))],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def missing_column_blocks():
|
|
"""Fixture for blocks where one is missing a struct column entirely."""
|
|
# Block 1: Has struct column
|
|
t1 = pa.table(
|
|
{
|
|
"struct": pa.array([{"a": 1, "b": "x"}, {"a": 2, "b": "y"}]),
|
|
"other": pa.array([10, 20]),
|
|
}
|
|
)
|
|
|
|
# Block 2: Missing struct column entirely
|
|
t2 = pa.table({"other": pa.array([30, 40])})
|
|
|
|
return t1, t2
|
|
|
|
|
|
@pytest.fixture
|
|
def missing_column_schema():
|
|
"""Fixture for schema with struct column that may be missing."""
|
|
return pa.schema(
|
|
[
|
|
("struct", pa.struct([("a", pa.int64()), ("b", pa.string())])),
|
|
("other", pa.int64()),
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def multiple_struct_blocks():
|
|
"""Fixture for blocks with multiple struct columns."""
|
|
# Block 1: Two struct columns with different field sets
|
|
struct1_data1 = [{"a": 1, "b": "x"}, {"a": 2, "b": "y"}]
|
|
struct2_data1 = [{"p": 10, "q": True}, {"p": 20, "q": False}]
|
|
|
|
# Block 2: Same struct columns but with different/missing fields
|
|
struct1_data2 = [{"a": 3, "c": 1.5}, {"a": 4, "c": 2.5}] # missing 'b', has 'c'
|
|
struct2_data2 = [
|
|
{"p": 30, "r": "alpha"},
|
|
{"p": 40, "r": "beta"},
|
|
] # missing 'q', has 'r'
|
|
|
|
t1 = pa.table(
|
|
{
|
|
"struct1": pa.array(struct1_data1),
|
|
"struct2": pa.array(struct2_data1),
|
|
}
|
|
)
|
|
|
|
t2 = pa.table(
|
|
{
|
|
"struct1": pa.array(struct1_data2),
|
|
"struct2": pa.array(struct2_data2),
|
|
}
|
|
)
|
|
|
|
return t1, t2
|
|
|
|
|
|
@pytest.fixture
|
|
def multiple_struct_schema():
|
|
"""Fixture for schema with multiple struct columns."""
|
|
struct1_fields = [("a", pa.int64()), ("b", pa.string()), ("c", pa.float64())]
|
|
struct2_fields = [("p", pa.int64()), ("q", pa.bool_()), ("r", pa.string())]
|
|
|
|
return pa.schema(
|
|
[
|
|
("struct1", pa.struct(struct1_fields)),
|
|
("struct2", pa.struct(struct2_fields)),
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mixed_column_blocks():
|
|
"""Fixture for blocks with mix of struct and non-struct columns."""
|
|
# Block 1: Mix of struct and non-struct columns
|
|
struct_data1 = [{"a": 1, "b": "x"}, {"a": 2, "b": "y"}]
|
|
int_col1 = [10, 20]
|
|
string_col1 = ["foo", "bar"]
|
|
|
|
# Block 2: Same structure
|
|
struct_data2 = [{"a": 3, "c": True}, {"a": 4, "c": False}] # missing 'b', has 'c'
|
|
int_col2 = [30, 40]
|
|
string_col2 = ["baz", "qux"]
|
|
|
|
t1 = pa.table(
|
|
{
|
|
"struct": pa.array(struct_data1),
|
|
"int_col": pa.array(int_col1),
|
|
"string_col": pa.array(string_col1),
|
|
}
|
|
)
|
|
|
|
t2 = pa.table(
|
|
{
|
|
"struct": pa.array(struct_data2),
|
|
"int_col": pa.array(int_col2),
|
|
"string_col": pa.array(string_col2),
|
|
}
|
|
)
|
|
|
|
return t1, t2
|
|
|
|
|
|
@pytest.fixture
|
|
def mixed_column_schema():
|
|
"""Fixture for schema with mix of struct and non-struct columns."""
|
|
struct_fields = [("a", pa.int64()), ("b", pa.string()), ("c", pa.bool_())]
|
|
|
|
return pa.schema(
|
|
[
|
|
("struct", pa.struct(struct_fields)),
|
|
("int_col", pa.int64()),
|
|
("string_col", pa.string()),
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def empty_block_blocks():
|
|
"""Fixture for blocks where one is empty."""
|
|
# Empty block
|
|
empty_struct_type = pa.struct([("a", pa.int64()), ("b", pa.string())])
|
|
t1 = pa.table({"struct": pa.array([], type=empty_struct_type)})
|
|
|
|
# Non-empty block
|
|
struct_data2 = [{"a": 1, "c": True}, {"a": 2, "c": False}] # missing 'b', has 'c'
|
|
t2 = pa.table({"struct": pa.array(struct_data2)})
|
|
|
|
return t1, t2
|
|
|
|
|
|
@pytest.fixture
|
|
def empty_block_schema():
|
|
"""Fixture for schema used with empty blocks."""
|
|
struct_fields = [("a", pa.int64()), ("b", pa.string()), ("c", pa.bool_())]
|
|
return _create_struct_schema(struct_fields, include_id=False)
|
|
|
|
|
|
@pytest.fixture
|
|
def already_aligned_blocks():
|
|
"""Fixture for blocks that are already aligned."""
|
|
# Both blocks have identical schemas
|
|
struct_data1 = [{"a": 1, "b": "x"}, {"a": 2, "b": "y"}]
|
|
struct_data2 = [{"a": 3, "b": "z"}, {"a": 4, "b": "w"}]
|
|
|
|
return _create_basic_struct_blocks(
|
|
struct_data1, struct_data2, id_data1=None, id_data2=None
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def already_aligned_schema():
|
|
"""Fixture for schema used with already aligned blocks."""
|
|
struct_fields = [("a", pa.int64()), ("b", pa.string())]
|
|
return _create_struct_schema(struct_fields, include_id=False)
|
|
|
|
|
|
@pytest.fixture
|
|
def no_struct_blocks():
|
|
"""Fixture for blocks with no struct columns."""
|
|
# Blocks with no struct columns
|
|
int_col1 = [1, 2]
|
|
string_col1 = ["a", "b"]
|
|
int_col2 = [3, 4]
|
|
string_col2 = ["c", "d"]
|
|
|
|
t1 = pa.table({"int_col": pa.array(int_col1), "string_col": pa.array(string_col1)})
|
|
t2 = pa.table({"int_col": pa.array(int_col2), "string_col": pa.array(string_col2)})
|
|
|
|
return t1, t2
|
|
|
|
|
|
@pytest.fixture
|
|
def no_struct_schema():
|
|
"""Fixture for schema with no struct columns."""
|
|
return pa.schema([("int_col", pa.int64()), ("string_col", pa.string())])
|
|
|
|
|
|
@pytest.fixture
|
|
def deep_nesting_blocks():
|
|
"""Fixture for blocks with deeply nested structs."""
|
|
# Block 1: Deeply nested struct
|
|
struct_data1 = [
|
|
{"level2": {"level3": {"a": 1, "b": "x"}}},
|
|
{"level2": {"level3": {"a": 2, "b": "y"}}},
|
|
]
|
|
|
|
# Block 2: Same structure but missing some fields
|
|
struct_data2 = [
|
|
{"level2": {"level3": {"a": 3, "c": True}}}, # missing 'b', has 'c'
|
|
{"level2": {"level3": {"a": 4, "c": False}}},
|
|
]
|
|
|
|
return _create_basic_struct_blocks(
|
|
struct_data1, struct_data2, column_name="level1", id_data1=None, id_data2=None
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def deep_nesting_schema():
|
|
"""Fixture for schema with deeply nested structs."""
|
|
level3_fields = [("a", pa.int64()), ("b", pa.string()), ("c", pa.bool_())]
|
|
level2_fields = [("level3", pa.struct(level3_fields))]
|
|
level1_fields = [("level2", pa.struct(level2_fields))]
|
|
|
|
return pa.schema([("level1", pa.struct(level1_fields))])
|
|
|
|
|
|
def test_align_struct_fields_simple(simple_struct_blocks, simple_struct_schema):
|
|
"""Test basic struct field alignment with missing fields."""
|
|
t1, t2 = simple_struct_blocks
|
|
|
|
aligned_blocks = _align_struct_fields([t1, t2], simple_struct_schema)
|
|
|
|
assert len(aligned_blocks) == 2
|
|
|
|
# Check first block - should have 'c' field filled with None
|
|
result1 = aligned_blocks[0]
|
|
assert result1.schema == simple_struct_schema
|
|
assert result1["struct"].to_pylist() == [
|
|
{"a": 1, "b": "x", "c": None},
|
|
{"a": 2, "b": "y", "c": None},
|
|
]
|
|
|
|
# Check second block - should have 'b' field filled with None
|
|
result2 = aligned_blocks[1]
|
|
assert result2.schema == simple_struct_schema
|
|
assert result2["struct"].to_pylist() == [
|
|
{"a": 3, "b": None, "c": True},
|
|
{"a": 4, "b": None, "c": False},
|
|
]
|
|
|
|
|
|
def test_align_struct_fields_nested(nested_struct_blocks, nested_struct_schema):
|
|
"""Test nested struct field alignment."""
|
|
t1, t2 = nested_struct_blocks
|
|
|
|
aligned_blocks = _align_struct_fields([t1, t2], nested_struct_schema)
|
|
|
|
assert len(aligned_blocks) == 2
|
|
|
|
# Check first block - should have 'z' field filled with None
|
|
result1 = aligned_blocks[0]
|
|
assert result1.schema == nested_struct_schema
|
|
assert result1["outer"].to_pylist() == [
|
|
{"inner": {"x": 1, "y": "a", "z": None}},
|
|
{"inner": {"x": 2, "y": "b", "z": None}},
|
|
]
|
|
|
|
# Check second block - should have 'y' field filled with None
|
|
result2 = aligned_blocks[1]
|
|
assert result2.schema == nested_struct_schema
|
|
assert result2["outer"].to_pylist() == [
|
|
{"inner": {"x": 3, "y": None, "z": 1.5}},
|
|
{"inner": {"x": 4, "y": None, "z": 2.5}},
|
|
]
|
|
|
|
|
|
def test_align_struct_fields_missing_column(
|
|
missing_column_blocks, missing_column_schema
|
|
):
|
|
"""Test alignment when a struct column is missing from some blocks."""
|
|
t1, t2 = missing_column_blocks
|
|
|
|
aligned_blocks = _align_struct_fields([t1, t2], missing_column_schema)
|
|
|
|
assert len(aligned_blocks) == 2
|
|
|
|
# Check first block - should be unchanged
|
|
result1 = aligned_blocks[0]
|
|
assert result1.schema == missing_column_schema
|
|
assert result1["struct"].to_pylist() == [{"a": 1, "b": "x"}, {"a": 2, "b": "y"}]
|
|
assert result1["other"].to_pylist() == [10, 20]
|
|
|
|
# Check second block - should have null struct column
|
|
result2 = aligned_blocks[1]
|
|
assert result2.schema == missing_column_schema
|
|
assert result2["struct"].to_pylist() == [None, None]
|
|
assert result2["other"].to_pylist() == [30, 40]
|
|
|
|
|
|
def test_align_struct_fields_multiple_structs(
|
|
multiple_struct_blocks, multiple_struct_schema
|
|
):
|
|
"""Test alignment with multiple struct columns."""
|
|
t1, t2 = multiple_struct_blocks
|
|
|
|
aligned_blocks = _align_struct_fields([t1, t2], multiple_struct_schema)
|
|
|
|
assert len(aligned_blocks) == 2
|
|
|
|
# Check first block
|
|
result1 = aligned_blocks[0]
|
|
assert result1.schema == multiple_struct_schema
|
|
assert result1["struct1"].to_pylist() == [
|
|
{"a": 1, "b": "x", "c": None},
|
|
{"a": 2, "b": "y", "c": None},
|
|
]
|
|
assert result1["struct2"].to_pylist() == [
|
|
{"p": 10, "q": True, "r": None},
|
|
{"p": 20, "q": False, "r": None},
|
|
]
|
|
|
|
# Check second block
|
|
result2 = aligned_blocks[1]
|
|
assert result2.schema == multiple_struct_schema
|
|
assert result2["struct1"].to_pylist() == [
|
|
{"a": 3, "b": None, "c": 1.5},
|
|
{"a": 4, "b": None, "c": 2.5},
|
|
]
|
|
assert result2["struct2"].to_pylist() == [
|
|
{"p": 30, "q": None, "r": "alpha"},
|
|
{"p": 40, "q": None, "r": "beta"},
|
|
]
|
|
|
|
|
|
def test_align_struct_fields_non_struct_columns(
|
|
mixed_column_blocks, mixed_column_schema
|
|
):
|
|
"""Test that non-struct columns are left unchanged."""
|
|
t1, t2 = mixed_column_blocks
|
|
|
|
aligned_blocks = _align_struct_fields([t1, t2], mixed_column_schema)
|
|
|
|
assert len(aligned_blocks) == 2
|
|
|
|
# Check that non-struct columns are unchanged
|
|
for i, block in enumerate(aligned_blocks):
|
|
assert block["int_col"].to_pylist() == [10 + i * 20, 20 + i * 20]
|
|
assert (
|
|
block["string_col"].to_pylist() == ["foo", "bar"]
|
|
if i == 0
|
|
else ["baz", "qux"]
|
|
)
|
|
|
|
|
|
def test_align_struct_fields_empty_blocks(empty_block_blocks, empty_block_schema):
|
|
"""Test alignment with empty blocks."""
|
|
t1, t2 = empty_block_blocks
|
|
|
|
aligned_blocks = _align_struct_fields([t1, t2], empty_block_schema)
|
|
|
|
assert len(aligned_blocks) == 2
|
|
|
|
# Check empty block
|
|
result1 = aligned_blocks[0]
|
|
assert result1.schema == empty_block_schema
|
|
assert len(result1) == 0
|
|
|
|
# Check non-empty block
|
|
result2 = aligned_blocks[1]
|
|
assert result2.schema == empty_block_schema
|
|
assert result2["struct"].to_pylist() == [
|
|
{"a": 1, "b": None, "c": True},
|
|
{"a": 2, "b": None, "c": False},
|
|
]
|
|
|
|
|
|
def test_align_struct_fields_already_aligned(
|
|
already_aligned_blocks, already_aligned_schema
|
|
):
|
|
"""Test that already aligned blocks are returned unchanged."""
|
|
t1, t2 = already_aligned_blocks
|
|
|
|
aligned_blocks = _align_struct_fields([t1, t2], already_aligned_schema)
|
|
|
|
# Should return the original blocks unchanged
|
|
assert aligned_blocks == [t1, t2]
|
|
|
|
|
|
def test_align_struct_fields_no_struct_columns(no_struct_blocks, no_struct_schema):
|
|
"""Test alignment when there are no struct columns in the schema."""
|
|
t1, t2 = no_struct_blocks
|
|
|
|
aligned_blocks = _align_struct_fields([t1, t2], no_struct_schema)
|
|
|
|
# Should return the original blocks unchanged
|
|
assert aligned_blocks == [t1, t2]
|
|
|
|
|
|
def test_align_struct_fields_deep_nesting(deep_nesting_blocks, deep_nesting_schema):
|
|
"""Test alignment with deeply nested structs."""
|
|
t1, t2 = deep_nesting_blocks
|
|
|
|
aligned_blocks = _align_struct_fields([t1, t2], deep_nesting_schema)
|
|
|
|
assert len(aligned_blocks) == 2
|
|
|
|
# Check first block - should have 'c' field filled with None
|
|
result1 = aligned_blocks[0]
|
|
assert result1.schema == deep_nesting_schema
|
|
assert result1["level1"].to_pylist() == [
|
|
{"level2": {"level3": {"a": 1, "b": "x", "c": None}}},
|
|
{"level2": {"level3": {"a": 2, "b": "y", "c": None}}},
|
|
]
|
|
|
|
# Check second block - should have 'b' field filled with None
|
|
result2 = aligned_blocks[1]
|
|
assert result2.schema == deep_nesting_schema
|
|
assert result2["level1"].to_pylist() == [
|
|
{"level2": {"level3": {"a": 3, "b": None, "c": True}}},
|
|
{"level2": {"level3": {"a": 4, "b": None, "c": False}}},
|
|
]
|
|
|
|
|
|
# Test fixtures for tensor-related tests
|
|
@pytest.fixture
|
|
def uniform_tensor_blocks(tensor_format_context):
|
|
"""Fixture for uniform tensor blocks with same shape."""
|
|
# Block 1: Fixed shape tensors (2x2)
|
|
a1 = np.arange(12).reshape((3, 2, 2))
|
|
t1 = pa.table({"a": ArrowTensorArray.from_numpy(a1)})
|
|
|
|
# Block 2: Fixed shape tensors (2x2)
|
|
a2 = np.arange(12, 24).reshape((3, 2, 2))
|
|
t2 = pa.table({"a": ArrowTensorArray.from_numpy(a2)})
|
|
|
|
return t1, t2
|
|
|
|
|
|
@pytest.fixture
|
|
def uniform_tensor_expected(tensor_format_context):
|
|
"""Fixture for expected results from uniform tensor concatenation."""
|
|
t = create_arrow_fixed_shape_tensor_type((2, 2), pa.int64())
|
|
expected_schema = pa.schema([("a", t)])
|
|
expected_length = 6
|
|
expected_chunks = 2
|
|
|
|
# Expected content
|
|
a1 = np.arange(12).reshape((3, 2, 2))
|
|
a2 = np.arange(12, 24).reshape((3, 2, 2))
|
|
|
|
return {
|
|
"schema": expected_schema,
|
|
"length": expected_length,
|
|
"chunks": expected_chunks,
|
|
"content": [a1, a2],
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def variable_shaped_tensor_blocks():
|
|
"""Fixture for variable-shaped tensor blocks."""
|
|
# Block 1: Variable shape tensors
|
|
a1 = np.array(
|
|
[np.arange(4).reshape((2, 2)), np.arange(4, 13).reshape((3, 3))], dtype=object
|
|
)
|
|
t1 = pa.table({"a": ArrowTensorArray.from_numpy(a1)})
|
|
|
|
# Block 2: Variable shape tensors
|
|
a2 = np.array(
|
|
[np.arange(4).reshape((2, 2)), np.arange(4, 13).reshape((3, 3))], dtype=object
|
|
)
|
|
t2 = pa.table({"a": ArrowTensorArray.from_numpy(a2)})
|
|
|
|
return t1, t2
|
|
|
|
|
|
@pytest.fixture
|
|
def variable_shaped_tensor_expected():
|
|
"""Fixture for expected results from variable-shaped tensor concatenation."""
|
|
expected_schema = pa.schema([("a", ArrowVariableShapedTensorType(pa.int64(), 2))])
|
|
expected_length = 4
|
|
expected_chunks = 2
|
|
|
|
# Expected content
|
|
a1 = np.array(
|
|
[np.arange(4).reshape((2, 2)), np.arange(4, 13).reshape((3, 3))], dtype=object
|
|
)
|
|
a2 = np.array(
|
|
[np.arange(4).reshape((2, 2)), np.arange(4, 13).reshape((3, 3))], dtype=object
|
|
)
|
|
|
|
return {
|
|
"schema": expected_schema,
|
|
"length": expected_length,
|
|
"chunks": expected_chunks,
|
|
"content": [a1, a2],
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def mixed_tensor_blocks():
|
|
"""Fixture for mixed fixed-shape and variable-shaped tensor blocks."""
|
|
# Block 1: Fixed shape tensors
|
|
a1 = np.arange(12).reshape((3, 2, 2))
|
|
t1 = pa.table({"a": ArrowTensorArray.from_numpy(a1)})
|
|
|
|
# Block 2: Variable shape tensors
|
|
a2 = np.array(
|
|
[np.arange(4).reshape((2, 2)), np.arange(4, 13).reshape((3, 3))], dtype=object
|
|
)
|
|
t2 = pa.table({"a": ArrowTensorArray.from_numpy(a2)})
|
|
|
|
return t1, t2
|
|
|
|
|
|
@pytest.fixture
|
|
def mixed_tensor_expected():
|
|
"""Fixture for expected results from mixed tensor concatenation."""
|
|
expected_schema = pa.schema([("a", ArrowVariableShapedTensorType(pa.int64(), 2))])
|
|
expected_length = 5
|
|
expected_chunks = 2
|
|
|
|
# Expected content
|
|
a1 = np.arange(12).reshape((3, 2, 2))
|
|
a2 = np.array(
|
|
[np.arange(4).reshape((2, 2)), np.arange(4, 13).reshape((3, 3))], dtype=object
|
|
)
|
|
|
|
return {
|
|
"schema": expected_schema,
|
|
"length": expected_length,
|
|
"chunks": expected_chunks,
|
|
"content": [a1, a2],
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def different_shape_tensor_blocks():
|
|
"""Fixture for tensor blocks with different fixed shapes."""
|
|
# Block 1: Fixed shape tensors (2x2)
|
|
a1 = np.arange(12).reshape((3, 2, 2))
|
|
t1 = pa.table({"a": ArrowTensorArray.from_numpy(a1)})
|
|
|
|
# Block 2: Fixed shape tensors (3x3)
|
|
a2 = np.arange(12, 39).reshape((3, 3, 3))
|
|
t2 = pa.table({"a": ArrowTensorArray.from_numpy(a2)})
|
|
|
|
return t1, t2
|
|
|
|
|
|
@pytest.fixture
|
|
def different_shape_tensor_expected():
|
|
"""Fixture for expected results from different shape tensor concatenation."""
|
|
expected_schema = pa.schema([("a", ArrowVariableShapedTensorType(pa.int64(), 2))])
|
|
expected_length = 6
|
|
expected_chunks = 2
|
|
|
|
# Expected content
|
|
a1 = np.arange(12).reshape((3, 2, 2))
|
|
a2 = np.arange(12, 39).reshape((3, 3, 3))
|
|
|
|
return {
|
|
"schema": expected_schema,
|
|
"length": expected_length,
|
|
"chunks": expected_chunks,
|
|
"content": [a1, a2],
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def mixed_tensor_types_same_dtype_blocks():
|
|
"""Fixture for mixed tensor types with same dtype but different shapes."""
|
|
# Block 1: Fixed shape tensors with float32
|
|
tensor_data1 = np.ones((2, 2), dtype=np.float32)
|
|
|
|
# Block 2: Variable shape tensors with float32
|
|
tensor_data2 = np.array(
|
|
[
|
|
np.ones((3, 3), dtype=np.float32),
|
|
np.zeros((1, 4), dtype=np.float32),
|
|
],
|
|
dtype=object,
|
|
)
|
|
|
|
return _create_tensor_blocks(tensor_data1, tensor_data2, "fixed", "variable")
|
|
|
|
|
|
@pytest.fixture
|
|
def mixed_tensor_types_same_dtype_expected():
|
|
"""Fixture for expected results from mixed tensor types with same dtype."""
|
|
expected_schema = _create_tensor_schema(struct_name="tensor")
|
|
expected_tensors = [
|
|
# First 2 were converted to var-shaped with their shape expanded
|
|
# with singleton axis: from (2,) to (1, 2)
|
|
np.ones((1, 2), dtype=np.float32),
|
|
np.ones((1, 2), dtype=np.float32),
|
|
# Last 2 were left intact
|
|
np.ones((3, 3), dtype=np.float32),
|
|
np.zeros((1, 4), dtype=np.float32),
|
|
]
|
|
|
|
return _create_expected_result(expected_schema, 4, tensor_values=expected_tensors)
|
|
|
|
|
|
@pytest.fixture
|
|
def mixed_tensor_types_fixed_shape_blocks():
|
|
"""Fixture for mixed tensor types with different fixed shapes."""
|
|
# Block 1: Fixed shape tensors (2x2)
|
|
tensor_data1 = np.ones((2, 2), dtype=np.float32)
|
|
|
|
# Block 2: Fixed shape tensors (3x3)
|
|
tensor_data2 = np.zeros((3, 3), dtype=np.float32)
|
|
|
|
return _create_tensor_blocks(
|
|
tensor_data1, tensor_data2, "fixed", "fixed", id_data2=[3, 4, 5]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mixed_tensor_types_fixed_shape_expected():
|
|
"""Fixture for expected results from mixed tensor types with different fixed shapes."""
|
|
expected_schema = _create_tensor_schema(struct_name="tensor", ndim=1)
|
|
expected_tensors = [
|
|
np.ones((2,), dtype=np.float32), # First 2 converted to variable-shaped
|
|
np.ones((2,), dtype=np.float32),
|
|
np.zeros((3,), dtype=np.float32), # Last 3 variable-shaped
|
|
np.zeros((3,), dtype=np.float32),
|
|
np.zeros((3,), dtype=np.float32),
|
|
]
|
|
|
|
return _create_expected_result(expected_schema, 5, tensor_values=expected_tensors)
|
|
|
|
|
|
@pytest.fixture
|
|
def mixed_tensor_types_variable_shaped_blocks():
|
|
"""Fixture for mixed tensor types with variable-shaped tensors."""
|
|
# Block 1: Variable shape tensors
|
|
tensor_data1 = np.array(
|
|
[
|
|
np.ones((2, 2), dtype=np.float32),
|
|
np.zeros((3, 3), dtype=np.float32),
|
|
],
|
|
dtype=object,
|
|
)
|
|
|
|
# Block 2: Variable shape tensors with different shapes
|
|
tensor_data2 = np.array(
|
|
[
|
|
np.ones((1, 4), dtype=np.float32),
|
|
np.zeros((2, 1), dtype=np.float32),
|
|
],
|
|
dtype=object,
|
|
)
|
|
|
|
return _create_tensor_blocks(tensor_data1, tensor_data2, "variable", "variable")
|
|
|
|
|
|
@pytest.fixture
|
|
def mixed_tensor_types_variable_shaped_expected():
|
|
"""Fixture for expected results from mixed variable-shaped tensor types."""
|
|
expected_schema = _create_tensor_schema(struct_name="tensor")
|
|
expected_tensors = [
|
|
np.ones((2, 2), dtype=np.float32),
|
|
np.zeros((3, 3), dtype=np.float32),
|
|
np.ones((1, 4), dtype=np.float32),
|
|
np.zeros((2, 1), dtype=np.float32),
|
|
]
|
|
|
|
return _create_expected_result(expected_schema, 4, tensor_values=expected_tensors)
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_with_mixed_tensor_types_blocks():
|
|
"""Fixture for struct blocks with mixed tensor types."""
|
|
# Block 1: Struct with fixed-shape tensor
|
|
tensor_data1 = np.ones((2, 2), dtype=np.float32)
|
|
|
|
# Block 2: Struct with variable-shaped tensor
|
|
tensor_data2 = np.array(
|
|
[
|
|
np.ones((3, 3), dtype=np.float32),
|
|
np.zeros((1, 4), dtype=np.float32),
|
|
],
|
|
dtype=object,
|
|
)
|
|
|
|
return _create_struct_tensor_blocks(tensor_data1, tensor_data2, "fixed", "variable")
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_with_mixed_tensor_types_expected():
|
|
"""Fixture for expected results from struct with mixed tensor types."""
|
|
expected_schema = _create_tensor_schema(struct_name="struct")
|
|
expected_struct_values = [
|
|
{"value": 1}, # First two from fixed-shape tensor struct
|
|
{"value": 2},
|
|
{"value": 3}, # Last two from variable-shaped tensor struct
|
|
{"value": 4},
|
|
]
|
|
|
|
return _create_expected_result(
|
|
expected_schema, 4, struct_values=expected_struct_values
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def nested_struct_with_mixed_tensor_types_blocks():
|
|
"""Fixture for nested struct blocks with mixed tensor types."""
|
|
# Block 1: Nested struct with fixed-shape tensors
|
|
tensor_data1 = np.ones((2, 2), dtype=np.float32)
|
|
tensor_array1 = _create_tensor_array(tensor_data1, "fixed")
|
|
inner_struct1 = pa.StructArray.from_arrays(
|
|
[tensor_array1, pa.array([10, 20], type=pa.int64())],
|
|
names=["inner_tensor", "inner_value"],
|
|
)
|
|
outer_tensor1 = _create_tensor_array(np.zeros((2, 1), dtype=np.float32), "fixed")
|
|
outer_struct1 = pa.StructArray.from_arrays(
|
|
[inner_struct1, outer_tensor1, pa.array([1, 2], type=pa.int64())],
|
|
names=["nested", "outer_tensor", "outer_value"],
|
|
)
|
|
t1 = pa.table({"id": [1, 2], "complex_struct": outer_struct1})
|
|
|
|
# Block 2: Nested struct with variable-shaped tensors
|
|
tensor_data2 = np.array(
|
|
[
|
|
np.ones((3, 3), dtype=np.float32),
|
|
np.zeros((1, 4), dtype=np.float32),
|
|
],
|
|
dtype=object,
|
|
)
|
|
tensor_array2 = _create_tensor_array(tensor_data2, "variable")
|
|
inner_struct2 = pa.StructArray.from_arrays(
|
|
[tensor_array2, pa.array([30, 40], type=pa.int64())],
|
|
names=["inner_tensor", "inner_value"],
|
|
)
|
|
outer_tensor2 = _create_tensor_array(
|
|
np.array(
|
|
[np.ones((2, 2), dtype=np.float32), np.zeros((1, 3), dtype=np.float32)],
|
|
dtype=object,
|
|
),
|
|
"variable",
|
|
)
|
|
outer_struct2 = pa.StructArray.from_arrays(
|
|
[inner_struct2, outer_tensor2, pa.array([3, 4], type=pa.int64())],
|
|
names=["nested", "outer_tensor", "outer_value"],
|
|
)
|
|
t2 = pa.table({"id": [3, 4], "complex_struct": outer_struct2})
|
|
|
|
return t1, t2
|
|
|
|
|
|
@pytest.fixture
|
|
def nested_struct_with_mixed_tensor_types_expected():
|
|
"""Fixture for expected results from nested struct with mixed tensor types."""
|
|
expected_schema = pa.schema(
|
|
[
|
|
("id", pa.int64()),
|
|
(
|
|
"complex_struct",
|
|
pa.struct(
|
|
[
|
|
(
|
|
"nested",
|
|
pa.struct(
|
|
[
|
|
(
|
|
"inner_tensor",
|
|
ArrowVariableShapedTensorType(pa.float32(), 2),
|
|
),
|
|
("inner_value", pa.int64()),
|
|
]
|
|
),
|
|
),
|
|
(
|
|
"outer_tensor",
|
|
ArrowVariableShapedTensorType(pa.float32(), 2),
|
|
),
|
|
("outer_value", pa.int64()),
|
|
]
|
|
),
|
|
),
|
|
]
|
|
)
|
|
expected_fields = [
|
|
"nested",
|
|
"outer_tensor",
|
|
"outer_value",
|
|
"inner_tensor",
|
|
"inner_value",
|
|
]
|
|
|
|
return _create_expected_result(expected_schema, 4, expected_fields=expected_fields)
|
|
|
|
|
|
@pytest.fixture
|
|
def multiple_tensor_fields_struct_blocks():
|
|
"""Fixture for struct blocks with multiple tensor fields."""
|
|
# Block 1: Struct with multiple fixed-shape tensors
|
|
tensor1_data = np.ones((2, 2), dtype=np.float32)
|
|
tensor1_array = _create_tensor_array(tensor1_data, "fixed")
|
|
tensor2_data = np.zeros((2, 3), dtype=np.int32)
|
|
tensor2_array = _create_tensor_array(tensor2_data, "fixed")
|
|
struct_array1 = pa.StructArray.from_arrays(
|
|
[tensor1_array, tensor2_array, pa.array([1, 2], type=pa.int64())],
|
|
names=["tensor1", "tensor2", "value"],
|
|
)
|
|
t1 = pa.table({"id": [1, 2], "multi_tensor_struct": struct_array1})
|
|
|
|
# Block 2: Struct with multiple variable-shaped tensors
|
|
tensor1_data2 = np.array(
|
|
[
|
|
np.ones((3, 3), dtype=np.float32),
|
|
np.zeros((1, 4), dtype=np.float32),
|
|
],
|
|
dtype=object,
|
|
)
|
|
tensor1_array2 = _create_tensor_array(tensor1_data2, "variable")
|
|
tensor2_data2 = np.array(
|
|
[
|
|
np.ones((2, 2), dtype=np.int32),
|
|
np.zeros((3, 1), dtype=np.int32),
|
|
],
|
|
dtype=object,
|
|
)
|
|
tensor2_array2 = _create_tensor_array(tensor2_data2, "variable")
|
|
struct_array2 = pa.StructArray.from_arrays(
|
|
[tensor1_array2, tensor2_array2, pa.array([3, 4], type=pa.int64())],
|
|
names=["tensor1", "tensor2", "value"],
|
|
)
|
|
t2 = pa.table({"id": [3, 4], "multi_tensor_struct": struct_array2})
|
|
|
|
return t1, t2
|
|
|
|
|
|
@pytest.fixture
|
|
def multiple_tensor_fields_struct_expected():
|
|
"""Fixture for expected results from struct with multiple tensor fields."""
|
|
expected_schema = pa.schema(
|
|
[
|
|
("id", pa.int64()),
|
|
(
|
|
"multi_tensor_struct",
|
|
pa.struct(
|
|
[
|
|
("tensor1", ArrowVariableShapedTensorType(pa.float32(), 2)),
|
|
("tensor2", ArrowVariableShapedTensorType(pa.int32(), 2)),
|
|
("value", pa.int64()),
|
|
]
|
|
),
|
|
),
|
|
]
|
|
)
|
|
expected_fields = ["tensor1", "tensor2", "value"]
|
|
|
|
return _create_expected_result(expected_schema, 4, expected_fields=expected_fields)
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_with_additional_fields_blocks():
|
|
"""Fixture for struct blocks where some have additional fields."""
|
|
# Block 1: Struct with tensor field and basic fields
|
|
tensor_data1 = np.ones((2, 2), dtype=np.float32)
|
|
|
|
# Block 2: Struct with tensor field and additional fields
|
|
tensor_data2 = np.array(
|
|
[
|
|
np.ones((3, 3), dtype=np.float32),
|
|
np.zeros((1, 4), dtype=np.float32),
|
|
],
|
|
dtype=object,
|
|
)
|
|
|
|
return _create_struct_tensor_blocks(
|
|
tensor_data1, tensor_data2, "fixed", "variable", extra_data2=["a", "b"]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_with_additional_fields_expected():
|
|
"""Fixture for expected results from struct with additional fields."""
|
|
expected_schema = _create_tensor_schema(struct_name="struct", include_extra=True)
|
|
expected_field_presence = {"tensor": True, "value": True, "extra": True}
|
|
expected_extra_values = [None, None, "a", "b"]
|
|
|
|
return _create_expected_result(
|
|
expected_schema,
|
|
4,
|
|
field_presence=expected_field_presence,
|
|
extra_values=expected_extra_values,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_with_null_tensor_values_blocks():
|
|
"""Fixture for struct blocks where some fields are missing and get filled with nulls."""
|
|
# Block 1: Struct with tensor and value fields
|
|
tensor_data1 = np.ones((2, 2), dtype=np.float32)
|
|
tensor_array1 = ArrowTensorArray.from_numpy(tensor_data1)
|
|
value_array1 = pa.array([1, 2], type=pa.int64())
|
|
struct_array1 = pa.StructArray.from_arrays(
|
|
[tensor_array1, value_array1], names=["tensor", "value"]
|
|
)
|
|
t1 = pa.table({"id": [1, 2], "struct": struct_array1})
|
|
|
|
# Block 2: Struct with only value field (missing tensor field)
|
|
value_array2 = pa.array([3], type=pa.int64())
|
|
struct_array2 = pa.StructArray.from_arrays([value_array2], names=["value"])
|
|
t2 = pa.table({"id": [3], "struct": struct_array2})
|
|
|
|
return t1, t2
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_with_null_tensor_values_expected():
|
|
"""Fixture for expected results from struct with null tensor values."""
|
|
expected_schema = pa.schema(
|
|
[
|
|
("id", pa.int64()),
|
|
(
|
|
"struct",
|
|
pa.struct(
|
|
[
|
|
("tensor", ArrowTensorTypeV2((2,), pa.float32())),
|
|
("value", pa.int64()),
|
|
]
|
|
),
|
|
),
|
|
]
|
|
)
|
|
expected_length = 3
|
|
expected_ids = [1, 2, 3]
|
|
|
|
# Expected value field values
|
|
expected_values = [1, 2, 3]
|
|
|
|
# Expected tensor field validity
|
|
expected_tensor_validity = [True, True, False]
|
|
|
|
return {
|
|
"schema": expected_schema,
|
|
"length": expected_length,
|
|
"ids": expected_ids,
|
|
"values": expected_values,
|
|
"tensor_validity": expected_tensor_validity,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def basic_concat_blocks():
|
|
"""Fixture for basic concat test data."""
|
|
t1 = pa.table({"a": [1, 2], "b": [5, 6]})
|
|
t2 = pa.table({"a": [3, 4], "b": [7, 8]})
|
|
return [t1, t2]
|
|
|
|
|
|
@pytest.fixture
|
|
def basic_concat_expected():
|
|
"""Fixture for basic concat expected results."""
|
|
return {
|
|
"length": 4,
|
|
"column_names": ["a", "b"],
|
|
"schema_types": [pa.int64(), pa.int64()],
|
|
"chunks": 2,
|
|
"content": {"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]},
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def null_promotion_blocks():
|
|
"""Fixture for null promotion test data."""
|
|
t1 = pa.table({"a": [None, None], "b": [5, 6]})
|
|
t2 = pa.table({"a": [3, 4], "b": [None, None]})
|
|
return [t1, t2]
|
|
|
|
|
|
@pytest.fixture
|
|
def null_promotion_expected():
|
|
"""Fixture for null promotion expected results."""
|
|
return {
|
|
"length": 4,
|
|
"column_names": ["a", "b"],
|
|
"schema_types": [pa.int64(), pa.int64()],
|
|
"chunks": 2,
|
|
"content": {"a": [None, None, 3, 4], "b": [5, 6, None, None]},
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_different_field_names_blocks():
|
|
"""Fixture for struct with different field names test data."""
|
|
struct_data1 = [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}]
|
|
struct_data2 = [{"x": 3, "z": "c"}]
|
|
|
|
struct_type1 = pa.struct([("x", pa.int32()), ("y", pa.string())])
|
|
struct_type2 = pa.struct([("x", pa.int32()), ("z", pa.string())])
|
|
|
|
additional_columns1 = {"a": [1, 2]}
|
|
additional_columns2 = {"a": [3]}
|
|
|
|
return _create_struct_blocks_with_columns(
|
|
struct_data1,
|
|
struct_data2,
|
|
struct_type1,
|
|
struct_type2,
|
|
additional_columns1,
|
|
additional_columns2,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_different_field_names_expected():
|
|
"""Fixture for struct with different field names expected results."""
|
|
field_names = ["x", "y", "z"]
|
|
field_types = [pa.int32(), pa.string(), pa.string()]
|
|
additional_fields = [("a", pa.int64())]
|
|
|
|
schema = _create_simple_struct_schema(field_names, field_types, additional_fields)
|
|
|
|
content = {
|
|
"a": [1, 2, 3],
|
|
"d": [
|
|
{"x": 1, "y": "a", "z": None},
|
|
{"x": 2, "y": "b", "z": None},
|
|
{"x": 3, "y": None, "z": "c"},
|
|
],
|
|
}
|
|
|
|
return _create_struct_expected_result(schema, 3, content)
|
|
|
|
|
|
@pytest.fixture
|
|
def nested_structs_blocks():
|
|
"""Fixture for nested structs test data."""
|
|
t1 = pa.table(
|
|
{
|
|
"a": [1],
|
|
"d": pa.array(
|
|
[
|
|
{
|
|
"x": {
|
|
"y": {"p": 1}, # Missing "q"
|
|
"z": {"m": 3}, # Missing "n"
|
|
},
|
|
"w": 5,
|
|
}
|
|
],
|
|
type=pa.struct(
|
|
[
|
|
(
|
|
"x",
|
|
pa.struct(
|
|
[
|
|
(
|
|
"y",
|
|
pa.struct([("p", pa.int32())]), # Only "p"
|
|
),
|
|
(
|
|
"z",
|
|
pa.struct([("m", pa.int32())]), # Only "m"
|
|
),
|
|
]
|
|
),
|
|
),
|
|
("w", pa.int32()),
|
|
]
|
|
),
|
|
),
|
|
}
|
|
)
|
|
t2 = pa.table(
|
|
{
|
|
"a": [2],
|
|
"d": pa.array(
|
|
[
|
|
{
|
|
"x": {
|
|
"y": {"q": 7}, # Missing "p"
|
|
"z": {"n": 9}, # Missing "m"
|
|
},
|
|
"w": 10,
|
|
}
|
|
],
|
|
type=pa.struct(
|
|
[
|
|
(
|
|
"x",
|
|
pa.struct(
|
|
[
|
|
(
|
|
"y",
|
|
pa.struct([("q", pa.int32())]), # Only "q"
|
|
),
|
|
(
|
|
"z",
|
|
pa.struct([("n", pa.int32())]), # Only "n"
|
|
),
|
|
]
|
|
),
|
|
),
|
|
("w", pa.int32()),
|
|
]
|
|
),
|
|
),
|
|
}
|
|
)
|
|
return [t1, t2]
|
|
|
|
|
|
@pytest.fixture
|
|
def nested_structs_expected():
|
|
"""Fixture for nested structs expected results."""
|
|
return {
|
|
"length": 2,
|
|
"schema": pa.schema(
|
|
[
|
|
("a", pa.int64()),
|
|
(
|
|
"d",
|
|
pa.struct(
|
|
[
|
|
(
|
|
"x",
|
|
pa.struct(
|
|
[
|
|
(
|
|
"y",
|
|
pa.struct(
|
|
[("p", pa.int32()), ("q", pa.int32())]
|
|
),
|
|
),
|
|
(
|
|
"z",
|
|
pa.struct(
|
|
[("m", pa.int32()), ("n", pa.int32())]
|
|
),
|
|
),
|
|
]
|
|
),
|
|
),
|
|
("w", pa.int32()),
|
|
]
|
|
),
|
|
),
|
|
]
|
|
),
|
|
"content": {
|
|
"a": [1, 2],
|
|
"d": [
|
|
{
|
|
"x": {
|
|
"y": {"p": 1, "q": None}, # Missing "q" filled with None
|
|
"z": {"m": 3, "n": None}, # Missing "n" filled with None
|
|
},
|
|
"w": 5,
|
|
},
|
|
{
|
|
"x": {
|
|
"y": {"p": None, "q": 7}, # Missing "p" filled with None
|
|
"z": {"m": None, "n": 9}, # Missing "m" filled with None
|
|
},
|
|
"w": 10,
|
|
},
|
|
],
|
|
},
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_null_values_blocks():
|
|
"""Fixture for struct with null values test data."""
|
|
struct_data1 = [{"x": 1, "y": "a"}, None] # Second row is null
|
|
struct_data2 = [None] # Entire struct is null
|
|
|
|
field_names = ["x", "y"]
|
|
field_types = [pa.int32(), pa.string()]
|
|
additional_columns1 = {"a": [1, 2]}
|
|
additional_columns2 = {"a": [3]}
|
|
|
|
return _create_simple_struct_blocks(
|
|
struct_data1,
|
|
struct_data2,
|
|
field_names,
|
|
field_types,
|
|
additional_columns1,
|
|
additional_columns2,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_null_values_expected():
|
|
"""Fixture for struct with null values expected results."""
|
|
field_names = ["x", "y"]
|
|
field_types = [pa.int32(), pa.string()]
|
|
additional_fields = [("a", pa.int64())]
|
|
|
|
schema = _create_simple_struct_schema(field_names, field_types, additional_fields)
|
|
|
|
content = {
|
|
"a": [1, 2, 3],
|
|
"d": [
|
|
{"x": 1, "y": "a"},
|
|
None, # Entire struct is None, not {"x": None, "y": None}
|
|
None, # Entire struct is None, not {"x": None, "y": None}
|
|
],
|
|
}
|
|
|
|
return _create_struct_expected_result(schema, 3, content)
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_mismatched_lengths_blocks():
|
|
"""Fixture for struct with mismatched lengths test data."""
|
|
struct_data1 = [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}]
|
|
struct_data2 = [{"x": 3, "y": "c"}]
|
|
|
|
field_names = ["x", "y"]
|
|
field_types = [pa.int32(), pa.string()]
|
|
additional_columns1 = {"a": [1, 2]}
|
|
additional_columns2 = {"a": [3]}
|
|
|
|
return _create_simple_struct_blocks(
|
|
struct_data1,
|
|
struct_data2,
|
|
field_names,
|
|
field_types,
|
|
additional_columns1,
|
|
additional_columns2,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_mismatched_lengths_expected():
|
|
"""Fixture for struct with mismatched lengths expected results."""
|
|
field_names = ["x", "y"]
|
|
field_types = [pa.int32(), pa.string()]
|
|
additional_fields = [("a", pa.int64())]
|
|
|
|
schema = _create_simple_struct_schema(field_names, field_types, additional_fields)
|
|
|
|
content = {
|
|
"a": [1, 2, 3],
|
|
"d": [
|
|
{"x": 1, "y": "a"},
|
|
{"x": 2, "y": "b"},
|
|
{"x": 3, "y": "c"},
|
|
],
|
|
}
|
|
|
|
return _create_struct_expected_result(schema, 3, content)
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_empty_arrays_blocks():
|
|
"""Fixture for struct with empty arrays test data."""
|
|
struct_data1 = [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}]
|
|
|
|
# Define the second table with null struct value (empty arrays for fields)
|
|
x_array = pa.array([None], type=pa.int32())
|
|
y_array = pa.array([None], type=pa.string())
|
|
|
|
# Create a struct array from null field arrays
|
|
null_struct_array = pa.StructArray.from_arrays(
|
|
[x_array, y_array],
|
|
["x", "y"],
|
|
mask=pa.array([True]),
|
|
)
|
|
|
|
t1 = pa.table(
|
|
{
|
|
"a": [1, 2],
|
|
"d": pa.array(
|
|
struct_data1, type=pa.struct([("x", pa.int32()), ("y", pa.string())])
|
|
),
|
|
}
|
|
)
|
|
|
|
t2 = pa.table({"a": [3], "d": null_struct_array})
|
|
return [t1, t2]
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_empty_arrays_expected():
|
|
"""Fixture for struct with empty arrays expected results."""
|
|
field_names = ["x", "y"]
|
|
field_types = [pa.int32(), pa.string()]
|
|
additional_fields = [("a", pa.int64())]
|
|
|
|
schema = _create_simple_struct_schema(field_names, field_types, additional_fields)
|
|
|
|
content = {
|
|
"a": [1, 2, 3],
|
|
"d": [
|
|
{"x": 1, "y": "a"},
|
|
{"x": 2, "y": "b"},
|
|
None, # Entire struct is None, as PyArrow handles it
|
|
],
|
|
}
|
|
|
|
return _create_struct_expected_result(schema, 3, content)
|
|
|
|
|
|
@pytest.fixture
|
|
def unify_schemas_basic_schemas():
|
|
"""Fixture for basic unify schemas test data."""
|
|
tensor_arr_1 = pa.schema([("tensor_arr", ArrowTensorType((3, 5), pa.int32()))])
|
|
tensor_arr_2 = pa.schema([("tensor_arr", ArrowTensorType((2, 1), pa.int32()))])
|
|
tensor_arr_3 = pa.schema([("tensor_arr", ArrowTensorType((3, 5), pa.int32()))])
|
|
var_tensor_arr = pa.schema(
|
|
[
|
|
("tensor_arr", ArrowVariableShapedTensorType(pa.int32(), 2)),
|
|
]
|
|
)
|
|
var_tensor_arr_1d = pa.schema(
|
|
[
|
|
("tensor_arr", ArrowVariableShapedTensorType(pa.int32(), 1)),
|
|
]
|
|
)
|
|
var_tensor_arr_3d = pa.schema(
|
|
[
|
|
("tensor_arr", ArrowVariableShapedTensorType(pa.int32(), 3)),
|
|
]
|
|
)
|
|
return {
|
|
"tensor_arr_1": tensor_arr_1,
|
|
"tensor_arr_2": tensor_arr_2,
|
|
"tensor_arr_3": tensor_arr_3,
|
|
"var_tensor_arr": var_tensor_arr,
|
|
"var_tensor_arr_1d": var_tensor_arr_1d,
|
|
"var_tensor_arr_3d": var_tensor_arr_3d,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def unify_schemas_multicol_schemas():
|
|
"""Fixture for multi-column unify schemas test data."""
|
|
multicol_schema_1 = pa.schema(
|
|
[
|
|
("col_int", pa.int32()),
|
|
("col_fixed_tensor", ArrowTensorType((4, 2), pa.int32())),
|
|
("col_var_tensor", ArrowVariableShapedTensorType(pa.int16(), 5)),
|
|
]
|
|
)
|
|
multicol_schema_2 = pa.schema(
|
|
[
|
|
("col_int", pa.int32()),
|
|
("col_fixed_tensor", ArrowTensorType((4, 2), pa.int32())),
|
|
("col_var_tensor", ArrowTensorType((9, 4, 1, 0, 5), pa.int16())),
|
|
]
|
|
)
|
|
multicol_schema_3 = pa.schema(
|
|
[
|
|
("col_int", pa.int32()),
|
|
("col_fixed_tensor", ArrowVariableShapedTensorType(pa.int32(), 3)),
|
|
("col_var_tensor", ArrowVariableShapedTensorType(pa.int16(), 5)),
|
|
]
|
|
)
|
|
return {
|
|
"multicol_schema_1": multicol_schema_1,
|
|
"multicol_schema_2": multicol_schema_2,
|
|
"multicol_schema_3": multicol_schema_3,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def object_concat_blocks():
|
|
"""Fixture for object concat test data."""
|
|
obj = types.SimpleNamespace(a=1, b="test")
|
|
t1 = pa.table({"a": [3, 4], "b": [7, 8]})
|
|
t2 = pa.table({"a": ArrowPythonObjectArray.from_objects([obj, obj]), "b": [0, 1]})
|
|
return [t1, t2]
|
|
|
|
|
|
@pytest.fixture
|
|
def object_concat_expected():
|
|
"""Fixture for object concat expected results."""
|
|
obj = types.SimpleNamespace(a=1, b="test")
|
|
return {
|
|
"length": 4,
|
|
"a_type": ArrowPythonObjectType,
|
|
"b_type": pa.types.is_integer,
|
|
"content": {"a": [3, 4, obj, obj], "b": [7, 8, 0, 1]},
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_variable_shaped_tensor_blocks():
|
|
"""Fixture for struct with variable shaped tensor test data."""
|
|
# Create variable-shaped tensor data for the first table
|
|
tensor_data1 = np.array(
|
|
[
|
|
np.ones((2, 2), dtype=np.float32),
|
|
np.zeros((3, 3), dtype=np.float32),
|
|
],
|
|
dtype=object,
|
|
)
|
|
tensor_array1 = ArrowVariableShapedTensorArray.from_numpy(tensor_data1)
|
|
|
|
# Create struct data with tensor field for the first table
|
|
metadata_array1 = pa.array(["row1", "row2"])
|
|
struct_array1 = pa.StructArray.from_arrays(
|
|
[metadata_array1, tensor_array1], names=["metadata", "tensor"]
|
|
)
|
|
|
|
t1 = pa.table({"id": [1, 2], "struct_with_tensor": struct_array1})
|
|
|
|
# Create variable-shaped tensor data for the second table
|
|
tensor_data2 = np.array(
|
|
[
|
|
np.ones((1, 4), dtype=np.float32),
|
|
np.zeros((2, 1), dtype=np.float32),
|
|
],
|
|
dtype=object,
|
|
)
|
|
tensor_array2 = ArrowVariableShapedTensorArray.from_numpy(tensor_data2)
|
|
|
|
# Create struct data with tensor field for the second table
|
|
metadata_array2 = pa.array(["row3", "row4"])
|
|
struct_array2 = pa.StructArray.from_arrays(
|
|
[metadata_array2, tensor_array2], names=["metadata", "tensor"]
|
|
)
|
|
|
|
t2 = pa.table({"id": [3, 4], "struct_with_tensor": struct_array2})
|
|
return [t1, t2]
|
|
|
|
|
|
@pytest.fixture
|
|
def struct_variable_shaped_tensor_expected():
|
|
"""Fixture for struct with variable shaped tensor expected results."""
|
|
return {
|
|
"length": 4,
|
|
"schema": pa.schema(
|
|
[
|
|
("id", pa.int64()),
|
|
(
|
|
"struct_with_tensor",
|
|
pa.struct(
|
|
[
|
|
("metadata", pa.string()),
|
|
("tensor", ArrowVariableShapedTensorType(pa.float32(), 2)),
|
|
]
|
|
),
|
|
),
|
|
]
|
|
),
|
|
"content": {"id": [1, 2, 3, 4]},
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def unify_schemas_object_types_schemas():
|
|
"""Fixture for object types unify schemas test data."""
|
|
from ray.data._internal.object_extensions.arrow import ArrowPythonObjectType
|
|
|
|
schema1 = pa.schema([("obj_col", ArrowPythonObjectType())])
|
|
schema2 = pa.schema([("obj_col", pa.int32())])
|
|
schema3 = pa.schema([("obj_col", pa.float64())])
|
|
expected = pa.schema([("obj_col", ArrowPythonObjectType())])
|
|
|
|
return {
|
|
"object_schema": schema1,
|
|
"int_schema": schema2,
|
|
"float_schema": schema3,
|
|
"expected": expected,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def unify_schemas_incompatible_tensor_schemas():
|
|
"""Fixture for incompatible tensor dtypes unify schemas test data."""
|
|
schema1 = pa.schema([("tensor", ArrowTensorType((2, 2), pa.int32()))])
|
|
schema2 = pa.schema([("tensor", ArrowTensorType((2, 2), pa.float32()))])
|
|
return [schema1, schema2]
|
|
|
|
|
|
@pytest.fixture
|
|
def unify_schemas_objects_and_tensors_schemas():
|
|
"""Fixture for objects and tensors unify schemas test data."""
|
|
from ray.data._internal.object_extensions.arrow import ArrowPythonObjectType
|
|
|
|
schema1 = pa.schema([("col", ArrowPythonObjectType())])
|
|
schema2 = pa.schema([("col", ArrowTensorType((2, 2), pa.int32()))])
|
|
return [schema1, schema2]
|
|
|
|
|
|
@pytest.fixture
|
|
def unify_schemas_missing_tensor_fields_schemas():
|
|
"""Fixture for missing tensor fields unify schemas test data."""
|
|
schema1 = pa.schema(
|
|
[
|
|
(
|
|
"struct",
|
|
pa.struct(
|
|
[
|
|
("tensor", ArrowTensorType((2, 2), pa.int32())),
|
|
("value", pa.int64()),
|
|
]
|
|
),
|
|
)
|
|
]
|
|
)
|
|
schema2 = pa.schema(
|
|
[("struct", pa.struct([("value", pa.int64())]))] # Missing tensor field
|
|
)
|
|
expected = pa.schema(
|
|
[
|
|
(
|
|
"struct",
|
|
pa.struct(
|
|
[
|
|
("tensor", ArrowTensorType((2, 2), pa.int32())),
|
|
("value", pa.int64()),
|
|
]
|
|
),
|
|
)
|
|
]
|
|
)
|
|
return {"with_tensor": schema1, "without_tensor": schema2, "expected": expected}
|
|
|
|
|
|
@pytest.fixture
|
|
def unify_schemas_nested_struct_tensors_schemas():
|
|
"""Fixture for nested struct tensors unify schemas test data."""
|
|
schema1 = pa.schema(
|
|
[
|
|
(
|
|
"outer",
|
|
pa.struct(
|
|
[
|
|
(
|
|
"inner",
|
|
pa.struct(
|
|
[
|
|
("tensor", ArrowTensorType((3, 3), pa.float32())),
|
|
("data", pa.string()),
|
|
]
|
|
),
|
|
),
|
|
("id", pa.int64()),
|
|
]
|
|
),
|
|
)
|
|
]
|
|
)
|
|
schema2 = pa.schema(
|
|
[
|
|
(
|
|
"outer",
|
|
pa.struct(
|
|
[
|
|
(
|
|
"inner",
|
|
pa.struct([("data", pa.string())]), # Missing tensor field
|
|
),
|
|
("id", pa.int64()),
|
|
]
|
|
),
|
|
)
|
|
]
|
|
)
|
|
expected = pa.schema(
|
|
[
|
|
(
|
|
"outer",
|
|
pa.struct(
|
|
[
|
|
(
|
|
"inner",
|
|
pa.struct(
|
|
[
|
|
(
|
|
"tensor",
|
|
ArrowTensorType((3, 3), pa.float32()),
|
|
),
|
|
("data", pa.string()),
|
|
]
|
|
),
|
|
),
|
|
("id", pa.int64()),
|
|
]
|
|
),
|
|
)
|
|
]
|
|
)
|
|
return {"with_tensor": schema1, "without_tensor": schema2, "expected": expected}
|
|
|
|
|
|
@pytest.fixture
|
|
def object_with_tensor_fails_blocks():
|
|
"""Blocks that should fail when concatenating objects with tensors."""
|
|
obj = types.SimpleNamespace(a=1, b="test")
|
|
t1 = pa.table({"a": ArrowPythonObjectArray.from_objects([obj, obj])})
|
|
# Create tensor array with proper extension type
|
|
tensor_array = ArrowTensorArray.from_numpy(np.array([[1, 2], [3, 4]]))
|
|
t2 = pa.table({"a": tensor_array})
|
|
return [t1, t2]
|
|
|
|
|
|
@pytest.fixture
|
|
def simple_concat_data():
|
|
"""Test data for simple concat operations."""
|
|
return {"empty": [], "single_block": pa.table({"a": [1, 2]})}
|
|
|
|
|
|
# Helper function for creating tensor arrays
|
|
def _create_tensor_array(data, tensor_type="fixed"):
|
|
"""Helper function to create tensor arrays with consistent patterns."""
|
|
if tensor_type == "fixed":
|
|
return ArrowTensorArray.from_numpy(data)
|
|
elif tensor_type == "variable":
|
|
return ArrowVariableShapedTensorArray.from_numpy(data)
|
|
else:
|
|
raise ValueError(f"Unknown tensor type: {tensor_type}")
|
|
|
|
|
|
# Helper function for creating expected results
|
|
def _create_expected_result(schema, length, **kwargs):
|
|
"""Helper function to create expected result dictionaries."""
|
|
result = {"schema": schema, "length": length}
|
|
result.update(kwargs)
|
|
return result
|
|
|
|
|
|
# Helper function for creating tensor blocks
|
|
def _create_tensor_blocks(
|
|
tensor_data1,
|
|
tensor_data2,
|
|
tensor_type1="fixed",
|
|
tensor_type2="variable",
|
|
id_data1=None,
|
|
id_data2=None,
|
|
column_name="tensor",
|
|
):
|
|
"""Helper function to create tensor blocks with consistent patterns."""
|
|
if id_data1 is None:
|
|
id_data1 = [1, 2]
|
|
if id_data2 is None:
|
|
id_data2 = [3, 4]
|
|
|
|
tensor_array1 = _create_tensor_array(tensor_data1, tensor_type1)
|
|
tensor_array2 = _create_tensor_array(tensor_data2, tensor_type2)
|
|
|
|
t1 = pa.table({"id": id_data1, column_name: tensor_array1})
|
|
t2 = pa.table({"id": id_data2, column_name: tensor_array2})
|
|
|
|
return t1, t2
|
|
|
|
|
|
# Helper function for creating struct blocks with tensors
|
|
def _create_struct_tensor_blocks(
|
|
tensor_data1,
|
|
tensor_data2,
|
|
tensor_type1="fixed",
|
|
tensor_type2="variable",
|
|
value_data1=None,
|
|
value_data2=None,
|
|
extra_data2=None,
|
|
struct_name="struct",
|
|
id_data1=None,
|
|
id_data2=None,
|
|
):
|
|
"""Helper function to create struct blocks with tensor fields."""
|
|
if value_data1 is None:
|
|
value_data1 = [1, 2]
|
|
if value_data2 is None:
|
|
value_data2 = [3, 4]
|
|
if id_data1 is None:
|
|
id_data1 = [1, 2]
|
|
if id_data2 is None:
|
|
id_data2 = [3, 4]
|
|
|
|
tensor_array1 = _create_tensor_array(tensor_data1, tensor_type1)
|
|
tensor_array2 = _create_tensor_array(tensor_data2, tensor_type2)
|
|
|
|
value_array1 = pa.array(value_data1, type=pa.int64())
|
|
value_array2 = pa.array(value_data2, type=pa.int64())
|
|
|
|
if extra_data2 is not None:
|
|
extra_array2 = pa.array(extra_data2, type=pa.string())
|
|
struct_array1 = pa.StructArray.from_arrays(
|
|
[tensor_array1, value_array1], names=["tensor", "value"]
|
|
)
|
|
struct_array2 = pa.StructArray.from_arrays(
|
|
[tensor_array2, value_array2, extra_array2],
|
|
names=["tensor", "value", "extra"],
|
|
)
|
|
else:
|
|
struct_array1 = pa.StructArray.from_arrays(
|
|
[tensor_array1, value_array1], names=["tensor", "value"]
|
|
)
|
|
struct_array2 = pa.StructArray.from_arrays(
|
|
[tensor_array2, value_array2], names=["tensor", "value"]
|
|
)
|
|
|
|
t1 = pa.table({"id": id_data1, struct_name: struct_array1})
|
|
t2 = pa.table({"id": id_data2, struct_name: struct_array2})
|
|
|
|
return t1, t2
|
|
|
|
|
|
# Helper function for creating expected tensor schemas
|
|
def _create_tensor_schema(
|
|
tensor_type=ArrowVariableShapedTensorType,
|
|
dtype=pa.float32(),
|
|
ndim=2,
|
|
include_id=True,
|
|
struct_name="struct",
|
|
include_extra=False,
|
|
):
|
|
"""Helper function to create expected tensor schemas."""
|
|
fields = []
|
|
if include_id:
|
|
fields.append(("id", pa.int64()))
|
|
|
|
if struct_name == "struct":
|
|
struct_fields = [
|
|
("tensor", tensor_type(dtype, ndim)),
|
|
("value", pa.int64()),
|
|
]
|
|
if include_extra:
|
|
struct_fields.append(("extra", pa.string()))
|
|
fields.append((struct_name, pa.struct(struct_fields)))
|
|
else:
|
|
fields.append(("tensor", tensor_type(dtype, ndim)))
|
|
|
|
return pa.schema(fields)
|
|
|
|
|
|
# Helper function for creating basic struct blocks
|
|
def _create_basic_struct_blocks(
|
|
struct_data1,
|
|
struct_data2,
|
|
column_name="struct",
|
|
id_data1=None,
|
|
id_data2=None,
|
|
other_columns=None,
|
|
):
|
|
"""Helper function to create basic struct blocks."""
|
|
struct_array1 = pa.array(struct_data1)
|
|
struct_array2 = pa.array(struct_data2)
|
|
|
|
t1_data = {column_name: struct_array1}
|
|
t2_data = {column_name: struct_array2}
|
|
|
|
# Only add id columns if they are provided
|
|
if id_data1 is not None:
|
|
t1_data["id"] = id_data1
|
|
if id_data2 is not None:
|
|
t2_data["id"] = id_data2
|
|
|
|
if other_columns:
|
|
t1_data.update(other_columns.get("t1", {}))
|
|
t2_data.update(other_columns.get("t2", {}))
|
|
|
|
t1 = pa.table(t1_data)
|
|
t2 = pa.table(t2_data)
|
|
|
|
return t1, t2
|
|
|
|
|
|
# Helper function for creating struct schemas
|
|
def _create_struct_schema(struct_fields, include_id=True, other_fields=None):
|
|
"""Helper function to create struct schemas."""
|
|
fields = []
|
|
if include_id:
|
|
fields.append(("id", pa.int64()))
|
|
|
|
fields.append(("struct", pa.struct(struct_fields)))
|
|
|
|
if other_fields:
|
|
fields.extend(other_fields)
|
|
|
|
return pa.schema(fields)
|
|
|
|
|
|
# Helper function for creating struct blocks with additional columns
|
|
def _create_struct_blocks_with_columns(
|
|
struct_data1,
|
|
struct_data2,
|
|
struct_type1,
|
|
struct_type2,
|
|
additional_columns1=None,
|
|
additional_columns2=None,
|
|
struct_column="d",
|
|
):
|
|
"""Helper function to create struct blocks with additional columns."""
|
|
t1_data = {}
|
|
t2_data = {}
|
|
|
|
# Add additional columns first to maintain expected order
|
|
if additional_columns1:
|
|
t1_data.update(additional_columns1)
|
|
if additional_columns2:
|
|
t2_data.update(additional_columns2)
|
|
|
|
# Add struct column
|
|
t1_data[struct_column] = pa.array(struct_data1, type=struct_type1)
|
|
t2_data[struct_column] = pa.array(struct_data2, type=struct_type2)
|
|
|
|
t1 = pa.table(t1_data)
|
|
t2 = pa.table(t2_data)
|
|
|
|
return t1, t2
|
|
|
|
|
|
# Helper function for creating expected results for struct tests
|
|
def _create_struct_expected_result(schema, length, content):
|
|
"""Helper function to create expected results for struct tests."""
|
|
return {
|
|
"length": length,
|
|
"schema": schema,
|
|
"content": content,
|
|
}
|
|
|
|
|
|
# Helper function for creating struct blocks with simple field patterns
|
|
def _create_simple_struct_blocks(
|
|
struct_data1,
|
|
struct_data2,
|
|
field_names,
|
|
field_types,
|
|
additional_columns1=None,
|
|
additional_columns2=None,
|
|
struct_column="d",
|
|
):
|
|
"""Helper function to create struct blocks with simple field patterns."""
|
|
struct_type = pa.struct(list(zip(field_names, field_types)))
|
|
|
|
return _create_struct_blocks_with_columns(
|
|
struct_data1,
|
|
struct_data2,
|
|
struct_type,
|
|
struct_type,
|
|
additional_columns1,
|
|
additional_columns2,
|
|
struct_column,
|
|
)
|
|
|
|
|
|
# Helper function for creating simple struct schemas
|
|
def _create_simple_struct_schema(field_names, field_types, additional_fields=None):
|
|
"""Helper function to create simple struct schemas."""
|
|
struct_fields = list(zip(field_names, field_types))
|
|
|
|
fields = []
|
|
if additional_fields:
|
|
fields.extend(additional_fields)
|
|
fields.append(("d", pa.struct(struct_fields)))
|
|
|
|
return pa.schema(fields)
|
|
|
|
|
|
@pytest.fixture
|
|
def unify_schemas_edge_cases_data():
|
|
"""Test data for unify schemas edge cases."""
|
|
return {
|
|
"empty_schemas": [],
|
|
"single_schema": pa.schema([("col", pa.int32())]),
|
|
"no_common_columns": {
|
|
"schema1": pa.schema([("col1", pa.int32())]),
|
|
"schema2": pa.schema([("col2", pa.string())]),
|
|
"expected": pa.schema([("col1", pa.int32()), ("col2", pa.string())]),
|
|
},
|
|
"all_null_schemas": {
|
|
"schema1": pa.schema([("col", pa.null())]),
|
|
"schema2": pa.schema([("col", pa.null())]),
|
|
},
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def unify_schemas_mixed_tensor_data():
|
|
"""Test data for mixed tensor types in unify schemas."""
|
|
return {
|
|
"fixed_shape": pa.schema([("tensor", ArrowTensorType((2, 2), pa.int32()))]),
|
|
"variable_shaped": pa.schema(
|
|
[("tensor", ArrowVariableShapedTensorType(pa.int32(), 2))]
|
|
),
|
|
"different_shape": pa.schema([("tensor", ArrowTensorType((3, 3), pa.int32()))]),
|
|
"expected_variable": pa.schema(
|
|
[("tensor", ArrowVariableShapedTensorType(pa.int32(), 2))]
|
|
),
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def unify_schemas_type_promotion_data():
|
|
"""Test data for type promotion scenarios."""
|
|
return {
|
|
"non_null": pa.schema([pa.field("A", pa.int32())]),
|
|
"nullable": pa.schema([pa.field("A", pa.int32(), nullable=True)]),
|
|
"int64": pa.schema([pa.field("A", pa.int64())]),
|
|
"float64": pa.schema([pa.field("A", pa.float64())]),
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def block_select_data():
|
|
"""Test data for block select operations."""
|
|
df = pd.DataFrame({"one": [10, 11, 12], "two": [11, 12, 13], "three": [14, 15, 16]})
|
|
table = pa.Table.from_pandas(df)
|
|
return {
|
|
"table": table,
|
|
"df": df,
|
|
"single_column": {
|
|
"columns": ["two"],
|
|
"expected_schema": pa.schema([("two", pa.int64())]),
|
|
},
|
|
"multiple_columns": {
|
|
"columns": ["two", "one"],
|
|
"expected_schema": pa.schema([("two", pa.int64()), ("one", pa.int64())]),
|
|
},
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def block_slice_data():
|
|
"""Test data for block slice operations."""
|
|
n = 20
|
|
df = pd.DataFrame(
|
|
{"one": list(range(n)), "two": ["a"] * n, "three": [np.nan] + [1.5] * (n - 1)}
|
|
)
|
|
table = pa.Table.from_pandas(df)
|
|
empty_df = pd.DataFrame({"one": []})
|
|
empty_table = pa.Table.from_pandas(empty_df)
|
|
return {
|
|
"normal": {"table": table, "df": df, "slice_params": {"a": 5, "b": 10}},
|
|
"empty": {"table": empty_table, "slice_params": {"a": 0, "b": 0}},
|
|
}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|