wehub-resource-sync
349 lines
15 KiB
Python
349 lines
15 KiB
Python
"""Unit tests for Parquet predicate splitting optimization.
|
|
|
|
This module tests the _split_predicate_by_columns function which optimizes
|
|
predicate pushdown for partitioned Parquet datasets by splitting predicates
|
|
into data-column, partition-column, and residual parts.
|
|
"""
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Optional, Set
|
|
|
|
import pytest
|
|
|
|
from ray.data._internal.datasource.parquet_datasource import (
|
|
_split_predicate_by_columns,
|
|
)
|
|
from ray.data.expressions import Expr, col
|
|
|
|
|
|
@dataclass
|
|
class PredicateSplitTestCase:
|
|
"""Test case for predicate splitting."""
|
|
|
|
predicate: Expr
|
|
partition_cols: Set[str]
|
|
expected_data_predicate: Optional[Expr]
|
|
expected_partition_predicate: Optional[Expr]
|
|
description: str
|
|
expected_residual_predicate: Optional[Expr] = None
|
|
|
|
|
|
# fmt: off
|
|
TEST_CASES = [
|
|
# ====================================================================
|
|
# Pure data predicates - should push down everything
|
|
# ====================================================================
|
|
PredicateSplitTestCase(
|
|
predicate=col("data1") > 5,
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=col("data1") > 5,
|
|
expected_partition_predicate=None,
|
|
description="Simple data predicate",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=(col("data1") > 5) & (col("data2") == "x"),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=(col("data1") > 5) & (col("data2") == "x"),
|
|
expected_partition_predicate=None,
|
|
description="AND of data predicates",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=(col("data1") > 5) | (col("data2") == "x"),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=(col("data1") > 5) | (col("data2") == "x"),
|
|
expected_partition_predicate=None,
|
|
description="OR of data predicates",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=~(col("data1") > 5),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=~(col("data1") > 5),
|
|
expected_partition_predicate=None,
|
|
description="NOT of data predicate",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=col("data1").is_null(),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=col("data1").is_null(),
|
|
expected_partition_predicate=None,
|
|
description="IS_NULL of data column",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=col("data1").is_not_null(),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=col("data1").is_not_null(),
|
|
expected_partition_predicate=None,
|
|
description="IS_NOT_NULL of data column",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=((col("data1") > 5) & (col("data2") < 10)) | (col("data3") == "test"),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=((col("data1") > 5) & (col("data2") < 10)) | (col("data3") == "test"),
|
|
expected_partition_predicate=None,
|
|
description="Complex nested data predicates",
|
|
),
|
|
# ====================================================================
|
|
# Pure partition predicates - should enable pruning
|
|
# ====================================================================
|
|
PredicateSplitTestCase(
|
|
predicate=col("partition_col") == "US",
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=None,
|
|
expected_partition_predicate=col("partition_col") == "US",
|
|
description="Simple partition predicate",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=(col("partition1") == "US") & (col("partition2") == "2020"),
|
|
partition_cols={"partition1", "partition2"},
|
|
expected_data_predicate=None,
|
|
expected_partition_predicate=(col("partition1") == "US") & (col("partition2") == "2020"),
|
|
description="AND of partition predicates",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=(col("partition1") == "US") | (col("partition2") == "2020"),
|
|
partition_cols={"partition1", "partition2"},
|
|
expected_data_predicate=None,
|
|
expected_partition_predicate=(col("partition1") == "US") | (col("partition2") == "2020"),
|
|
description="OR of partition predicates",
|
|
),
|
|
# ====================================================================
|
|
# Mixed predicates with AND - should split both parts
|
|
# ====================================================================
|
|
PredicateSplitTestCase(
|
|
predicate=(col("data1") > 5) & (col("partition_col") == "US"),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=col("data1") > 5,
|
|
expected_partition_predicate=col("partition_col") == "US",
|
|
description="Simple AND with data and partition",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=(col("partition_col") == "US") & (col("data1") > 5),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=col("data1") > 5,
|
|
expected_partition_predicate=col("partition_col") == "US",
|
|
description="Simple AND with partition and data (reversed order)",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=(col("data1") > 5) & (col("data2") < 10) & (col("partition_col") == "US"),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=(col("data1") > 5) & (col("data2") < 10),
|
|
expected_partition_predicate=col("partition_col") == "US",
|
|
description="Multiple data predicates AND partition",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=(col("data1") > 5) & (col("partition1") == "US") & (col("data2") < 10) & (col("partition2") == "2020"),
|
|
partition_cols={"partition1", "partition2"},
|
|
expected_data_predicate=(col("data1") > 5) & (col("data2") < 10),
|
|
expected_partition_predicate=(col("partition1") == "US") & (col("partition2") == "2020"),
|
|
description="Interleaved data and partition predicates",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=((col("data1") > 5) & (col("data2") < 10)) & (col("partition_col") == "US"),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=(col("data1") > 5) & (col("data2") < 10),
|
|
expected_partition_predicate=col("partition_col") == "US",
|
|
description="Nested AND of data predicates with partition",
|
|
),
|
|
# ====================================================================
|
|
# Mixed predicates with OR - CANNOT split safely
|
|
# ====================================================================
|
|
PredicateSplitTestCase(
|
|
predicate=(col("data1") > 5) | (col("partition_col") == "US"),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=None,
|
|
expected_partition_predicate=None,
|
|
expected_residual_predicate=(col("data1") > 5) | (col("partition_col") == "US"),
|
|
description="OR with data and partition (unsafe to split)",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=(col("partition_col") == "US") | (col("data1") > 5),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=None,
|
|
expected_partition_predicate=None,
|
|
expected_residual_predicate=(col("partition_col") == "US") | (col("data1") > 5),
|
|
description="OR with partition and data (unsafe to split)",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=((col("data1") > 5) & (col("data2") < 10)) | (col("partition_col") == "US"),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=None,
|
|
expected_partition_predicate=None,
|
|
expected_residual_predicate=((col("data1") > 5) & (col("data2") < 10)) | (col("partition_col") == "US"),
|
|
description="OR with complex data predicate and partition",
|
|
),
|
|
# ====================================================================
|
|
# Mixed predicates with NOT - CANNOT split safely
|
|
# ====================================================================
|
|
PredicateSplitTestCase(
|
|
predicate=~(col("partition_col") == "US"),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=None,
|
|
expected_partition_predicate=~(col("partition_col") == "US"),
|
|
description="NOT of partition predicate",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=~((col("data1") > 5) & (col("partition_col") == "US")),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=None,
|
|
expected_partition_predicate=None,
|
|
expected_residual_predicate=~((col("data1") > 5) & (col("partition_col") == "US")),
|
|
description="NOT of mixed AND (becomes OR via De Morgan, unsafe)",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=(col("data1") > 5) & ~(col("partition_col") == "US"),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=col("data1") > 5,
|
|
expected_partition_predicate=~(col("partition_col") == "US"),
|
|
description="AND with NOT of partition predicate (can extract both parts)",
|
|
),
|
|
# ====================================================================
|
|
# Complex nested scenarios
|
|
# ====================================================================
|
|
PredicateSplitTestCase(
|
|
predicate=((col("data1") > 5) & (col("data2") < 10)) & ((col("data3") == "test") & (col("partition_col") == "US")),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=(col("data1") > 5) & (col("data2") < 10) & (col("data3") == "test"),
|
|
expected_partition_predicate=col("partition_col") == "US",
|
|
description="Deeply nested ANDs with mixed columns",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=((col("data1") > 5) | (col("data2") < 10)) & (col("partition_col") == "US"),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=(col("data1") > 5) | (col("data2") < 10),
|
|
expected_partition_predicate=col("partition_col") == "US",
|
|
description="AND of complex data predicate (with OR) and partition",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=(col("data1") > 5) & ((col("data2") < 10) & (col("partition_col") == "US")),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=(col("data1") > 5) & (col("data2") < 10),
|
|
expected_partition_predicate=col("partition_col") == "US",
|
|
description="Left-nested data with right-nested mixed",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=((col("data1") > 5) & (col("partition1") == "US")) & ((col("data2") < 10) & (col("partition2") == "2020")),
|
|
partition_cols={"partition1", "partition2"},
|
|
expected_data_predicate=(col("data1") > 5) & (col("data2") < 10),
|
|
expected_partition_predicate=(col("partition1") == "US") & (col("partition2") == "2020"),
|
|
description="Both sides have mixed predicates",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=((col("data1") > 5) | (col("partition_col") == "US")) & (col("data2") < 10),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=col("data2") < 10,
|
|
expected_partition_predicate=None,
|
|
expected_residual_predicate=(col("data1") > 5) | (col("partition_col") == "US"),
|
|
description="AND with left side having OR with partition (residual carries the unsplittable OR)",
|
|
),
|
|
# ====================================================================
|
|
# Edge cases
|
|
# ====================================================================
|
|
PredicateSplitTestCase(
|
|
predicate=(col("data1") > 5) & (col("data1") < 10),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=(col("data1") > 5) & (col("data1") < 10),
|
|
expected_partition_predicate=None,
|
|
description="Same column referenced multiple times (data)",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=(col("partition_col") == "US") & (col("partition_col") != "UK"),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=None,
|
|
expected_partition_predicate=(col("partition_col") == "US") & (col("partition_col") != "UK"),
|
|
description="Same partition column referenced multiple times",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=col("data1").is_null() & (col("partition_col") == "US"),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=col("data1").is_null(),
|
|
expected_partition_predicate=col("partition_col") == "US",
|
|
description="IS_NULL data predicate with partition",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=col("partition_col").is_null(),
|
|
partition_cols={"partition_col"},
|
|
expected_data_predicate=None,
|
|
expected_partition_predicate=col("partition_col").is_null(),
|
|
description="IS_NULL on partition column",
|
|
),
|
|
# ====================================================================
|
|
# No partition columns in dataset
|
|
# ====================================================================
|
|
PredicateSplitTestCase(
|
|
predicate=col("data1") > 5,
|
|
partition_cols=set(),
|
|
expected_data_predicate=col("data1") > 5,
|
|
expected_partition_predicate=None,
|
|
description="No partition columns in dataset",
|
|
),
|
|
PredicateSplitTestCase(
|
|
predicate=(col("data1") > 5) & (col("data2") < 10),
|
|
partition_cols=set(),
|
|
expected_data_predicate=(col("data1") > 5) & (col("data2") < 10),
|
|
expected_partition_predicate=None,
|
|
description="Complex predicate with no partition columns",
|
|
),
|
|
# ====================================================================
|
|
# All columns are partition columns
|
|
# ====================================================================
|
|
PredicateSplitTestCase(
|
|
predicate=col("partition1") > 5,
|
|
partition_cols={"partition1", "partition2"},
|
|
expected_data_predicate=None,
|
|
expected_partition_predicate=col("partition1") > 5,
|
|
description="All columns are partition columns",
|
|
),
|
|
]
|
|
|
|
def _assert_predicate_matches(
|
|
actual: Optional[Expr], expected: Optional[Expr], pred_type: str, description: str
|
|
):
|
|
"""Helper to assert predicate matches expected value."""
|
|
if expected is None:
|
|
assert actual is None, (
|
|
f"{description}: Expected no {pred_type} predicate (None), got {actual}"
|
|
)
|
|
else:
|
|
assert actual is not None, f"{description}: Expected {pred_type} predicate, got None"
|
|
|
|
|
|
@pytest.mark.parametrize("test_case", TEST_CASES, ids=lambda tc: tc.description)
|
|
def test_split_predicate_by_columns(test_case: PredicateSplitTestCase):
|
|
"""Test predicate splitting for various scenarios.
|
|
|
|
This test covers:
|
|
- Pure data predicates (should extract data part only)
|
|
- Pure partition predicates (should extract partition part only)
|
|
- Mixed predicates with AND (should split both parts)
|
|
- Mixed predicates with OR (kept as residual)
|
|
- Mixed predicates with NOT (varies by case)
|
|
- Complex nested scenarios with residual carry-over
|
|
- Edge cases
|
|
"""
|
|
result = _split_predicate_by_columns(test_case.predicate, test_case.partition_cols)
|
|
|
|
_assert_predicate_matches(
|
|
result.data_predicate,
|
|
test_case.expected_data_predicate,
|
|
"data",
|
|
test_case.description,
|
|
)
|
|
_assert_predicate_matches(
|
|
result.partition_predicate,
|
|
test_case.expected_partition_predicate,
|
|
"partition",
|
|
test_case.description,
|
|
)
|
|
_assert_predicate_matches(
|
|
result.residual_predicate,
|
|
test_case.expected_residual_predicate,
|
|
"residual",
|
|
test_case.description,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|