| ################################################################################ | |
| # Licensed to the Apache Software Foundation (ASF) under one | |
| # or more contributor license agreements. See the NOTICE file | |
| # distributed with this work for additional information | |
| # regarding copyright ownership. The ASF licenses this file | |
| # to you under the Apache License, Version 2.0 (the | |
| # "License"); you may not use this file except in compliance | |
| # with the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| ################################################################################ | |
| import datetime | |
| import json | |
| import unittest | |
| import pyarrow as pa | |
| from pypaimon.common.predicate_json_parser import ( | |
| _apply_predicate_transform, | |
| _convert_literal, | |
| _paimon_type_to_arrow, | |
| extract_referenced_fields, | |
| parse_predicate_to_batch_filter, | |
| ) | |
| def _make_leaf(field_name, function, literals=None, field_type="INT"): | |
| d = { | |
| "kind": "LEAF", | |
| "transform": { | |
| "name": "FIELD_REF", | |
| "fieldRef": {"index": 0, "name": field_name, "type": field_type}, | |
| }, | |
| "function": function, | |
| } | |
| if literals is not None: | |
| d["literals"] = literals | |
| return json.dumps(d) | |
| class TestLeafFunctions(unittest.TestCase): | |
| def setUp(self): | |
| self.batch = pa.RecordBatch.from_pydict({ | |
| "id": [1, 2, 3, 4, 5], | |
| "name": ["alice", "bob", "charlie", "alice", "eve"], | |
| "score": [85.5, 92.0, 78.3, 95.1, 60.0], | |
| }) | |
| def _filter(self, json_str): | |
| return parse_predicate_to_batch_filter(json_str)(self.batch).to_pylist() | |
| def test_equal(self): | |
| self.assertEqual( | |
| self._filter(_make_leaf("id", "EQUAL", [3])), | |
| [False, False, True, False, False], | |
| ) | |
| def test_not_equal(self): | |
| self.assertEqual( | |
| self._filter(_make_leaf("id", "NOT_EQUAL", [3])), | |
| [True, True, False, True, True], | |
| ) | |
| def test_less_than(self): | |
| self.assertEqual( | |
| self._filter(_make_leaf("id", "LESS_THAN", [3])), | |
| [True, True, False, False, False], | |
| ) | |
| def test_less_or_equal(self): | |
| self.assertEqual( | |
| self._filter(_make_leaf("id", "LESS_OR_EQUAL", [3])), | |
| [True, True, True, False, False], | |
| ) | |
| def test_greater_than(self): | |
| self.assertEqual( | |
| self._filter(_make_leaf("id", "GREATER_THAN", [3])), | |
| [False, False, False, True, True], | |
| ) | |
| def test_greater_or_equal(self): | |
| self.assertEqual( | |
| self._filter(_make_leaf("id", "GREATER_OR_EQUAL", [3])), | |
| [False, False, True, True, True], | |
| ) | |
| def test_is_null(self): | |
| batch = pa.RecordBatch.from_pydict({"val": [1, None, 3, None, 5]}) | |
| f = parse_predicate_to_batch_filter(_make_leaf("val", "IS_NULL")) | |
| self.assertEqual(f(batch).to_pylist(), [False, True, False, True, False]) | |
| def test_is_not_null(self): | |
| batch = pa.RecordBatch.from_pydict({"val": [1, None, 3, None, 5]}) | |
| f = parse_predicate_to_batch_filter(_make_leaf("val", "IS_NOT_NULL")) | |
| self.assertEqual(f(batch).to_pylist(), [True, False, True, False, True]) | |
| def test_in(self): | |
| self.assertEqual( | |
| self._filter(_make_leaf("id", "IN", [1, 3, 5])), | |
| [True, False, True, False, True], | |
| ) | |
| def test_not_in(self): | |
| self.assertEqual( | |
| self._filter(_make_leaf("id", "NOT_IN", [1, 3, 5])), | |
| [False, True, False, True, False], | |
| ) | |
| def test_between(self): | |
| self.assertEqual( | |
| self._filter(_make_leaf("id", "BETWEEN", [2, 4])), | |
| [False, True, True, True, False], | |
| ) | |
| def test_not_between(self): | |
| self.assertEqual( | |
| self._filter(_make_leaf("id", "NOT_BETWEEN", [2, 4])), | |
| [True, False, False, False, True], | |
| ) | |
| def test_starts_with(self): | |
| self.assertEqual( | |
| self._filter(_make_leaf("name", "STARTS_WITH", ["al"], "STRING")), | |
| [True, False, False, True, False], | |
| ) | |
| def test_ends_with(self): | |
| self.assertEqual( | |
| self._filter(_make_leaf("name", "ENDS_WITH", ["e"], "STRING")), | |
| [True, False, True, True, True], | |
| ) | |
| def test_contains(self): | |
| self.assertEqual( | |
| self._filter(_make_leaf("name", "CONTAINS", ["li"], "STRING")), | |
| [True, False, True, True, False], | |
| ) | |
| def test_like(self): | |
| like_json = json.dumps({ | |
| "kind": "LEAF", | |
| "transform": {"name": "FIELD_REF", "fieldRef": {"index": 0, "name": "name", "type": "STRING"}}, | |
| "function": "LIKE", | |
| "literals": ["%li%"], | |
| }) | |
| self.assertEqual( | |
| parse_predicate_to_batch_filter(like_json)(self.batch).to_pylist(), | |
| [True, False, True, True, False], | |
| ) | |
| def test_true(self): | |
| self.assertEqual( | |
| self._filter(_make_leaf("id", "TRUE")), | |
| [True, True, True, True, True], | |
| ) | |
| def test_false(self): | |
| self.assertEqual( | |
| self._filter(_make_leaf("id", "FALSE")), | |
| [False, False, False, False, False], | |
| ) | |
| class TestCompoundPredicates(unittest.TestCase): | |
| def setUp(self): | |
| self.batch = pa.RecordBatch.from_pydict({ | |
| "id": [1, 2, 3, 4, 5], | |
| }) | |
| def test_and(self): | |
| pred = json.dumps({ | |
| "kind": "COMPOUND", | |
| "function": "AND", | |
| "children": [ | |
| {"kind": "LEAF", | |
| "transform": {"name": "FIELD_REF", | |
| "fieldRef": {"index": 0, "name": "id", "type": "INT"}}, | |
| "function": "GREATER_THAN", "literals": [2]}, | |
| {"kind": "LEAF", | |
| "transform": {"name": "FIELD_REF", | |
| "fieldRef": {"index": 0, "name": "id", "type": "INT"}}, | |
| "function": "LESS_THAN", "literals": [5]}, | |
| ], | |
| }) | |
| f = parse_predicate_to_batch_filter(pred) | |
| self.assertEqual(f(self.batch).to_pylist(), [False, False, True, True, False]) | |
| def test_or(self): | |
| pred = json.dumps({ | |
| "kind": "COMPOUND", | |
| "function": "OR", | |
| "children": [ | |
| {"kind": "LEAF", | |
| "transform": {"name": "FIELD_REF", | |
| "fieldRef": {"index": 0, "name": "id", "type": "INT"}}, | |
| "function": "EQUAL", "literals": [1]}, | |
| {"kind": "LEAF", | |
| "transform": {"name": "FIELD_REF", | |
| "fieldRef": {"index": 0, "name": "id", "type": "INT"}}, | |
| "function": "EQUAL", "literals": [5]}, | |
| ], | |
| }) | |
| f = parse_predicate_to_batch_filter(pred) | |
| self.assertEqual(f(self.batch).to_pylist(), [True, False, False, False, True]) | |
| class TestPredicateTransforms(unittest.TestCase): | |
| def setUp(self): | |
| self.batch = pa.RecordBatch.from_pydict({ | |
| "id": [1, 2, 3, 4, 5], | |
| "name": ["alice", "bob", "charlie", "alice", "eve"], | |
| }) | |
| def test_upper_transform(self): | |
| pred = json.dumps({ | |
| "kind": "LEAF", | |
| "transform": {"name": "UPPER", "inputs": [{"index": 1, "name": "name", "type": "STRING"}]}, | |
| "function": "EQUAL", | |
| "literals": ["ALICE"], | |
| }) | |
| f = parse_predicate_to_batch_filter(pred) | |
| self.assertEqual(f(self.batch).to_pylist(), [True, False, False, True, False]) | |
| def test_lower_transform(self): | |
| pred = json.dumps({ | |
| "kind": "LEAF", | |
| "transform": {"name": "LOWER", "inputs": [{"index": 1, "name": "name", "type": "STRING"}]}, | |
| "function": "EQUAL", | |
| "literals": ["bob"], | |
| }) | |
| f = parse_predicate_to_batch_filter(pred) | |
| self.assertEqual(f(self.batch).to_pylist(), [False, True, False, False, False]) | |
| def test_cast_transform(self): | |
| pred = json.dumps({ | |
| "kind": "LEAF", | |
| "transform": {"name": "CAST", "fieldRef": {"index": 0, "name": "id", "type": "INT"}, "type": "BIGINT"}, | |
| "function": "GREATER_THAN", | |
| "literals": [3], | |
| }) | |
| f = parse_predicate_to_batch_filter(pred) | |
| self.assertEqual(f(self.batch).to_pylist(), [False, False, False, True, True]) | |
| def test_null_transform(self): | |
| pred = json.dumps({ | |
| "kind": "LEAF", | |
| "transform": {"name": "NULL"}, | |
| "function": "IS_NULL", | |
| }) | |
| f = parse_predicate_to_batch_filter(pred) | |
| self.assertEqual(f(self.batch).to_pylist(), [True, True, True, True, True]) | |
| def test_concat_transform_in_predicate(self): | |
| batch = pa.RecordBatch.from_pydict({ | |
| "first": ["john", "jane"], | |
| "last": ["doe", "smith"], | |
| }) | |
| pred = json.dumps({ | |
| "kind": "LEAF", | |
| "transform": { | |
| "name": "CONCAT", | |
| "inputs": [ | |
| {"index": 0, "name": "first", "type": "STRING"}, | |
| {"index": 1, "name": "last", "type": "STRING"}, | |
| ], | |
| }, | |
| "function": "EQUAL", | |
| "literals": ["johndoe"], | |
| }) | |
| f = parse_predicate_to_batch_filter(pred) | |
| self.assertEqual(f(batch).to_pylist(), [True, False]) | |
| def test_concat_ws_transform_in_predicate(self): | |
| batch = pa.RecordBatch.from_pydict({ | |
| "first": ["john", "jane"], | |
| "last": ["doe", "smith"], | |
| }) | |
| pred = json.dumps({ | |
| "kind": "LEAF", | |
| "transform": { | |
| "name": "CONCAT_WS", | |
| "inputs": [ | |
| " ", | |
| {"index": 0, "name": "first", "type": "STRING"}, | |
| {"index": 1, "name": "last", "type": "STRING"}, | |
| ], | |
| }, | |
| "function": "EQUAL", | |
| "literals": ["john doe"], | |
| }) | |
| f = parse_predicate_to_batch_filter(pred) | |
| self.assertEqual(f(batch).to_pylist(), [True, False]) | |
| class TestPaimonTypeToArrow(unittest.TestCase): | |
| def test_simple_types(self): | |
| self.assertEqual(_paimon_type_to_arrow("INT"), pa.int32()) | |
| self.assertEqual(_paimon_type_to_arrow("BIGINT"), pa.int64()) | |
| self.assertEqual(_paimon_type_to_arrow("SMALLINT"), pa.int16()) | |
| self.assertEqual(_paimon_type_to_arrow("TINYINT"), pa.int8()) | |
| self.assertEqual(_paimon_type_to_arrow("FLOAT"), pa.float32()) | |
| self.assertEqual(_paimon_type_to_arrow("DOUBLE"), pa.float64()) | |
| self.assertEqual(_paimon_type_to_arrow("STRING"), pa.string()) | |
| self.assertEqual(_paimon_type_to_arrow("BOOLEAN"), pa.bool_()) | |
| self.assertEqual(_paimon_type_to_arrow("BYTES"), pa.binary()) | |
| self.assertEqual(_paimon_type_to_arrow("DATE"), pa.date32()) | |
| def test_varchar_char(self): | |
| self.assertEqual(_paimon_type_to_arrow("VARCHAR"), pa.string()) | |
| self.assertEqual(_paimon_type_to_arrow("VARCHAR(100)"), pa.string()) | |
| self.assertEqual(_paimon_type_to_arrow("CHAR(10)"), pa.string()) | |
| def test_timestamp(self): | |
| self.assertEqual(_paimon_type_to_arrow("TIMESTAMP(3)"), pa.timestamp("ms")) | |
| self.assertEqual(_paimon_type_to_arrow("TIMESTAMP(6)"), pa.timestamp("us")) | |
| self.assertEqual(_paimon_type_to_arrow("TIMESTAMP(9)"), pa.timestamp("ns")) | |
| self.assertEqual(_paimon_type_to_arrow("TIMESTAMP(0)"), pa.timestamp("s")) | |
| def test_timestamp_ltz(self): | |
| self.assertEqual( | |
| _paimon_type_to_arrow("TIMESTAMP WITH LOCAL TIME ZONE(6)"), | |
| pa.timestamp("us", tz="UTC"), | |
| ) | |
| def test_decimal(self): | |
| self.assertEqual( | |
| _paimon_type_to_arrow("DECIMAL(10, 2)"), | |
| pa.decimal128(10, 2), | |
| ) | |
| def test_unsupported_type_raises(self): | |
| with self.assertRaises(ValueError): | |
| _paimon_type_to_arrow("ARRAY<INT>") | |
| class TestConvertLiteral(unittest.TestCase): | |
| def test_none_literal(self): | |
| self.assertIsNone(_convert_literal(None, pa.int32())) | |
| def test_int_passthrough(self): | |
| self.assertEqual(_convert_literal(42, pa.int32()), 42) | |
| def test_string_passthrough(self): | |
| self.assertEqual(_convert_literal("hello", pa.string()), "hello") | |
| def test_timestamp_literal(self): | |
| result = _convert_literal("2024-01-15T10:30:00", pa.timestamp("us")) | |
| self.assertIsInstance(result, pa.Scalar) | |
| def test_timestamp_z_suffix(self): | |
| result = _convert_literal("2024-01-15T10:30:00Z", pa.timestamp("us", tz="UTC")) | |
| self.assertIsInstance(result, pa.Scalar) | |
| def test_date_literal(self): | |
| result = _convert_literal("2024-01-15", pa.date32()) | |
| self.assertIsInstance(result, pa.Scalar) | |
| def test_decimal_literal(self): | |
| result = _convert_literal(123.45, pa.decimal128(10, 2)) | |
| self.assertIsInstance(result, pa.Scalar) | |
| def test_timestamp_ltz_literal_uses_epoch_seconds(self): | |
| # Jackson serializes java.time.Instant as epoch seconds (with an | |
| # optional fractional part for nanos), not epoch milliseconds. | |
| target_type = pa.timestamp("us", tz="UTC") | |
| result = _convert_literal(1720612345.123456, target_type) | |
| self.assertIsInstance(result, pa.Scalar) | |
| expected = pa.scalar( | |
| datetime.datetime.fromtimestamp(1720612345.123456, tz=datetime.timezone.utc), | |
| type=target_type, | |
| ) | |
| self.assertEqual(result.value, expected.value) | |
| def test_time_literal_with_nanos(self): | |
| # Jackson serializes java.time.LocalTime as [hour, minute, second, | |
| # nanoOfSecond]; the nanosecond component must survive as | |
| # sub-second precision on the resulting time scalar. | |
| result = _convert_literal([12, 30, 45, 123456000], pa.time64("us")) | |
| self.assertIsInstance(result, pa.Scalar) | |
| expected = pa.scalar( | |
| datetime.time(12, 30, 45, microsecond=123456), type=pa.time64("us")) | |
| self.assertEqual(result.value, expected.value) | |
| def test_time_literal_without_nanos(self): | |
| result = _convert_literal([12, 30, 45], pa.time64("us")) | |
| self.assertIsInstance(result, pa.Scalar) | |
| expected = pa.scalar(datetime.time(12, 30, 45), type=pa.time64("us")) | |
| self.assertEqual(result.value, expected.value) | |
| class TestExtractReferencedFields(unittest.TestCase): | |
| def test_leaf_field_ref(self): | |
| refs = extract_referenced_fields(json.dumps({ | |
| "kind": "LEAF", | |
| "transform": {"name": "FIELD_REF", "fieldRef": {"index": 0, "name": "col1", "type": "INT"}}, | |
| "function": "EQUAL", | |
| "literals": [1], | |
| })) | |
| self.assertEqual(refs, {"col1"}) | |
| def test_leaf_cast(self): | |
| refs = extract_referenced_fields(json.dumps({ | |
| "kind": "LEAF", | |
| "transform": {"name": "CAST", "fieldRef": {"index": 0, "name": "col1", "type": "INT"}, "type": "BIGINT"}, | |
| "function": "EQUAL", | |
| "literals": [1], | |
| })) | |
| self.assertEqual(refs, {"col1"}) | |
| def test_compound_collects_all(self): | |
| refs = extract_referenced_fields(json.dumps({ | |
| "kind": "COMPOUND", | |
| "function": "AND", | |
| "children": [ | |
| {"kind": "LEAF", | |
| "transform": {"name": "FIELD_REF", | |
| "fieldRef": {"index": 0, "name": "a", "type": "INT"}}, | |
| "function": "EQUAL", "literals": [1]}, | |
| {"kind": "LEAF", | |
| "transform": {"name": "UPPER", | |
| "inputs": [{"index": 1, "name": "b", "type": "STRING"}]}, | |
| "function": "EQUAL", "literals": ["X"]}, | |
| ], | |
| })) | |
| self.assertEqual(refs, {"a", "b"}) | |
| def test_concat_inputs(self): | |
| refs = extract_referenced_fields(json.dumps({ | |
| "kind": "LEAF", | |
| "transform": { | |
| "name": "CONCAT", | |
| "inputs": [ | |
| {"index": 0, "name": "first", "type": "STRING"}, | |
| "literal", | |
| {"index": 1, "name": "last", "type": "STRING"}, | |
| ], | |
| }, | |
| "function": "EQUAL", | |
| "literals": ["x"], | |
| })) | |
| self.assertEqual(refs, {"first", "last"}) | |
| class TestLikeEdgeCases(unittest.TestCase): | |
| def test_dot_is_literal(self): | |
| batch = pa.RecordBatch.from_pydict({"v": ["a.b", "axb", "a.bc"]}) | |
| f = parse_predicate_to_batch_filter(json.dumps({ | |
| "kind": "LEAF", | |
| "transform": {"name": "FIELD_REF", "fieldRef": {"index": 0, "name": "v", "type": "STRING"}}, | |
| "function": "LIKE", | |
| "literals": ["a.b"], | |
| })) | |
| self.assertEqual(f(batch).to_pylist(), [True, False, False]) | |
| def test_underscore_matches_single_char(self): | |
| batch = pa.RecordBatch.from_pydict({"v": ["abc", "axc", "ac", "abbc"]}) | |
| f = parse_predicate_to_batch_filter(json.dumps({ | |
| "kind": "LEAF", | |
| "transform": {"name": "FIELD_REF", "fieldRef": {"index": 0, "name": "v", "type": "STRING"}}, | |
| "function": "LIKE", | |
| "literals": ["a_c"], | |
| })) | |
| self.assertEqual(f(batch).to_pylist(), [True, True, False, False]) | |
| def test_percent_matches_any(self): | |
| batch = pa.RecordBatch.from_pydict({"v": ["abc", "ac", "axyzc", "def"]}) | |
| f = parse_predicate_to_batch_filter(json.dumps({ | |
| "kind": "LEAF", | |
| "transform": {"name": "FIELD_REF", "fieldRef": {"index": 0, "name": "v", "type": "STRING"}}, | |
| "function": "LIKE", | |
| "literals": ["a%c"], | |
| })) | |
| self.assertEqual(f(batch).to_pylist(), [True, True, True, False]) | |
| class TestLikeEscapeSemantics(unittest.TestCase): | |
| def _make_like_batch(self, values): | |
| return pa.RecordBatch.from_pydict( | |
| {"name": values}, schema=pa.schema([("name", pa.string())])) | |
| def _make_like_predicate(self, pattern): | |
| return json.dumps({ | |
| "kind": "LEAF", | |
| "transform": {"name": "FIELD_REF", "fieldRef": {"name": "name"}}, | |
| "function": "LIKE", | |
| "literals": [pattern] | |
| }) | |
| def test_like_simple_percent(self): | |
| batch = self._make_like_batch(["admin", "admin_foo", "user"]) | |
| fn = parse_predicate_to_batch_filter(self._make_like_predicate("admin%")) | |
| self.assertEqual(fn(batch).to_pylist(), [True, True, False]) | |
| def test_like_simple_underscore(self): | |
| batch = self._make_like_batch(["ab", "abc", "a"]) | |
| fn = parse_predicate_to_batch_filter(self._make_like_predicate("a_")) | |
| self.assertEqual(fn(batch).to_pylist(), [True, False, False]) | |
| def test_like_escaped_underscore(self): | |
| batch = self._make_like_batch(["admin_foo", "adminXfoo", "admin"]) | |
| fn = parse_predicate_to_batch_filter(self._make_like_predicate("admin\\_%")) | |
| self.assertEqual(fn(batch).to_pylist(), [True, False, False]) | |
| def test_like_escaped_percent(self): | |
| batch = self._make_like_batch(["100%", "100X", "100"]) | |
| fn = parse_predicate_to_batch_filter(self._make_like_predicate("100\\%")) | |
| self.assertEqual(fn(batch).to_pylist(), [True, False, False]) | |
| def test_like_escaped_backslash(self): | |
| batch = self._make_like_batch(["a\\b", "a/b", "ab"]) | |
| fn = parse_predicate_to_batch_filter(self._make_like_predicate("a\\\\b")) | |
| self.assertEqual(fn(batch).to_pylist(), [True, False, False]) | |
| def test_like_regex_special_chars(self): | |
| batch = self._make_like_batch(["a.b", "axb", "a..b"]) | |
| fn = parse_predicate_to_batch_filter(self._make_like_predicate("a.b")) | |
| self.assertEqual(fn(batch).to_pylist(), [True, False, False]) | |
| def test_like_percent_matches_newline(self): | |
| batch = self._make_like_batch(["hello\nworld", "hello"]) | |
| fn = parse_predicate_to_batch_filter(self._make_like_predicate("hello%")) | |
| self.assertEqual(fn(batch).to_pylist(), [True, True]) | |
| def test_like_invalid_escape_sequence(self): | |
| import pytest | |
| fn = parse_predicate_to_batch_filter(self._make_like_predicate("admin\\x")) | |
| batch = self._make_like_batch(["admin"]) | |
| with pytest.raises((RuntimeError, ValueError), match="Invalid escape sequence"): | |
| fn(batch) | |
| class TestNotInNullSemantics(unittest.TestCase): | |
| """Java NOT_IN: null field -> false; any null literal -> false for all.""" | |
| def test_not_in_null_field_returns_false(self): | |
| batch = pa.RecordBatch.from_pydict( | |
| {"dept": ["eng", None, "sales"]}, | |
| schema=pa.schema([("dept", pa.string())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter( | |
| _make_leaf("dept", "NOT_IN", ["blocked"], field_type="STRING")) | |
| result = fn(batch).to_pylist() | |
| # null dept must return false, not true | |
| self.assertEqual(result, [True, False, True]) | |
| def test_not_in_null_literal_returns_false_for_all(self): | |
| batch = pa.RecordBatch.from_pydict( | |
| {"dept": ["eng", None, "sales"]}, | |
| schema=pa.schema([("dept", pa.string())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter( | |
| _make_leaf("dept", "NOT_IN", ["blocked", None], | |
| field_type="STRING")) | |
| result = fn(batch).to_pylist() | |
| # any null literal makes NOT_IN false for all rows | |
| self.assertEqual(result, [False, False, False]) | |
| def test_not_in_normal(self): | |
| batch = pa.RecordBatch.from_pydict( | |
| {"id": [1, 2, 3]}, | |
| schema=pa.schema([("id", pa.int64())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter( | |
| _make_leaf("id", "NOT_IN", [2, 3])) | |
| result = fn(batch).to_pylist() | |
| self.assertEqual(result, [True, False, False]) | |
| class TestIsNaN(unittest.TestCase): | |
| def test_is_nan_float(self): | |
| batch = pa.RecordBatch.from_pydict( | |
| {"val": [1.0, float("nan"), 3.0, None]}, | |
| schema=pa.schema([("val", pa.float64())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter( | |
| _make_leaf("val", "IS_NAN", field_type="DOUBLE")) | |
| result = fn(batch).to_pylist() | |
| self.assertEqual(result, [False, True, False, False]) | |
| def test_is_nan_no_nans(self): | |
| batch = pa.RecordBatch.from_pydict( | |
| {"val": [1.0, 2.0]}, | |
| schema=pa.schema([("val", pa.float64())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter( | |
| _make_leaf("val", "IS_NAN", field_type="DOUBLE")) | |
| result = fn(batch).to_pylist() | |
| self.assertEqual(result, [False, False]) | |
| class TestNullLiteralDefense(unittest.TestCase): | |
| """Java LeafBinaryFunction/LeafTernaryFunction: null literal -> false.""" | |
| def test_starts_with_null_literal(self): | |
| batch = pa.RecordBatch.from_pydict( | |
| {"name": ["alice", "bob", None]}, | |
| schema=pa.schema([("name", pa.string())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter( | |
| _make_leaf("name", "STARTS_WITH", [None], field_type="STRING")) | |
| result = fn(batch).to_pylist() | |
| self.assertEqual(result, [False, False, False]) | |
| def test_ends_with_null_literal(self): | |
| batch = pa.RecordBatch.from_pydict( | |
| {"name": ["alice", "bob", None]}, | |
| schema=pa.schema([("name", pa.string())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter( | |
| _make_leaf("name", "ENDS_WITH", [None], field_type="STRING")) | |
| result = fn(batch).to_pylist() | |
| self.assertEqual(result, [False, False, False]) | |
| def test_contains_null_literal(self): | |
| batch = pa.RecordBatch.from_pydict( | |
| {"name": ["alice", "bob", None]}, | |
| schema=pa.schema([("name", pa.string())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter( | |
| _make_leaf("name", "CONTAINS", [None], field_type="STRING")) | |
| result = fn(batch).to_pylist() | |
| self.assertEqual(result, [False, False, False]) | |
| def test_like_null_literal(self): | |
| batch = pa.RecordBatch.from_pydict( | |
| {"name": ["alice", "bob", None]}, | |
| schema=pa.schema([("name", pa.string())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter( | |
| _make_leaf("name", "LIKE", [None], field_type="STRING")) | |
| result = fn(batch).to_pylist() | |
| self.assertEqual(result, [False, False, False]) | |
| def test_between_null_lower_literal(self): | |
| batch = pa.RecordBatch.from_pydict( | |
| {"val": [1, 5, 10, None]}, | |
| schema=pa.schema([("val", pa.int64())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter( | |
| _make_leaf("val", "BETWEEN", [None, 10])) | |
| result = fn(batch).to_pylist() | |
| self.assertEqual(result, [False, False, False, False]) | |
| def test_between_null_upper_literal(self): | |
| batch = pa.RecordBatch.from_pydict( | |
| {"val": [1, 5, 10, None]}, | |
| schema=pa.schema([("val", pa.int64())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter( | |
| _make_leaf("val", "BETWEEN", [1, None])) | |
| result = fn(batch).to_pylist() | |
| self.assertEqual(result, [False, False, False, False]) | |
| def test_not_between_null_literal(self): | |
| batch = pa.RecordBatch.from_pydict( | |
| {"val": [1, 5, 10, None]}, | |
| schema=pa.schema([("val", pa.int64())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter( | |
| _make_leaf("val", "NOT_BETWEEN", [None, 10])) | |
| result = fn(batch).to_pylist() | |
| self.assertEqual(result, [False, False, False, False]) | |
| def test_starts_with_normal(self): | |
| """Verify normal case still works after adding null guard.""" | |
| batch = pa.RecordBatch.from_pydict( | |
| {"name": ["alice", "anna", "bob", None]}, | |
| schema=pa.schema([("name", pa.string())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter( | |
| _make_leaf("name", "STARTS_WITH", ["a"], field_type="STRING")) | |
| result = fn(batch).to_pylist() | |
| self.assertEqual(result, [True, True, False, False]) | |
| def test_between_normal(self): | |
| """Verify normal BETWEEN still works.""" | |
| batch = pa.RecordBatch.from_pydict( | |
| {"val": [1, 5, 10, None]}, | |
| schema=pa.schema([("val", pa.int64())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter( | |
| _make_leaf("val", "BETWEEN", [3, 8])) | |
| result = fn(batch).to_pylist() | |
| self.assertEqual(result, [False, True, False, False]) | |
| if __name__ == "__main__": | |
| unittest.main() | |
| class TestInNullHandling(unittest.TestCase): | |
| def test_in_with_null_literal_drops_null(self): | |
| batch = pa.RecordBatch.from_pydict( | |
| {"dept": ["eng", "sales", None]}, | |
| schema=pa.schema([("dept", pa.string())]), | |
| ) | |
| pred = _make_leaf("dept", "IN", ["eng", None], field_type="STRING") | |
| result = parse_predicate_to_batch_filter(pred)(batch).to_pylist() | |
| self.assertEqual(result, [True, False, False]) | |
| def test_in_all_null_literals(self): | |
| batch = pa.RecordBatch.from_pydict( | |
| {"dept": ["eng", "sales"]}, | |
| schema=pa.schema([("dept", pa.string())]), | |
| ) | |
| pred = _make_leaf("dept", "IN", [None, None], field_type="STRING") | |
| result = parse_predicate_to_batch_filter(pred)(batch).to_pylist() | |
| self.assertEqual(result, [False, False]) | |
| def test_in_null_field_value_returns_false(self): | |
| batch = pa.RecordBatch.from_pydict( | |
| {"dept": [None, "eng"]}, | |
| schema=pa.schema([("dept", pa.string())]), | |
| ) | |
| pred = _make_leaf("dept", "IN", ["eng", "sales"], field_type="STRING") | |
| result = parse_predicate_to_batch_filter(pred)(batch).to_pylist() | |
| self.assertEqual(result, [False, True]) | |
| class TestComparisonNullLiteral(unittest.TestCase): | |
| def test_equal_null_literal_returns_false(self): | |
| batch = pa.RecordBatch.from_pydict({"id": [1, 2, 3]}) | |
| pred = _make_leaf("id", "EQUAL", [None]) | |
| result = parse_predicate_to_batch_filter(pred)(batch).to_pylist() | |
| self.assertEqual(result, [False, False, False]) | |
| def test_not_equal_null_literal_returns_false(self): | |
| batch = pa.RecordBatch.from_pydict({"id": [1, 2, 3]}) | |
| pred = _make_leaf("id", "NOT_EQUAL", [None]) | |
| result = parse_predicate_to_batch_filter(pred)(batch).to_pylist() | |
| self.assertEqual(result, [False, False, False]) | |
| def test_less_than_null_literal_returns_false(self): | |
| batch = pa.RecordBatch.from_pydict({"id": [1, 2, 3]}) | |
| pred = _make_leaf("id", "LESS_THAN", [None]) | |
| result = parse_predicate_to_batch_filter(pred)(batch).to_pylist() | |
| self.assertEqual(result, [False, False, False]) | |
| def test_greater_than_null_literal_returns_false(self): | |
| batch = pa.RecordBatch.from_pydict({"id": [1, 2, 3]}) | |
| pred = _make_leaf("id", "GREATER_THAN", [None]) | |
| result = parse_predicate_to_batch_filter(pred)(batch).to_pylist() | |
| self.assertEqual(result, [False, False, False]) | |
| class TestConcatWsAllNull(unittest.TestCase): | |
| def test_concat_ws_all_null_values_returns_empty(self): | |
| pred_json = json.dumps({ | |
| "kind": "LEAF", | |
| "transform": { | |
| "name": "CONCAT_WS", | |
| "inputs": [",", None, None], | |
| }, | |
| "function": "EQUAL", | |
| "literals": [""], | |
| }) | |
| batch = pa.RecordBatch.from_pydict( | |
| {"x": [1, 2, 3]}, | |
| schema=pa.schema([("x", pa.int64())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter(pred_json) | |
| result = fn(batch).to_pylist() | |
| self.assertEqual(result, [True, True, True]) | |
| def test_concat_ws_mixed_null_preserves_row_positions(self): | |
| pred_json = json.dumps({ | |
| "kind": "LEAF", | |
| "transform": { | |
| "name": "CONCAT_WS", | |
| "inputs": [ | |
| "-", | |
| {"index": 0, "name": "a", "type": "STRING"}, | |
| {"index": 1, "name": "b", "type": "STRING"}, | |
| ], | |
| }, | |
| "function": "EQUAL", | |
| "literals": ["x-y"], | |
| }) | |
| batch = pa.RecordBatch.from_pydict( | |
| {"a": [None, "x", None, "p"], "b": [None, "y", "z", None]}, | |
| schema=pa.schema([("a", pa.string()), ("b", pa.string())]), | |
| ) | |
| fn = parse_predicate_to_batch_filter(pred_json) | |
| result = fn(batch).to_pylist() | |
| self.assertEqual(result, [False, True, False, False]) | |
| class TestDateExtractTransforms(unittest.TestCase): | |
| """YEAR / MONTH / DAY / HOUR / MINUTE / SECOND / QUARTER / DAY_OF_YEAR / | |
| WEEKDAY / ISO_DAY_OF_WEEK / DAY_OF_WEEK / WEEK / YEAR_OF_WEEK, | |
| mirroring Java's DateExtractTransform subclasses.""" | |
| def _ts_batch(self): | |
| # 2024-01-01 00:00:00 and 2024-07-01 02:24:05 (us), plus null. | |
| return pa.RecordBatch.from_pydict({ | |
| "t": pa.array([1704067200000000, 1719800645000000, None], | |
| type=pa.timestamp("us")), | |
| }) | |
| def _apply(self, name, batch, field="t", ftype="TIMESTAMP(6)"): | |
| transform = {"name": name, | |
| "fieldRef": {"index": 0, "name": field, "type": ftype}} | |
| return _apply_predicate_transform(transform, batch).to_pylist() | |
| def test_timestamp_parts(self): | |
| batch = self._ts_batch() | |
| self.assertEqual(self._apply("YEAR", batch), [2024, 2024, None]) | |
| self.assertEqual(self._apply("MONTH", batch), [1, 7, None]) | |
| self.assertEqual(self._apply("DAY", batch), [1, 1, None]) | |
| self.assertEqual(self._apply("HOUR", batch), [0, 2, None]) | |
| self.assertEqual(self._apply("MINUTE", batch), [0, 24, None]) | |
| self.assertEqual(self._apply("SECOND", batch), [0, 5, None]) | |
| self.assertEqual(self._apply("QUARTER", batch), [1, 3, None]) | |
| self.assertEqual(self._apply("DAY_OF_YEAR", batch), [1, 183, None]) | |
| def test_week_family_parts(self): | |
| # 2024-01-01 Monday, 2023-01-01 Sunday, 2021-01-01 Friday (ISO week 53 | |
| # of the 2020 week-based year), plus null. Numbering matches java.time. | |
| batch = pa.RecordBatch.from_pydict({ | |
| "d": pa.array( | |
| [datetime.date(2024, 1, 1), datetime.date(2023, 1, 1), | |
| datetime.date(2021, 1, 1), None], | |
| type=pa.date32()), | |
| }) | |
| # WEEKDAY Monday=0..Sunday=6 | |
| self.assertEqual(self._apply("WEEKDAY", batch, "d", "DATE"), [0, 6, 4, None]) | |
| # ISO_DAY_OF_WEEK Monday=1..Sunday=7 | |
| self.assertEqual(self._apply("ISO_DAY_OF_WEEK", batch, "d", "DATE"), [1, 7, 5, None]) | |
| # DAY_OF_WEEK Sunday=1..Saturday=7 | |
| self.assertEqual(self._apply("DAY_OF_WEEK", batch, "d", "DATE"), [2, 1, 6, None]) | |
| # WEEK: ISO week of the week-based year | |
| self.assertEqual(self._apply("WEEK", batch, "d", "DATE"), [1, 52, 53, None]) | |
| # YEAR_OF_WEEK: ISO week-based year | |
| self.assertEqual(self._apply("YEAR_OF_WEEK", batch, "d", "DATE"), [2024, 2022, 2020, None]) | |
| def test_date_field_starts_at_midnight(self): | |
| # epoch day 19723 == 2024-01-01; a DATE has no time, so HOUR is 0. | |
| batch = pa.RecordBatch.from_pydict( | |
| {"d": pa.array([19723, None], type=pa.date32())}) | |
| self.assertEqual(self._apply("YEAR", batch, "d", "DATE"), [2024, None]) | |
| self.assertEqual(self._apply("MONTH", batch, "d", "DATE"), [1, None]) | |
| self.assertEqual(self._apply("HOUR", batch, "d", "DATE"), [0, None]) | |
| def test_wired_through_a_leaf_filter(self): | |
| pred = json.dumps({ | |
| "kind": "LEAF", | |
| "transform": {"name": "MONTH", | |
| "fieldRef": {"index": 0, "name": "t", "type": "TIMESTAMP(6)"}}, | |
| "function": "EQUAL", | |
| "literals": [7], | |
| }) | |
| self.assertEqual( | |
| parse_predicate_to_batch_filter(pred)(self._ts_batch()).to_pylist(), | |
| [False, True, False]) | |
| class TestStringTransforms(unittest.TestCase): | |
| """LENGTH / BIT_LENGTH / TRANSLATE / OVERLAY / PAD, mirroring the Java | |
| transforms of the same name (see TransformJsonSerdeTest for the wire shape).""" | |
| def _apply(self, transform, column): | |
| batch = pa.RecordBatch.from_pydict({"s": column}) | |
| return _apply_predicate_transform(transform, batch).to_pylist() | |
| def _field(self): | |
| return {"index": 0, "name": "s", "type": "STRING"} | |
| def test_length(self): | |
| out = self._apply({"name": "LENGTH", "inputs": [self._field()]}, | |
| ["hello", "hi", None]) | |
| self.assertEqual(out, [5, 2, None]) | |
| def test_bit_length(self): | |
| # "é" is two UTF-8 bytes -> 16 bits. | |
| out = self._apply({"name": "BIT_LENGTH", "inputs": [self._field()]}, | |
| ["hello", "é", None]) | |
| self.assertEqual(out, [40, 16, None]) | |
| def test_translate_maps_and_deletes(self): | |
| mapped = self._apply( | |
| {"name": "TRANSLATE", "inputs": [self._field(), "el", "ip"]}, | |
| ["hello", None]) | |
| self.assertEqual(mapped, ["hippo", None]) | |
| # A shorter replacement deletes the unmatched matching characters. | |
| deleted = self._apply( | |
| {"name": "TRANSLATE", "inputs": [self._field(), "l", ""]}, | |
| ["hello"]) | |
| self.assertEqual(deleted, ["heo"]) | |
| def test_overlay_with_and_without_length(self): | |
| four = self._apply( | |
| {"name": "OVERLAY", "inputs": [self._field(), "XX", 2, 1]}, | |
| ["hello", None]) | |
| self.assertEqual(four, ["hXXllo", None]) | |
| three = self._apply( | |
| {"name": "OVERLAY", "inputs": [self._field(), "XX", 2]}, | |
| ["hello"]) | |
| self.assertEqual(three, ["hXXlo"]) | |
| def test_overlay_position_arithmetic_overflows_like_java(self): | |
| # pos + replaced overflows 32-bit int in Java's OverlayTransform: | |
| # OVERLAY('hello', 'x', 1, 2147483647) is 'xhell' there (the wrapped | |
| # end position lands inside the string), not 'x'. The auth rule must | |
| # admit the same rows in PyPaimon. | |
| out = self._apply( | |
| {"name": "OVERLAY", "inputs": [self._field(), "x", 1, 2147483647]}, | |
| ["hello"]) | |
| self.assertEqual(out, ["xhell"]) | |
| def test_pad_left_right_and_truncate(self): | |
| left = self._apply( | |
| {"name": "PAD", "inputs": [self._field(), 5, "ab"], "direction": "LEFT"}, | |
| ["hi", None]) | |
| self.assertEqual(left, ["abahi", None]) | |
| right = self._apply( | |
| {"name": "PAD", "inputs": [self._field(), 5, "ab"], "direction": "RIGHT"}, | |
| ["hi"]) | |
| self.assertEqual(right, ["hiaba"]) | |
| # A target shorter than the source truncates. | |
| self.assertEqual( | |
| self._apply( | |
| {"name": "PAD", "inputs": [self._field(), 3, "x"], "direction": "LEFT"}, | |
| ["hello"]), | |
| ["hel"]) | |
| def test_length_wired_through_a_leaf_filter(self): | |
| pred = json.dumps({ | |
| "kind": "LEAF", | |
| "transform": {"name": "LENGTH", "inputs": [self._field()]}, | |
| "function": "GREATER_THAN", | |
| "literals": [3], | |
| }) | |
| batch = pa.RecordBatch.from_pydict({"s": ["hello", "hi", "abcd"]}) | |
| self.assertEqual( | |
| parse_predicate_to_batch_filter(pred)(batch).to_pylist(), | |
| [True, False, True]) | |
| def test_unknown_transform_still_raises(self): | |
| with self.assertRaises(ValueError): | |
| self._apply({"name": "NO_SUCH_TRANSFORM", "inputs": [self._field()]}, | |
| ["hello"]) |