wehub-resource-sync
1484 lines
56 KiB
Python
1484 lines
56 KiB
Python
from dataclasses import dataclass
|
|
from typing import Dict, List, Set
|
|
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pyarrow.compute as pc
|
|
import pyarrow.parquet as pq
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data._internal.logical.interfaces import LogicalPlan
|
|
from ray.data._internal.logical.operators.input_data_operator import InputData
|
|
from ray.data._internal.logical.operators.map_operator import Project
|
|
from ray.data._internal.logical.operators.one_to_one_operator import AbstractOneToOne
|
|
from ray.data._internal.logical.optimizers import LogicalOptimizer
|
|
from ray.data._internal.logical.rules import (
|
|
ProjectionPushdown,
|
|
)
|
|
from ray.data._internal.util import rows_same
|
|
from ray.data.aggregate import Mean, Sum
|
|
from ray.data.context import DataContext
|
|
from ray.data.expressions import DataType, StarExpr, col, star, udf
|
|
|
|
|
|
@dataclass
|
|
class FusionTestCase:
|
|
"""Test case for projection fusion scenarios."""
|
|
|
|
name: str
|
|
expressions_list: List[Dict[str, str]] # List of {name: expression_desc}
|
|
expected_levels: int
|
|
expected_level_contents: List[Set[str]] # Expected expressions in each level
|
|
description: str
|
|
|
|
|
|
@dataclass
|
|
class DependencyTestCase:
|
|
"""Test case for dependency analysis."""
|
|
|
|
name: str
|
|
expression_desc: str
|
|
expected_refs: Set[str]
|
|
description: str
|
|
|
|
|
|
class TestProjectionFusion:
|
|
"""Test topological sorting in projection pushdown fusion."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def setup(self):
|
|
"""Set up test fixtures."""
|
|
self.context = DataContext.get_current()
|
|
|
|
# Create UDFs for testing
|
|
@udf(return_dtype=DataType.int64())
|
|
def multiply_by_two(x: pa.Array) -> pa.Array:
|
|
return pc.multiply(x, 2)
|
|
|
|
@udf(return_dtype=DataType.int64())
|
|
def add_one(x: pa.Array) -> pa.Array:
|
|
return pc.add(x, 1)
|
|
|
|
@udf(return_dtype=DataType.float64())
|
|
def divide_by_three(x: pa.Array) -> pa.Array:
|
|
# Convert to float to ensure floating point division
|
|
return pc.divide(pc.cast(x, pa.float64()), 3.0)
|
|
|
|
self.udfs = {
|
|
"multiply_by_two": multiply_by_two,
|
|
"add_one": add_one,
|
|
"divide_by_three": divide_by_three,
|
|
}
|
|
|
|
def _create_input_op(self):
|
|
"""Create a dummy input operator."""
|
|
return InputData(input_data=[])
|
|
|
|
def _parse_expression(self, expr_desc: str):
|
|
"""Parse expression description into actual expression object."""
|
|
# Enhanced parser for test expressions
|
|
expr_map = {
|
|
"col('id')": col("id"),
|
|
"col('id') + 10": col("id") + 10,
|
|
"col('id') * 2": col("id") * 2,
|
|
"col('id') - 5": col("id") - 5,
|
|
"col('id') + 1": col("id") + 1,
|
|
"col('id') - 1": col("id") - 1,
|
|
"col('id') - 3": col("id") - 3,
|
|
"col('step1') * 2": col("step1") * 2,
|
|
"col('step2') + 1": col("step2") + 1,
|
|
"col('a') + col('b')": col("a") + col("b"),
|
|
"col('c') + col('d')": col("c") + col("d"),
|
|
"col('e') * 3": col("e") * 3,
|
|
"col('a') + 1": col("a") + 1,
|
|
"multiply_by_two(col('id'))": self.udfs["multiply_by_two"](col("id")),
|
|
"multiply_by_two(col('id')) + col('plus_ten')": (
|
|
self.udfs["multiply_by_two"](col("id")) + col("plus_ten")
|
|
),
|
|
"col('times_three') > col('plus_ten')": (
|
|
col("times_three") > col("plus_ten")
|
|
),
|
|
"multiply_by_two(col('x'))": self.udfs["multiply_by_two"](col("x")),
|
|
"add_one(col('id'))": self.udfs["add_one"](col("id")),
|
|
"multiply_by_two(col('plus_one'))": self.udfs["multiply_by_two"](
|
|
col("plus_one")
|
|
),
|
|
"divide_by_three(col('times_two'))": self.udfs["divide_by_three"](
|
|
col("times_two")
|
|
),
|
|
}
|
|
|
|
if expr_desc in expr_map:
|
|
return expr_map[expr_desc]
|
|
else:
|
|
raise ValueError(f"Unknown expression: {expr_desc}")
|
|
|
|
def _create_project_chain(self, input_op, expressions_list: List[Dict[str, str]]):
|
|
"""Create a chain of Project operators from expression descriptions."""
|
|
current_op = input_op
|
|
|
|
for expr_dict in expressions_list:
|
|
# Convert dictionary to list of named expressions
|
|
exprs = []
|
|
for name, desc in expr_dict.items():
|
|
expr = self._parse_expression(desc)
|
|
named_expr = expr.alias(name)
|
|
exprs.append(named_expr)
|
|
|
|
current_op = Project(
|
|
exprs=[star()] + exprs,
|
|
input_dependencies=[current_op],
|
|
ray_remote_args={},
|
|
)
|
|
|
|
return current_op
|
|
|
|
def _extract_levels_from_plan(self, plan: LogicalPlan) -> List[Set[str]]:
|
|
"""Extract expression levels from optimized plan."""
|
|
current = plan.dag
|
|
levels = []
|
|
|
|
while isinstance(current, Project):
|
|
# Extract names, ignoring StarExpr (not a named column)
|
|
levels.append(
|
|
{expr.name for expr in current.exprs if not isinstance(expr, StarExpr)}
|
|
)
|
|
current = current.input_dependencies[0]
|
|
|
|
return list(reversed(levels)) # Return bottom-up order
|
|
|
|
def _count_project_operators(self, plan: LogicalPlan) -> int:
|
|
"""Count the number of Project operators in the plan."""
|
|
current = plan.dag
|
|
count = 0
|
|
|
|
while current:
|
|
if isinstance(current, Project):
|
|
count += 1
|
|
current = (
|
|
current.input_dependencies[0]
|
|
if isinstance(current, AbstractOneToOne)
|
|
else None
|
|
)
|
|
|
|
return count
|
|
|
|
def _describe_plan_structure(self, plan: LogicalPlan) -> str:
|
|
"""Generate a description of the plan structure."""
|
|
current = plan.dag
|
|
operators = []
|
|
|
|
while current:
|
|
if isinstance(current, Project):
|
|
expr_count = len(current.exprs) if current.exprs else 0
|
|
operators.append(f"Project({expr_count} exprs)")
|
|
else:
|
|
operators.append(current.__class__.__name__)
|
|
current = (
|
|
current.input_dependencies[0]
|
|
if isinstance(current, AbstractOneToOne)
|
|
else None
|
|
)
|
|
|
|
return " -> ".join(operators)
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
FusionTestCase(
|
|
name="no_dependencies",
|
|
expressions_list=[
|
|
{"doubled": "col('id') * 2", "plus_five": "col('id') + 10"},
|
|
{"minus_three": "col('id') - 3"},
|
|
],
|
|
expected_levels=1,
|
|
expected_level_contents=[{"doubled", "plus_five", "minus_three"}],
|
|
description="Independent expressions should fuse into single operator",
|
|
),
|
|
FusionTestCase(
|
|
name="simple_chain",
|
|
expressions_list=[
|
|
{"step1": "col('id') + 10"},
|
|
{"step2": "col('step1') * 2"},
|
|
{"step3": "col('step2') + 1"},
|
|
],
|
|
expected_levels=1,
|
|
expected_level_contents=[
|
|
{"step1", "step2", "step3"}
|
|
], # All in one level
|
|
description="All expressions fuse into single operator with OrderedDict preservation",
|
|
),
|
|
FusionTestCase(
|
|
name="mixed_udf_regular",
|
|
expressions_list=[
|
|
{"plus_ten": "col('id') + 10"},
|
|
{"times_three": "multiply_by_two(col('id'))"},
|
|
{"minus_five": "col('id') - 5"},
|
|
{
|
|
"udf_plus_regular": "multiply_by_two(col('id')) + col('plus_ten')"
|
|
},
|
|
{"comparison": "col('times_three') > col('plus_ten')"},
|
|
],
|
|
expected_levels=1,
|
|
expected_level_contents=[
|
|
{
|
|
"plus_ten",
|
|
"times_three",
|
|
"minus_five",
|
|
"udf_plus_regular",
|
|
"comparison",
|
|
}
|
|
],
|
|
description="All expressions fuse into single operator",
|
|
),
|
|
FusionTestCase(
|
|
name="complex_graph",
|
|
expressions_list=[
|
|
{"a": "col('id') + 1", "b": "col('id') * 2"},
|
|
{"c": "col('a') + col('b')"},
|
|
{"d": "col('id') - 1"},
|
|
{"e": "col('c') + col('d')"},
|
|
{"f": "col('e') * 3"},
|
|
],
|
|
expected_levels=1,
|
|
expected_level_contents=[{"a", "b", "c", "d", "e", "f"}],
|
|
description="All expressions fuse into single operator",
|
|
),
|
|
FusionTestCase(
|
|
name="udf_dependency_chain",
|
|
expressions_list=[
|
|
{"plus_one": "add_one(col('id'))"},
|
|
{"times_two": "multiply_by_two(col('plus_one'))"},
|
|
{"div_three": "divide_by_three(col('times_two'))"},
|
|
],
|
|
expected_levels=1, # Changed from 3 to 1
|
|
expected_level_contents=[{"plus_one", "times_two", "div_three"}],
|
|
description="All UDF expressions fuse into single operator with preserved order",
|
|
),
|
|
],
|
|
)
|
|
def test_fusion_scenarios(self, test_case: FusionTestCase):
|
|
"""Test various fusion scenarios with simplified single-operator fusion."""
|
|
input_op = self._create_input_op()
|
|
final_op = self._create_project_chain(input_op, test_case.expressions_list)
|
|
|
|
# Apply projection pushdown
|
|
plan = LogicalPlan(final_op, self.context)
|
|
rule = ProjectionPushdown()
|
|
optimized_plan = rule.apply(plan)
|
|
|
|
# Extract levels from optimized plan
|
|
actual_levels = self._extract_levels_from_plan(optimized_plan)
|
|
|
|
# Verify number of levels
|
|
assert len(actual_levels) == test_case.expected_levels, (
|
|
f"{test_case.name}: Expected {test_case.expected_levels} operators, "
|
|
f"got {len(actual_levels)}. Actual operators: {actual_levels}"
|
|
)
|
|
|
|
# Verify level contents (more flexible matching)
|
|
for i, expected_content in enumerate(test_case.expected_level_contents):
|
|
assert expected_content.issubset(actual_levels[i]), (
|
|
f"{test_case.name}: Operator {i} missing expressions. "
|
|
f"Expected {expected_content} to be subset of {actual_levels[i]}"
|
|
)
|
|
|
|
def test_pairwise_fusion_behavior(self, ray_start_regular_shared):
|
|
"""Test to understand how pairwise fusion works in practice."""
|
|
input_data = [{"id": i} for i in range(10)]
|
|
|
|
# Test with 2 operations (should fuse to 1)
|
|
ds2 = ray.data.from_items(input_data)
|
|
ds2 = ds2.with_column("col1", col("id") + 1)
|
|
ds2 = ds2.with_column("col2", col("id") * 2)
|
|
|
|
count2 = self._count_project_operators(ds2._logical_plan)
|
|
print(f"2 operations -> {count2} operators")
|
|
|
|
# Test with 3 operations
|
|
ds3 = ray.data.from_items(input_data)
|
|
ds3 = ds3.with_column("col1", col("id") + 1)
|
|
ds3 = ds3.with_column("col2", col("id") * 2)
|
|
ds3 = ds3.with_column("col3", col("id") - 1)
|
|
|
|
count3 = self._count_project_operators(ds3._logical_plan)
|
|
print(f"3 operations -> {count3} operators")
|
|
|
|
# Test with 4 operations
|
|
ds4 = ray.data.from_items(input_data)
|
|
ds4 = ds4.with_column("col1", col("id") + 1)
|
|
ds4 = ds4.with_column("col2", col("id") * 2)
|
|
ds4 = ds4.with_column("col3", col("id") - 1)
|
|
ds4 = ds4.with_column("col4", col("id") + 5)
|
|
|
|
count4 = self._count_project_operators(ds4._logical_plan)
|
|
print(f"4 operations -> {count4} operators")
|
|
|
|
# Verify that fusion is happening (fewer operators than original)
|
|
assert count2 <= 2, f"2 operations should result in ≤2 operators, got {count2}"
|
|
assert count3 <= 3, f"3 operations should result in ≤3 operators, got {count3}"
|
|
assert count4 <= 4, f"4 operations should result in ≤4 operators, got {count4}"
|
|
|
|
# Verify correctness
|
|
result2 = ds2.take(1)[0]
|
|
result3 = ds3.take(1)[0]
|
|
result4 = ds4.take(1)[0]
|
|
|
|
assert result2 == {"id": 0, "col1": 1, "col2": 0}
|
|
assert result3 == {"id": 0, "col1": 1, "col2": 0, "col3": -1}
|
|
assert result4 == {"id": 0, "col1": 1, "col2": 0, "col3": -1, "col4": 5}
|
|
|
|
def test_optimal_fusion_with_single_chain(self, ray_start_regular_shared):
|
|
"""Test fusion when all operations are added in a single chain (ideal case)."""
|
|
input_data = [{"id": i} for i in range(10)]
|
|
|
|
# Create a single Project operator with multiple expressions
|
|
# This simulates what would happen with perfect fusion
|
|
ds = ray.data.from_items(input_data)
|
|
|
|
# Apply multiple operations that should all be independent
|
|
expressions = {
|
|
"col1": col("id") + 1,
|
|
"col2": col("id") * 2,
|
|
"col3": col("id") - 1,
|
|
"col4": col("id") + 5,
|
|
"col5": col("id") * 3,
|
|
}
|
|
|
|
# Use map_batches to create a single operation that does everything
|
|
def apply_all_expressions(batch):
|
|
import pyarrow.compute as pc
|
|
|
|
result = batch.to_pydict()
|
|
result["col1"] = pc.add(batch["id"], 1)
|
|
result["col2"] = pc.multiply(batch["id"], 2)
|
|
result["col3"] = pc.subtract(batch["id"], 1)
|
|
result["col4"] = pc.add(batch["id"], 5)
|
|
result["col5"] = pc.multiply(batch["id"], 3)
|
|
return pa.table(result)
|
|
|
|
ds_optimal = ds.map_batches(apply_all_expressions, batch_format="pyarrow")
|
|
|
|
# Compare with the with_column approach
|
|
ds_with_column = ds
|
|
for col_name, expr in expressions.items():
|
|
ds_with_column = ds_with_column.with_column(col_name, expr)
|
|
|
|
# Convert both to pandas for reliable comparison
|
|
result_optimal_df = ds_optimal.to_pandas()
|
|
result_with_column_df = ds_with_column.to_pandas()
|
|
|
|
# Sort columns before comparison
|
|
result_optimal_df = result_optimal_df[sorted(result_optimal_df.columns)]
|
|
result_with_column_df = result_with_column_df[
|
|
sorted(result_with_column_df.columns)
|
|
]
|
|
|
|
# Compare using rows_same (deterministic, ignores order)
|
|
assert rows_same(result_optimal_df, result_with_column_df)
|
|
|
|
def test_basic_fusion_works(self, ray_start_regular_shared):
|
|
"""Test that basic fusion of two independent operations works."""
|
|
input_data = [{"id": i} for i in range(5)]
|
|
|
|
# Create dataset with two independent operations
|
|
ds = ray.data.from_items(input_data)
|
|
ds = ds.with_column("doubled", col("id") * 2)
|
|
ds = ds.with_column("plus_one", col("id") + 1)
|
|
|
|
# Check before optimization
|
|
original_count = self._count_project_operators(ds._logical_plan)
|
|
print(f"Before optimization: {original_count} operators")
|
|
|
|
# Apply optimization
|
|
rule = ProjectionPushdown()
|
|
optimized_plan = rule.apply(ds._logical_plan)
|
|
|
|
# Check after optimization
|
|
optimized_count = self._count_project_operators(optimized_plan)
|
|
print(f"After optimization: {optimized_count} operators")
|
|
|
|
# Two independent operations should fuse into one
|
|
assert (
|
|
optimized_count == 1
|
|
), f"Two independent operations should fuse to 1 operator, got {optimized_count}"
|
|
|
|
# Verify correctness using rows_same
|
|
from ray.data.dataset import Dataset
|
|
|
|
optimized_ds = Dataset._from_parent(ds, optimized_plan)
|
|
result_df = optimized_ds.to_pandas()
|
|
|
|
expected_df = pd.DataFrame(
|
|
{
|
|
"id": [0, 1, 2, 3, 4],
|
|
"doubled": [0, 2, 4, 6, 8],
|
|
"plus_one": [1, 2, 3, 4, 5],
|
|
}
|
|
)
|
|
|
|
# Sort columns before comparison
|
|
result_df = result_df[sorted(result_df.columns)]
|
|
expected_df = expected_df[sorted(expected_df.columns)]
|
|
assert rows_same(result_df, expected_df)
|
|
|
|
def test_dependency_prevents_fusion(self, ray_start_regular_shared):
|
|
"""Test that dependencies are handled in single operator with OrderedDict."""
|
|
input_data = [{"id": i} for i in range(5)]
|
|
|
|
# Create dataset with dependency chain
|
|
ds = ray.data.from_items(input_data)
|
|
ds = ds.with_column("doubled", col("id") * 2)
|
|
ds = ds.with_column(
|
|
"doubled_plus_one", col("doubled") + 1
|
|
) # Depends on doubled
|
|
|
|
# Check before optimization
|
|
original_count = self._count_project_operators(ds._logical_plan)
|
|
print(f"Before optimization: {original_count} operators")
|
|
|
|
# Apply optimization
|
|
rule = ProjectionPushdown()
|
|
optimized_plan = rule.apply(ds._logical_plan)
|
|
|
|
# Check after optimization
|
|
optimized_count = self._count_project_operators(optimized_plan)
|
|
print(f"After optimization: {optimized_count} operators")
|
|
|
|
# Should have 1 operator now (changed from 2)
|
|
assert (
|
|
optimized_count == 1
|
|
), f"All operations should fuse into 1 operator, got {optimized_count}"
|
|
|
|
# Verify correctness using rows_same
|
|
from ray.data.dataset import Dataset
|
|
|
|
optimized_ds = Dataset._from_parent(ds, optimized_plan)
|
|
result_df = optimized_ds.to_pandas()
|
|
|
|
expected_df = pd.DataFrame(
|
|
{
|
|
"id": [0, 1, 2, 3, 4],
|
|
"doubled": [0, 2, 4, 6, 8],
|
|
"doubled_plus_one": [1, 3, 5, 7, 9],
|
|
}
|
|
)
|
|
|
|
# Sort columns before comparison
|
|
result_df = result_df[sorted(result_df.columns)]
|
|
expected_df = expected_df[sorted(expected_df.columns)]
|
|
assert rows_same(result_df, expected_df)
|
|
|
|
def test_mixed_udf_regular_end_to_end(self, ray_start_regular_shared):
|
|
"""Test the exact failing scenario from the original issue."""
|
|
input_data = [{"id": i} for i in range(5)]
|
|
|
|
# Create dataset with mixed UDF and regular expressions (the failing test case)
|
|
ds = ray.data.from_items(input_data)
|
|
ds = ds.with_column("plus_ten", col("id") + 10)
|
|
ds = ds.with_column(
|
|
"times_three", self.udfs["multiply_by_two"](col("id"))
|
|
) # Actually multiply by 2
|
|
ds = ds.with_column("minus_five", col("id") - 5)
|
|
ds = ds.with_column(
|
|
"udf_plus_regular",
|
|
self.udfs["multiply_by_two"](col("id")) + col("plus_ten"),
|
|
)
|
|
ds = ds.with_column("comparison", col("times_three") > col("plus_ten"))
|
|
|
|
# Apply optimization
|
|
rule = ProjectionPushdown()
|
|
optimized_plan = rule.apply(ds._logical_plan)
|
|
|
|
# Verify execution correctness
|
|
from ray.data.dataset import Dataset
|
|
|
|
optimized_ds = Dataset._from_parent(ds, optimized_plan)
|
|
result_df = optimized_ds.to_pandas()
|
|
|
|
expected_df = pd.DataFrame(
|
|
{
|
|
"id": [0, 1, 2, 3, 4],
|
|
"plus_ten": [10, 11, 12, 13, 14], # id + 10
|
|
"times_three": [0, 2, 4, 6, 8], # id * 2 (multiply_by_two UDF)
|
|
"minus_five": [-5, -4, -3, -2, -1], # id - 5
|
|
"udf_plus_regular": [10, 13, 16, 19, 22], # (id * 2) + (id + 10)
|
|
"comparison": [
|
|
False,
|
|
False,
|
|
False,
|
|
False,
|
|
False,
|
|
], # times_three > plus_ten
|
|
}
|
|
)
|
|
|
|
# Sort columns before comparison
|
|
result_df = result_df[sorted(result_df.columns)]
|
|
expected_df = expected_df[sorted(expected_df.columns)]
|
|
assert rows_same(result_df, expected_df)
|
|
|
|
# Verify that we have 1 operator (changed from multiple)
|
|
optimized_count = self._count_project_operators(optimized_plan)
|
|
assert (
|
|
optimized_count == 1
|
|
), f"Expected 1 operator with all expressions fused, got {optimized_count}"
|
|
|
|
def test_optimal_fusion_comparison(self, ray_start_regular_shared):
|
|
"""Compare optimized with_column approach against manual map_batches."""
|
|
input_data = [{"id": i} for i in range(10)]
|
|
|
|
# Create dataset using with_column (will be optimized)
|
|
ds_with_column = ray.data.from_items(input_data)
|
|
ds_with_column = ds_with_column.with_column("col1", col("id") + 1)
|
|
ds_with_column = ds_with_column.with_column("col2", col("id") * 2)
|
|
ds_with_column = ds_with_column.with_column("col3", col("id") - 1)
|
|
ds_with_column = ds_with_column.with_column("col4", col("id") + 5)
|
|
ds_with_column = ds_with_column.with_column("col5", col("id") * 3)
|
|
|
|
# Apply optimization
|
|
rule = ProjectionPushdown()
|
|
optimized_plan = rule.apply(ds_with_column._logical_plan)
|
|
from ray.data.dataset import Dataset
|
|
|
|
optimized_ds = Dataset._from_parent(ds_with_column, optimized_plan)
|
|
|
|
# Create dataset using single map_batches (optimal case)
|
|
ds_optimal = ray.data.from_items(input_data)
|
|
|
|
def apply_all_expressions(batch):
|
|
import pyarrow.compute as pc
|
|
|
|
result = batch.to_pydict()
|
|
result["col1"] = pc.add(batch["id"], 1)
|
|
result["col2"] = pc.multiply(batch["id"], 2)
|
|
result["col3"] = pc.subtract(batch["id"], 1)
|
|
result["col4"] = pc.add(batch["id"], 5)
|
|
result["col5"] = pc.multiply(batch["id"], 3)
|
|
return pa.table(result)
|
|
|
|
ds_optimal = ds_optimal.map_batches(
|
|
apply_all_expressions, batch_format="pyarrow"
|
|
)
|
|
|
|
# Compare results using rows_same
|
|
result_optimized = optimized_ds.to_pandas()
|
|
result_optimal = ds_optimal.to_pandas()
|
|
|
|
# Sort columns before comparison
|
|
result_optimized = result_optimized[sorted(result_optimized.columns)]
|
|
result_optimal = result_optimal[sorted(result_optimal.columns)]
|
|
assert rows_same(result_optimized, result_optimal)
|
|
|
|
def test_chained_udf_dependencies(self, ray_start_regular_shared):
|
|
"""Test multiple non-vectorized UDFs in a dependency chain."""
|
|
input_data = [{"id": i} for i in range(5)]
|
|
|
|
# Create dataset with chained UDF dependencies
|
|
ds = ray.data.from_items(input_data)
|
|
ds = ds.with_column("plus_one", self.udfs["add_one"](col("id")))
|
|
ds = ds.with_column("times_two", self.udfs["multiply_by_two"](col("plus_one")))
|
|
ds = ds.with_column("div_three", self.udfs["divide_by_three"](col("times_two")))
|
|
|
|
# Apply optimization
|
|
rule = ProjectionPushdown()
|
|
optimized_plan = rule.apply(ds._logical_plan)
|
|
|
|
# Verify 1 operator (changed from 3)
|
|
assert self._count_project_operators(optimized_plan) == 1
|
|
assert (
|
|
self._describe_plan_structure(optimized_plan)
|
|
== "Project(4 exprs) -> FromItems" # Changed from multiple operators
|
|
)
|
|
|
|
# Verify execution correctness
|
|
from ray.data.dataset import Dataset
|
|
|
|
optimized_ds = Dataset._from_parent(ds, optimized_plan)
|
|
result_df = optimized_ds.to_pandas()
|
|
|
|
expected_df = pd.DataFrame(
|
|
{
|
|
"id": [0, 1, 2, 3, 4],
|
|
"plus_one": [1, 2, 3, 4, 5],
|
|
"times_two": [2, 4, 6, 8, 10],
|
|
"div_three": [2 / 3, 4 / 3, 2.0, 8 / 3, 10 / 3],
|
|
}
|
|
)
|
|
|
|
# Sort columns before comparison
|
|
result_df = result_df[sorted(result_df.columns)]
|
|
expected_df = expected_df[sorted(expected_df.columns)]
|
|
assert rows_same(result_df, expected_df)
|
|
|
|
def test_performance_impact_of_udf_chains(self, ray_start_regular_shared):
|
|
"""Test performance characteristics of UDF dependency chains vs independent UDFs."""
|
|
input_data = [{"id": i} for i in range(100)]
|
|
|
|
# Case 1: Independent UDFs (should fuse)
|
|
ds_independent = ray.data.from_items(input_data)
|
|
ds_independent = ds_independent.with_column(
|
|
"udf1", self.udfs["add_one"](col("id"))
|
|
)
|
|
ds_independent = ds_independent.with_column(
|
|
"udf2", self.udfs["multiply_by_two"](col("id"))
|
|
)
|
|
ds_independent = ds_independent.with_column(
|
|
"udf3", self.udfs["divide_by_three"](col("id"))
|
|
)
|
|
|
|
# Case 2: Chained UDFs (should also fuse now)
|
|
ds_chained = ray.data.from_items(input_data)
|
|
ds_chained = ds_chained.with_column("step1", self.udfs["add_one"](col("id")))
|
|
ds_chained = ds_chained.with_column(
|
|
"step2", self.udfs["multiply_by_two"](col("step1"))
|
|
)
|
|
ds_chained = ds_chained.with_column(
|
|
"step3", self.udfs["divide_by_three"](col("step2"))
|
|
)
|
|
|
|
# Apply optimization
|
|
rule = ProjectionPushdown()
|
|
optimized_independent = rule.apply(ds_independent._logical_plan)
|
|
optimized_chained = rule.apply(ds_chained._logical_plan)
|
|
|
|
# Verify fusion behavior (both should be 1 now)
|
|
assert self._count_project_operators(optimized_independent) == 1
|
|
assert (
|
|
self._count_project_operators(optimized_chained) == 1
|
|
) # Changed from 3 to 1
|
|
assert (
|
|
self._describe_plan_structure(optimized_independent)
|
|
== "Project(4 exprs) -> FromItems"
|
|
)
|
|
assert (
|
|
self._describe_plan_structure(optimized_chained)
|
|
== "Project(4 exprs) -> FromItems" # Changed from multiple operators
|
|
)
|
|
|
|
@pytest.mark.parametrize(
|
|
"operations,expected",
|
|
[
|
|
# Single operations
|
|
([("rename", {"a": "A"})], {"A": 1, "b": 2, "c": 3}),
|
|
([("select", ["a", "b"])], {"a": 1, "b": 2}),
|
|
([("with_column", "d", 4)], {"a": 1, "b": 2, "c": 3, "d": 4}),
|
|
# Two operations - rename then select
|
|
([("rename", {"a": "A"}), ("select", ["A"])], {"A": 1}),
|
|
([("rename", {"a": "A"}), ("select", ["b"])], {"b": 2}),
|
|
(
|
|
[("rename", {"a": "A", "b": "B"}), ("select", ["A", "B"])],
|
|
{"A": 1, "B": 2},
|
|
),
|
|
# Two operations - select then rename
|
|
([("select", ["a", "b"]), ("rename", {"a": "A"})], {"A": 1, "b": 2}),
|
|
([("select", ["a"]), ("rename", {"a": "x"})], {"x": 1}),
|
|
# Two operations - with_column combinations
|
|
([("with_column", "d", 4), ("select", ["a", "d"])], {"a": 1, "d": 4}),
|
|
([("select", ["a"]), ("with_column", "d", 4)], {"a": 1, "d": 4}),
|
|
(
|
|
[("rename", {"a": "A"}), ("with_column", "d", 4)],
|
|
{"A": 1, "b": 2, "c": 3, "d": 4},
|
|
),
|
|
(
|
|
[("with_column", "d", 4), ("rename", {"d": "D"})],
|
|
{"a": 1, "b": 2, "c": 3, "D": 4},
|
|
),
|
|
# Three operations
|
|
(
|
|
[
|
|
("rename", {"a": "A"}),
|
|
("select", ["A", "b"]),
|
|
("with_column", "d", 4),
|
|
],
|
|
{"A": 1, "b": 2, "d": 4},
|
|
),
|
|
(
|
|
[
|
|
("with_column", "d", 4),
|
|
("rename", {"a": "A"}),
|
|
("select", ["A", "d"]),
|
|
],
|
|
{"A": 1, "d": 4},
|
|
),
|
|
(
|
|
[
|
|
("select", ["a", "b"]),
|
|
("rename", {"a": "x"}),
|
|
("with_column", "d", 4),
|
|
],
|
|
{"x": 1, "b": 2, "d": 4},
|
|
),
|
|
# Column swap (no actual changes)
|
|
([("rename", {"a": "b", "b": "a"}), ("select", ["a"])], {"a": 2}),
|
|
([("rename", {"a": "b", "b": "a"}), ("select", ["b"])], {"b": 1}),
|
|
# Multiple same operations
|
|
(
|
|
[("rename", {"a": "x"}), ("rename", {"x": "y"})],
|
|
{"y": 1, "b": 2, "c": 3},
|
|
),
|
|
([("select", ["a", "b"]), ("select", ["a"])], {"a": 1}),
|
|
(
|
|
[("with_column", "d", 4), ("with_column", "e", 5)],
|
|
{"a": 1, "b": 2, "c": 3, "d": 4, "e": 5},
|
|
),
|
|
# Complex expressions with with_column
|
|
(
|
|
[("rename", {"a": "x"}), ("with_column_expr", "sum", "x", 10)],
|
|
{"x": 1, "b": 2, "c": 3, "sum": 10},
|
|
),
|
|
(
|
|
[
|
|
("with_column", "d", 4),
|
|
("with_column", "e", 5),
|
|
("select", ["d", "e"]),
|
|
],
|
|
{"d": 4, "e": 5},
|
|
),
|
|
],
|
|
)
|
|
def test_projection_operations_comprehensive(
|
|
self, ray_start_regular_shared, operations, expected
|
|
):
|
|
"""Comprehensive test for projection operations combinations."""
|
|
from ray.data.expressions import col, lit
|
|
|
|
# Create initial dataset
|
|
ds = ray.data.range(1).map(lambda row: {"a": 1, "b": 2, "c": 3})
|
|
|
|
# Apply operations
|
|
for op in operations:
|
|
if op[0] == "rename":
|
|
ds = ds.rename_columns(op[1])
|
|
elif op[0] == "select":
|
|
ds = ds.select_columns(op[1])
|
|
elif op[0] == "with_column":
|
|
ds = ds.with_column(op[1], lit(op[2]))
|
|
elif op[0] == "with_column_expr":
|
|
# Special case for expressions referencing columns
|
|
ds = ds.with_column(op[1], col(op[2]) * op[3])
|
|
|
|
# Verify result using rows_same
|
|
result_df = ds.to_pandas()
|
|
expected_df = pd.DataFrame([expected])
|
|
# Ensure columns are in the same order for comparison
|
|
result_df = result_df[sorted(result_df.columns)]
|
|
expected_df = expected_df[sorted(expected_df.columns)]
|
|
assert rows_same(result_df, expected_df)
|
|
|
|
def test_with_column_alias_then_rename_preserves_both_columns(
|
|
self, ray_start_regular_shared
|
|
):
|
|
"""Regression test for alias and rename referencing the same input column."""
|
|
ds = ray.data.from_items([{"a": 1}])
|
|
ds = ds.with_column("x", col("a")).rename_columns({"a": "b"})
|
|
|
|
optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
|
|
assert self._count_project_operators(optimized_plan) == 1
|
|
|
|
assert ds.take_all() == [{"x": 1, "b": 1}]
|
|
|
|
@pytest.mark.parametrize(
|
|
"operations,expected",
|
|
[
|
|
# Basic count operations
|
|
([("count",)], 3), # All 3 rows
|
|
([("rename", {"a": "A"}), ("count",)], 3),
|
|
([("select", ["a", "b"]), ("count",)], 3),
|
|
([("with_column", "d", 4), ("count",)], 3),
|
|
# Filter operations affecting count
|
|
([("filter", col("a") > 1), ("count",)], 2), # 2 rows have a > 1
|
|
([("filter", col("b") == 2), ("count",)], 3), # All rows have b == 2
|
|
([("filter", col("c") < 10), ("count",)], 3), # All rows have c < 10
|
|
([("filter", col("a") == 1), ("count",)], 1), # 1 row has a == 1
|
|
# Projection then filter then count
|
|
([("rename", {"a": "A"}), ("filter", col("A") > 1), ("count",)], 2),
|
|
([("select", ["a", "b"]), ("filter", col("a") > 1), ("count",)], 2),
|
|
([("with_column", "d", 4), ("filter", col("d") == 4), ("count",)], 3),
|
|
# Filter then projection then count
|
|
([("filter", col("a") > 1), ("rename", {"a": "A"}), ("count",)], 2),
|
|
([("filter", col("b") == 2), ("select", ["a", "b"]), ("count",)], 3),
|
|
([("filter", col("c") < 10), ("with_column", "d", 4), ("count",)], 3),
|
|
# Multiple projections with filter and count
|
|
(
|
|
[
|
|
("rename", {"a": "A"}),
|
|
("select", ["A", "b"]),
|
|
("filter", col("A") > 1),
|
|
("count",),
|
|
],
|
|
2,
|
|
),
|
|
(
|
|
[
|
|
("with_column", "d", 4),
|
|
("rename", {"d": "D"}),
|
|
("filter", col("D") == 4),
|
|
("count",),
|
|
],
|
|
3,
|
|
),
|
|
(
|
|
[
|
|
("select", ["a", "b"]),
|
|
("filter", col("a") > 1),
|
|
("rename", {"a": "x"}),
|
|
("count",),
|
|
],
|
|
2,
|
|
),
|
|
# Complex combinations
|
|
(
|
|
[
|
|
("filter", col("a") > 0),
|
|
("rename", {"b": "B"}),
|
|
("select", ["a", "B"]),
|
|
("filter", col("B") == 2),
|
|
("count",),
|
|
],
|
|
3,
|
|
),
|
|
(
|
|
[
|
|
("with_column", "sum", 99),
|
|
("filter", col("a") > 1),
|
|
("select", ["a", "sum"]),
|
|
("count",),
|
|
],
|
|
2,
|
|
),
|
|
(
|
|
[
|
|
("rename", {"a": "A", "b": "B"}),
|
|
("filter", (col("A") + col("B")) > 3),
|
|
("select", ["A"]),
|
|
("count",),
|
|
],
|
|
2,
|
|
),
|
|
],
|
|
)
|
|
def test_projection_fusion_with_count_and_filter(
|
|
self, ray_start_regular_shared, operations, expected
|
|
):
|
|
"""Test projection fusion with count operations including filters."""
|
|
from ray.data.expressions import lit
|
|
|
|
# Create dataset with 3 rows: {"a": 1, "b": 2, "c": 3}, {"a": 2, "b": 2, "c": 3}, {"a": 3, "b": 2, "c": 3}
|
|
ds = ray.data.from_items(
|
|
[
|
|
{"a": 1, "b": 2, "c": 3},
|
|
{"a": 2, "b": 2, "c": 3},
|
|
{"a": 3, "b": 2, "c": 3},
|
|
]
|
|
)
|
|
|
|
# Apply operations
|
|
for op in operations:
|
|
if op[0] == "rename":
|
|
ds = ds.rename_columns(op[1])
|
|
elif op[0] == "select":
|
|
ds = ds.select_columns(op[1])
|
|
elif op[0] == "with_column":
|
|
ds = ds.with_column(op[1], lit(op[2]))
|
|
elif op[0] == "filter":
|
|
# Use the predicate expression directly
|
|
ds = ds.filter(expr=op[1])
|
|
elif op[0] == "count":
|
|
# Count returns a scalar, not a dataset
|
|
result = ds.count()
|
|
assert result == expected
|
|
return # Early return since count() terminates the pipeline
|
|
|
|
# This should not be reached for count operations
|
|
assert False, "Count operation should have returned early"
|
|
|
|
@pytest.mark.parametrize(
|
|
"invalid_operations,error_type,error_message_contains",
|
|
[
|
|
# Try to filter on a column that doesn't exist yet
|
|
(
|
|
[("filter", col("d") > 0), ("with_column", "d", 4)],
|
|
(KeyError, ray.exceptions.RayTaskError),
|
|
"d",
|
|
),
|
|
# Try to filter on a renamed column before the rename
|
|
(
|
|
[("filter", col("A") > 1), ("rename", {"a": "A"})],
|
|
(KeyError, ray.exceptions.RayTaskError),
|
|
"A",
|
|
),
|
|
# Try to use a column that was removed by select
|
|
(
|
|
[("select", ["a"]), ("filter", col("b") == 2)],
|
|
(KeyError, ray.exceptions.RayTaskError),
|
|
"b",
|
|
),
|
|
# Try to filter on a column after it was removed by select
|
|
(
|
|
[("select", ["a", "b"]), ("filter", col("c") < 10)],
|
|
(KeyError, ray.exceptions.RayTaskError),
|
|
"c",
|
|
),
|
|
# Try to use with_column referencing a non-existent column
|
|
(
|
|
[("select", ["a"]), ("with_column", "new_col", col("b") + 1)],
|
|
(KeyError, ray.exceptions.RayTaskError),
|
|
"b",
|
|
),
|
|
# Try to filter on a column that was renamed away
|
|
(
|
|
[("rename", {"b": "B"}), ("filter", col("b") == 2)],
|
|
(KeyError, ray.exceptions.RayTaskError),
|
|
"b",
|
|
),
|
|
# Try to use with_column with old column name after rename
|
|
(
|
|
[("rename", {"a": "A"}), ("with_column", "result", col("a") + 1)],
|
|
(KeyError, ray.exceptions.RayTaskError),
|
|
"a",
|
|
),
|
|
# Try to select using old column name after rename
|
|
(
|
|
[("rename", {"b": "B"}), ("select", ["a", "b", "c"])],
|
|
(KeyError, ray.exceptions.RayTaskError),
|
|
"b",
|
|
),
|
|
# Try to filter on a computed column that was removed by select
|
|
(
|
|
[
|
|
("with_column", "d", 4),
|
|
("select", ["a", "b"]),
|
|
("filter", col("d") == 4),
|
|
],
|
|
(KeyError, ray.exceptions.RayTaskError),
|
|
"d",
|
|
),
|
|
# Try to rename a column that was removed by select
|
|
(
|
|
[("select", ["a", "b"]), ("rename", {"c": "C"})],
|
|
(KeyError, ray.exceptions.RayTaskError),
|
|
"c",
|
|
),
|
|
# Complex: rename, select (removing renamed source), then use old name
|
|
(
|
|
[
|
|
("rename", {"a": "A"}),
|
|
("select", ["b", "c"]),
|
|
("filter", col("a") > 0),
|
|
],
|
|
(KeyError, ray.exceptions.RayTaskError),
|
|
"a",
|
|
),
|
|
# Complex: with_column, select (keeping new column), filter on removed original
|
|
(
|
|
[
|
|
("with_column", "sum", col("a") + col("b")),
|
|
("select", ["sum"]),
|
|
("filter", col("a") > 0),
|
|
],
|
|
(KeyError, ray.exceptions.RayTaskError),
|
|
"a",
|
|
),
|
|
# Try to use column in with_column expression after it was removed
|
|
(
|
|
[
|
|
("select", ["a", "c"]),
|
|
("with_column", "result", col("a") + col("b")),
|
|
],
|
|
(KeyError, ray.exceptions.RayTaskError),
|
|
"b",
|
|
),
|
|
],
|
|
)
|
|
def test_projection_operations_invalid_order(
|
|
self,
|
|
ray_start_regular_shared,
|
|
invalid_operations,
|
|
error_type,
|
|
error_message_contains,
|
|
):
|
|
"""Test that operations fail gracefully when referencing non-existent columns."""
|
|
import ray
|
|
from ray.data.expressions import lit
|
|
|
|
# Create dataset with 3 rows: {"a": 1, "b": 2, "c": 3}, {"a": 2, "b": 2, "c": 3}, {"a": 3, "b": 2, "c": 3}
|
|
ds = ray.data.from_items(
|
|
[
|
|
{"a": 1, "b": 2, "c": 3},
|
|
{"a": 2, "b": 2, "c": 3},
|
|
{"a": 3, "b": 2, "c": 3},
|
|
]
|
|
)
|
|
|
|
# Apply operations and expect them to fail
|
|
with pytest.raises(error_type) as exc_info:
|
|
for op in invalid_operations:
|
|
if op[0] == "rename":
|
|
ds = ds.rename_columns(op[1])
|
|
elif op[0] == "select":
|
|
ds = ds.select_columns(op[1])
|
|
elif op[0] == "with_column":
|
|
if len(op) == 3 and not isinstance(op[2], (int, float, str)):
|
|
# Expression-based with_column (op[2] is an expression)
|
|
ds = ds.with_column(op[1], op[2])
|
|
else:
|
|
# Literal-based with_column
|
|
ds = ds.with_column(op[1], lit(op[2]))
|
|
elif op[0] == "filter":
|
|
ds = ds.filter(expr=op[1])
|
|
elif op[0] == "count":
|
|
ds.count()
|
|
return
|
|
|
|
# Force execution to trigger the error
|
|
result = ds.take_all()
|
|
print(f"Unexpected success: {result}")
|
|
|
|
# Verify the error message contains the expected column name
|
|
error_str = str(exc_info.value).lower()
|
|
assert (
|
|
error_message_contains.lower() in error_str
|
|
), f"Expected '{error_message_contains}' in error message: {error_str}"
|
|
|
|
@pytest.mark.parametrize(
|
|
"operations,expected_output",
|
|
[
|
|
# === Basic Select Operations ===
|
|
pytest.param(
|
|
[("select", ["a"])],
|
|
[{"a": 1}, {"a": 2}, {"a": 3}],
|
|
id="select_single_column",
|
|
),
|
|
pytest.param(
|
|
[("select", ["a", "b"])],
|
|
[{"a": 1, "b": 4}, {"a": 2, "b": 5}, {"a": 3, "b": 6}],
|
|
id="select_two_columns",
|
|
),
|
|
pytest.param(
|
|
[("select", ["a", "b", "c"])],
|
|
[
|
|
{"a": 1, "b": 4, "c": 7},
|
|
{"a": 2, "b": 5, "c": 8},
|
|
{"a": 3, "b": 6, "c": 9},
|
|
],
|
|
id="select_all_columns",
|
|
),
|
|
pytest.param(
|
|
[("select", ["c", "a"])],
|
|
[{"c": 7, "a": 1}, {"c": 8, "a": 2}, {"c": 9, "a": 3}],
|
|
id="select_reordered_columns",
|
|
),
|
|
# === Basic Rename Operations ===
|
|
pytest.param(
|
|
[("rename", {"a": "alpha"})],
|
|
[
|
|
{"alpha": 1, "b": 4, "c": 7},
|
|
{"alpha": 2, "b": 5, "c": 8},
|
|
{"alpha": 3, "b": 6, "c": 9},
|
|
],
|
|
id="rename_single_column",
|
|
),
|
|
pytest.param(
|
|
[("rename", {"a": "alpha", "b": "beta"})],
|
|
[
|
|
{"alpha": 1, "beta": 4, "c": 7},
|
|
{"alpha": 2, "beta": 5, "c": 8},
|
|
{"alpha": 3, "beta": 6, "c": 9},
|
|
],
|
|
id="rename_multiple_columns",
|
|
),
|
|
# === Basic with_column Operations ===
|
|
pytest.param(
|
|
[("with_column_expr", "sum", "add", "a", "b")],
|
|
[
|
|
{"a": 1, "b": 4, "c": 7, "sum": 5},
|
|
{"a": 2, "b": 5, "c": 8, "sum": 7},
|
|
{"a": 3, "b": 6, "c": 9, "sum": 9},
|
|
],
|
|
id="with_column_add_keep_all",
|
|
),
|
|
pytest.param(
|
|
[("with_column_expr", "product", "multiply", "b", "c")],
|
|
[
|
|
{"a": 1, "b": 4, "c": 7, "product": 28},
|
|
{"a": 2, "b": 5, "c": 8, "product": 40},
|
|
{"a": 3, "b": 6, "c": 9, "product": 54},
|
|
],
|
|
id="with_column_multiply_keep_all",
|
|
),
|
|
# === Chained Selects ===
|
|
pytest.param(
|
|
[("select", ["a", "b", "c"]), ("select", ["a", "b"])],
|
|
[{"a": 1, "b": 4}, {"a": 2, "b": 5}, {"a": 3, "b": 6}],
|
|
id="chained_selects_two_levels",
|
|
),
|
|
pytest.param(
|
|
[
|
|
("select", ["a", "b", "c"]),
|
|
("select", ["a", "b"]),
|
|
("select", ["a"]),
|
|
],
|
|
[{"a": 1}, {"a": 2}, {"a": 3}],
|
|
id="chained_selects_three_levels",
|
|
),
|
|
# === Rename → Select ===
|
|
pytest.param(
|
|
[("rename", {"a": "x"}), ("select", ["x", "b"])],
|
|
[{"x": 1, "b": 4}, {"x": 2, "b": 5}, {"x": 3, "b": 6}],
|
|
id="rename_then_select",
|
|
),
|
|
pytest.param(
|
|
[("rename", {"a": "x", "c": "z"}), ("select", ["x", "z"])],
|
|
[{"x": 1, "z": 7}, {"x": 2, "z": 8}, {"x": 3, "z": 9}],
|
|
id="rename_multiple_then_select",
|
|
),
|
|
# === Select → Rename ===
|
|
pytest.param(
|
|
[("select", ["a", "b"]), ("rename", {"a": "x"})],
|
|
[{"x": 1, "b": 4}, {"x": 2, "b": 5}, {"x": 3, "b": 6}],
|
|
id="select_then_rename",
|
|
),
|
|
pytest.param(
|
|
[("select", ["a", "b", "c"]), ("rename", {"a": "x", "b": "y"})],
|
|
[
|
|
{"x": 1, "y": 4, "c": 7},
|
|
{"x": 2, "y": 5, "c": 8},
|
|
{"x": 3, "y": 6, "c": 9},
|
|
],
|
|
id="select_all_then_rename_some",
|
|
),
|
|
# === Multiple Renames ===
|
|
pytest.param(
|
|
[("rename", {"a": "x"}), ("rename", {"x": "y"})],
|
|
[
|
|
{"y": 1, "b": 4, "c": 7},
|
|
{"y": 2, "b": 5, "c": 8},
|
|
{"y": 3, "b": 6, "c": 9},
|
|
],
|
|
id="chained_renames",
|
|
),
|
|
# === with_column → Select ===
|
|
pytest.param(
|
|
[("with_column_expr", "sum", "add", "a", "b"), ("select", ["sum"])],
|
|
[{"sum": 5}, {"sum": 7}, {"sum": 9}],
|
|
id="with_column_then_select_only_computed",
|
|
),
|
|
pytest.param(
|
|
[
|
|
("with_column_expr", "sum", "add", "a", "b"),
|
|
("select", ["a", "sum"]),
|
|
],
|
|
[{"a": 1, "sum": 5}, {"a": 2, "sum": 7}, {"a": 3, "sum": 9}],
|
|
id="with_column_then_select_mixed",
|
|
),
|
|
pytest.param(
|
|
[
|
|
("with_column_expr", "result", "multiply", "b", "c"),
|
|
("select", ["a", "result"]),
|
|
],
|
|
[
|
|
{"a": 1, "result": 28},
|
|
{"a": 2, "result": 40},
|
|
{"a": 3, "result": 54},
|
|
],
|
|
id="with_column_select_source_and_computed",
|
|
),
|
|
# === Multiple with_column Operations ===
|
|
pytest.param(
|
|
[
|
|
("with_column_expr", "sum", "add", "a", "b"),
|
|
("with_column_expr", "product", "multiply", "a", "c"),
|
|
],
|
|
[
|
|
{"a": 1, "b": 4, "c": 7, "sum": 5, "product": 7},
|
|
{"a": 2, "b": 5, "c": 8, "sum": 7, "product": 16},
|
|
{"a": 3, "b": 6, "c": 9, "sum": 9, "product": 27},
|
|
],
|
|
id="multiple_with_column_keep_all",
|
|
),
|
|
pytest.param(
|
|
[
|
|
("with_column_expr", "sum", "add", "a", "b"),
|
|
("with_column_expr", "product", "multiply", "a", "c"),
|
|
("select", ["sum", "product"]),
|
|
],
|
|
[
|
|
{"sum": 5, "product": 7},
|
|
{"sum": 7, "product": 16},
|
|
{"sum": 9, "product": 27},
|
|
],
|
|
id="multiple_with_column_then_select",
|
|
),
|
|
pytest.param(
|
|
[
|
|
("with_column_expr", "sum", "add", "a", "b"),
|
|
("with_column_expr", "diff", "add", "c", "a"),
|
|
("select", ["sum", "diff"]),
|
|
],
|
|
[{"sum": 5, "diff": 8}, {"sum": 7, "diff": 10}, {"sum": 9, "diff": 12}],
|
|
id="multiple_with_column_independent_sources",
|
|
),
|
|
# === with_column → Rename ===
|
|
pytest.param(
|
|
[
|
|
("with_column_expr", "sum", "add", "a", "b"),
|
|
("rename", {"sum": "total"}),
|
|
],
|
|
[
|
|
{"a": 1, "b": 4, "c": 7, "total": 5},
|
|
{"a": 2, "b": 5, "c": 8, "total": 7},
|
|
{"a": 3, "b": 6, "c": 9, "total": 9},
|
|
],
|
|
id="with_column_then_rename_computed",
|
|
),
|
|
# === Rename → with_column ===
|
|
pytest.param(
|
|
[
|
|
("rename", {"a": "x"}),
|
|
("with_column_expr", "x_plus_b", "add", "x", "b"),
|
|
],
|
|
[
|
|
{"x": 1, "b": 4, "c": 7, "x_plus_b": 5},
|
|
{"x": 2, "b": 5, "c": 8, "x_plus_b": 7},
|
|
{"x": 3, "b": 6, "c": 9, "x_plus_b": 9},
|
|
],
|
|
id="rename_then_with_column_using_renamed",
|
|
),
|
|
pytest.param(
|
|
[
|
|
("rename", {"a": "x"}),
|
|
("with_column_expr", "result", "add", "x", "b"),
|
|
("select", ["result"]),
|
|
],
|
|
[{"result": 5}, {"result": 7}, {"result": 9}],
|
|
id="rename_with_column_select_chain",
|
|
),
|
|
# === Select → with_column → Select ===
|
|
pytest.param(
|
|
[
|
|
("select", ["a", "b"]),
|
|
("with_column_expr", "sum", "add", "a", "b"),
|
|
("select", ["a", "sum"]),
|
|
],
|
|
[{"a": 1, "sum": 5}, {"a": 2, "sum": 7}, {"a": 3, "sum": 9}],
|
|
id="select_with_column_select_chain",
|
|
),
|
|
pytest.param(
|
|
[
|
|
("select", ["a", "b", "c"]),
|
|
("with_column_expr", "x", "add", "a", "b"),
|
|
("with_column_expr", "y", "multiply", "b", "c"),
|
|
("select", ["x", "y"]),
|
|
],
|
|
[{"x": 5, "y": 28}, {"x": 7, "y": 40}, {"x": 9, "y": 54}],
|
|
id="select_multiple_with_column_select_chain",
|
|
),
|
|
# === Complex Multi-Step Chains ===
|
|
pytest.param(
|
|
[
|
|
("select", ["a", "b", "c"]),
|
|
("rename", {"a": "x"}),
|
|
("with_column_expr", "result", "add", "x", "b"),
|
|
("select", ["result", "c"]),
|
|
],
|
|
[{"result": 5, "c": 7}, {"result": 7, "c": 8}, {"result": 9, "c": 9}],
|
|
id="complex_select_rename_with_column_select",
|
|
),
|
|
pytest.param(
|
|
[
|
|
("rename", {"a": "alpha", "b": "beta"}),
|
|
("select", ["alpha", "beta", "c"]),
|
|
("with_column_expr", "sum", "add", "alpha", "beta"),
|
|
("rename", {"sum": "total"}),
|
|
("select", ["total", "c"]),
|
|
],
|
|
[{"total": 5, "c": 7}, {"total": 7, "c": 8}, {"total": 9, "c": 9}],
|
|
id="complex_five_step_chain",
|
|
),
|
|
pytest.param(
|
|
[
|
|
("select", ["a", "b", "c"]),
|
|
("select", ["b", "c"]),
|
|
("select", ["c"]),
|
|
],
|
|
[{"c": 7}, {"c": 8}, {"c": 9}],
|
|
id="select_chain",
|
|
),
|
|
],
|
|
)
|
|
def test_projection_pushdown_into_parquet_read(
|
|
self, ray_start_regular_shared, tmp_path, operations, expected_output
|
|
):
|
|
"""Test that projection operations fuse and push down into parquet reads.
|
|
|
|
Verifies:
|
|
- Multiple projections fuse into single operator
|
|
- Fused projection pushes down into Read operator
|
|
- Only necessary columns are read from parquet
|
|
- Results are correct for select, rename, and with_column operations
|
|
"""
|
|
from ray.data.expressions import col
|
|
|
|
# Create test parquet file
|
|
df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]})
|
|
parquet_path = tmp_path / "test.parquet"
|
|
df.to_parquet(parquet_path, index=False)
|
|
|
|
# Build pipeline with operations
|
|
ds = ray.data.read_parquet(str(parquet_path))
|
|
|
|
for op_type, *op_args in operations:
|
|
if op_type == "select":
|
|
ds = ds.select_columns(op_args[0])
|
|
elif op_type == "rename":
|
|
ds = ds.rename_columns(op_args[0])
|
|
elif op_type == "with_column_expr":
|
|
col_name, operator, col1, col2 = op_args
|
|
if operator == "add":
|
|
ds = ds.with_column(col_name, col(col1) + col(col2))
|
|
elif operator == "multiply":
|
|
ds = ds.with_column(col_name, col(col1) * col(col2))
|
|
|
|
result_df = ds.to_pandas()
|
|
expected_df = pd.DataFrame(expected_output)
|
|
# Ensure columns are in the same order for comparison
|
|
result_df = result_df[sorted(result_df.columns)]
|
|
expected_df = expected_df[sorted(expected_df.columns)]
|
|
assert rows_same(result_df, expected_df)
|
|
|
|
|
|
@pytest.mark.parametrize("flavor", ["project_before", "project_after"])
|
|
def test_projection_pushdown_merge_rename_x(ray_start_regular_shared, flavor):
|
|
"""
|
|
Test that valid select and renaming merges correctly.
|
|
"""
|
|
path = "example://iris.parquet"
|
|
ds = ray.data.read_parquet(path)
|
|
ds = ds.map_batches(lambda d: d)
|
|
|
|
if flavor == "project_before":
|
|
ds = ds.select_columns(["sepal.length", "petal.width"])
|
|
|
|
# First projection renames 'sepal.length' to 'length'
|
|
ds = ds.rename_columns({"sepal.length": "length"})
|
|
|
|
# Second projection renames 'petal.width' to 'width'
|
|
ds = ds.rename_columns({"petal.width": "width"})
|
|
|
|
if flavor == "project_after":
|
|
ds = ds.select_columns(["length", "width"])
|
|
|
|
logical_plan = ds._logical_plan
|
|
op = logical_plan.dag
|
|
assert isinstance(op, Project), op.name
|
|
|
|
optimized_logical_plan = LogicalOptimizer().optimize(logical_plan)
|
|
assert isinstance(optimized_logical_plan.dag, Project)
|
|
|
|
select_op = optimized_logical_plan.dag
|
|
|
|
# Check that both "sepal.length" and "petal.width" are present in the columns,
|
|
# regardless of their order.
|
|
assert select_op.exprs == [
|
|
# TODO fix (renaming doesn't remove prev columns)
|
|
col("sepal.length").alias("length"),
|
|
col("petal.width").alias("width"),
|
|
]
|
|
|
|
|
|
def _leaf_op(dag):
|
|
"""Return the source (leaf) operator of a logical DAG (e.g. the read)."""
|
|
op = dag
|
|
while op.input_dependencies:
|
|
op = op.input_dependencies[0]
|
|
return op
|
|
|
|
|
|
def _find_op(dag, op_type):
|
|
"""Walk a single-input DAG and return the first op of ``op_type``."""
|
|
op = dag
|
|
while op is not None:
|
|
if isinstance(op, op_type):
|
|
return op
|
|
op = op.input_dependencies[0] if op.input_dependencies else None
|
|
return None
|
|
|
|
|
|
class TestAggregateInputPruning:
|
|
"""``ProjectionPushdown`` should prune the columns flowing into an
|
|
``Aggregate`` down to the ones it consumes (group keys + aggregation
|
|
targets), and push that pruning into the read."""
|
|
|
|
@pytest.fixture
|
|
def wide_parquet(self, tmp_path):
|
|
import os
|
|
|
|
path = os.path.join(str(tmp_path), "wide.parquet")
|
|
pq.write_table(
|
|
pa.table(
|
|
{
|
|
"k": [i % 3 for i in range(60)],
|
|
"a": [float(i) for i in range(60)],
|
|
"b": [0.5] * 60,
|
|
"unused": ["x" * 32] * 60, # wide column nothing consumes
|
|
}
|
|
),
|
|
path,
|
|
)
|
|
return path
|
|
|
|
@pytest.mark.parametrize(
|
|
"make_result, expected_input_cols",
|
|
[
|
|
pytest.param(
|
|
lambda ds: ds.with_column("revenue", col("a") * col("b"))
|
|
.groupby("k")
|
|
.sum("revenue"),
|
|
{"k", "revenue"},
|
|
id="groupby-sum-computed-column",
|
|
),
|
|
pytest.param(
|
|
lambda ds: ds.groupby("k").sum("a"),
|
|
{"k", "a"},
|
|
id="groupby-sum",
|
|
),
|
|
pytest.param(
|
|
lambda ds: ds.groupby("k").aggregate(Sum("a"), Mean("b")),
|
|
{"k", "a", "b"},
|
|
id="groupby-multi-agg",
|
|
),
|
|
pytest.param(
|
|
# Count reads no column, so only the group key is required.
|
|
lambda ds: ds.groupby("k").count(),
|
|
{"k"},
|
|
id="groupby-count",
|
|
),
|
|
],
|
|
)
|
|
def test_aggregate_input_pruned_to_required_columns(
|
|
self,
|
|
wide_parquet,
|
|
make_result,
|
|
expected_input_cols,
|
|
ray_start_regular_shared,
|
|
):
|
|
from ray.data._internal.logical.operators.all_to_all_operator import Aggregate
|
|
|
|
ds = make_result(ray.data.read_parquet(wide_parquet))
|
|
dag = LogicalOptimizer().optimize(ds._logical_plan).dag
|
|
|
|
# The op feeding the aggregate must carry only the consumed columns.
|
|
agg = _find_op(dag, Aggregate)
|
|
assert (
|
|
set(agg.input_dependencies[0].infer_schema().names) == expected_input_cols
|
|
)
|
|
|
|
# The wide unused column is always dropped from the read.
|
|
assert "unused" not in set(_leaf_op(dag).infer_schema().names)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|