wehub-resource-sync
141 lines
4.4 KiB
Python
141 lines
4.4 KiB
Python
import pyarrow as pa
|
|
import pytest
|
|
|
|
from ray.data._internal.arrow_ops.transform_pyarrow import (
|
|
unify_schemas,
|
|
)
|
|
from ray.data.extensions import (
|
|
ArrowPythonObjectType,
|
|
ArrowTensorType,
|
|
ArrowVariableShapedTensorType,
|
|
)
|
|
|
|
|
|
# Schema factory functions - just return schemas
|
|
def _create_simple_schema(num_columns):
|
|
return pa.schema([(f"col_{i}", pa.int64()) for i in range(num_columns)])
|
|
|
|
|
|
def _create_tensor_fixed_schema(num_columns):
|
|
return pa.schema(
|
|
[
|
|
(f"tensor_{i}", ArrowTensorType((2, 2), pa.float32()))
|
|
for i in range(num_columns)
|
|
]
|
|
)
|
|
|
|
|
|
def _create_tensor_variable_schema(num_columns):
|
|
return pa.schema(
|
|
[
|
|
(f"tensor_{i}", ArrowVariableShapedTensorType(pa.float32(), 2))
|
|
for i in range(num_columns)
|
|
]
|
|
)
|
|
|
|
|
|
def _create_object_schema(num_columns):
|
|
return pa.schema(
|
|
[(f"obj_{i}", ArrowPythonObjectType()) for i in range(num_columns)]
|
|
)
|
|
|
|
|
|
def _create_nested_struct_schema(num_columns):
|
|
fields = []
|
|
for i in range(num_columns):
|
|
inner_struct = pa.struct(
|
|
[("x", pa.int32()), ("y", pa.string()), ("z", pa.float64())]
|
|
)
|
|
fields.append((f"struct_{i}", inner_struct))
|
|
return pa.schema(fields)
|
|
|
|
|
|
def _create_deep_nested_schema(num_columns):
|
|
fields = []
|
|
for i in range(num_columns):
|
|
level4 = pa.struct([("data", pa.int32()), ("meta", pa.string())])
|
|
level3 = pa.struct([("level4", level4), ("id3", pa.int64())])
|
|
level2 = pa.struct([("level3", level3), ("id2", pa.int64())])
|
|
level1 = pa.struct([("level2", level2), ("id1", pa.int64())])
|
|
fields.append((f"deep_{i}", level1))
|
|
return pa.schema(fields)
|
|
|
|
|
|
def _create_mixed_complex_schema(num_columns):
|
|
fields = []
|
|
for i in range(num_columns):
|
|
field_type = i % 5
|
|
if field_type == 0:
|
|
fields.append((f"col_{i}", pa.int64()))
|
|
elif field_type == 1:
|
|
fields.append((f"col_{i}", ArrowTensorType((3, 3), pa.int32())))
|
|
elif field_type == 2:
|
|
fields.append((f"col_{i}", ArrowPythonObjectType()))
|
|
elif field_type == 3:
|
|
inner_struct = pa.struct([("a", pa.int32()), ("b", pa.string())])
|
|
fields.append((f"col_{i}", inner_struct))
|
|
else:
|
|
fields.append((f"col_{i}", pa.list_(pa.float64())))
|
|
return pa.schema(fields)
|
|
|
|
|
|
@pytest.mark.parametrize("num_schemas", [10, 100])
|
|
@pytest.mark.parametrize("num_columns", [10, 100, 1000, 5000])
|
|
@pytest.mark.parametrize(
|
|
"schema_factory,expected_time_per_schema_per_column",
|
|
[
|
|
(_create_simple_schema, 0.00001),
|
|
(_create_tensor_fixed_schema, 0.00005),
|
|
(_create_tensor_variable_schema, 0.00005),
|
|
(_create_object_schema, 0.00005),
|
|
(_create_nested_struct_schema, 0.0001),
|
|
(_create_deep_nested_schema, 0.0002),
|
|
(_create_mixed_complex_schema, 0.0002),
|
|
],
|
|
)
|
|
def test_unify_schemas_equivalent_performance(
|
|
num_schemas, num_columns, schema_factory, expected_time_per_schema_per_column
|
|
):
|
|
"""Stress test for unify_schemas when ALL schemas are equivalent (identical).
|
|
|
|
This tests the fast path where all schemas are the same and should be optimized
|
|
to return quickly without expensive comparisons.
|
|
"""
|
|
import time
|
|
|
|
# Create the base schema
|
|
base_schema = schema_factory(num_columns)
|
|
|
|
# Create list of identical schemas
|
|
schemas = [base_schema] * num_schemas
|
|
|
|
# Time the unification
|
|
start_time = time.time()
|
|
unified = unify_schemas(schemas)
|
|
elapsed_time = time.time() - start_time
|
|
|
|
# Verify the result is correct (should be identical to base schema)
|
|
assert unified == base_schema
|
|
|
|
# Performance assertions with scaling based on complexity
|
|
scale_factor = num_schemas * num_columns
|
|
max_allowed_time = expected_time_per_schema_per_column * scale_factor
|
|
buffer_factor = 2
|
|
assert elapsed_time < buffer_factor * max_allowed_time, (
|
|
f"unify_schemas took {elapsed_time:.4f}s for {num_schemas} identical "
|
|
f"{schema_factory.__name__} schemas with {num_columns} columns, "
|
|
f"should be < {max_allowed_time:.4f}s"
|
|
)
|
|
|
|
# Print timing info for large cases
|
|
if num_schemas >= 1000 or num_columns >= 100:
|
|
print(
|
|
f"\n{schema_factory.__name__}: {num_schemas} schemas x {num_columns} cols = {elapsed_time:.4f}s"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|