| ################################################################################ |
| # 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 ( |
| _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]) |
