| ################################################################################ | |
| # 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 | |
| from pathlib import Path | |
| import pyarrow as pa | |
| from pypaimon.common.predicate_json_parser import extract_referenced_fields | |
| from pypaimon.common.arrow_predicate import ( | |
| _apply_predicate_transform, | |
| _convert_literal, | |
| _paimon_type_to_arrow, | |
| 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_short_circuit_only_evaluates_rows_reaching_each_child(self): | |
| batch = pa.RecordBatch.from_pydict({'id': [1, 2], 's': ['abc', 'xyz'], 'position': [1, 2147483648]}) | |
| first = json.loads(_make_leaf('id', 'EQUAL', [1])) | |
| second = {'kind': 'LEAF', 'function': 'EQUAL', 'literals': ['a'], 'transform': { | |
| 'name': 'SUBSTRING', 'inputs': [ | |
| {'name': 's', 'type': 'STRING'}, {'name': 'position', 'type': 'BIGINT'}, 1]}} | |
| for function, compare, expected in [('AND', 1, [True, False]), ('OR', 2, [True, True])]: | |
| first['literals'] = [compare] | |
| rule = {'kind': 'COMPOUND', 'function': function, 'children': [first, second]} | |
| with self.subTest(function=function): | |
| self.assertEqual(parse_predicate_to_batch_filter(json.dumps(rule))(batch).to_pylist(), expected) | |
| 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_nested_compounds_keep_original_mixed_row_positions(self): | |
| batch = pa.RecordBatch.from_pydict({ | |
| 'id': [1, 2, 3, 4], 's': ['x', 'abc', '', None], 'position': [2147483648, 1, 0, 2147483648]}) | |
| def guard(id_): | |
| return {'kind': 'LEAF', 'function': 'EQUAL', 'literals': [id_], | |
| 'transform': {'name': 'FIELD_REF', 'fieldRef': {'name': 'id', 'type': 'BIGINT'}}} | |
| substring = {'kind': 'LEAF', 'function': 'EQUAL', 'literals': ['a'], 'transform': { | |
| 'name': 'SUBSTRING', 'inputs': [ | |
| {'name': 's', 'type': 'STRING'}, {'name': 'position', 'type': 'BIGINT'}, 1]}} | |
| rule = {'kind': 'COMPOUND', 'function': 'OR', 'children': [guard(1), { | |
| 'kind': 'COMPOUND', 'function': 'AND', 'children': [guard(2), substring]}]} | |
| self.assertEqual(parse_predicate_to_batch_filter(json.dumps(rule))(batch).to_pylist(), | |
| [True, True, False, 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]) | |
| def test_unused_literals_and_unreachable_children_still_validate(self): | |
| batch = pa.RecordBatch.from_pydict({'s': ['abc', None]}) | |
| for function in ['TRUE', 'FALSE', 'IS_NULL', 'IS_NOT_NULL', 'EQUAL', 'NOT_EQUAL', 'IN']: | |
| for literals in [['bad'], [None, 'bad']]: | |
| rule = {'kind': 'LEAF', 'transform': {'name': 'LENGTH', 'inputs': [{'name': 's', 'type': 'STRING'}]}, | |
| 'function': function, 'literals': literals} | |
| unreachable = {'kind': 'COMPOUND', 'function': 'OR', 'children': [ | |
| {'kind': 'LEAF', 'transform': {'name': 'NULL'}, 'function': 'TRUE', 'literals': []}, rule]} | |
| for definition in [rule, unreachable]: | |
| with self.subTest(function=function, literals=literals), self.assertRaises((ValueError, TypeError)): | |
| parse_predicate_to_batch_filter(json.dumps(definition))(batch) | |
| def test_unary_and_constant_literals_can_be_null_or_omitted(self): | |
| batch = pa.RecordBatch.from_pydict({'s': ['abc', None]}) | |
| for function, expected in [('TRUE', [True, True]), ('FALSE', [False, False]), | |
| ('IS_NULL', [False, True]), ('IS_NOT_NULL', [True, False])]: | |
| rule = {'kind': 'LEAF', 'transform': {'name': 'LENGTH', 'inputs': [{'name': 's', 'type': 'STRING'}]}, | |
| 'function': function, 'literals': None} | |
| self.assertEqual(parse_predicate_to_batch_filter(json.dumps(rule))(batch).to_pylist(), expected) | |
| rule.pop('literals') | |
| self.assertEqual(parse_predicate_to_batch_filter(json.dumps(rule))(batch).to_pylist(), expected) | |
| def test_invalid_compound_never_becomes_allow_all(self): | |
| for function, children in [('AND', []), ('OR', []), ('NOT', [])]: | |
| with self.assertRaises(ValueError): | |
| parse_predicate_to_batch_filter(json.dumps({'kind': 'COMPOUND', 'function': function, | |
| 'children': children})) | |
| def test_binary_predicate_literals_follow_java_rejection(self): | |
| batch = pa.RecordBatch.from_pydict({'s': ['abc', None]}) | |
| for type_ in ['BINARY(3)', 'VARBINARY(3)']: | |
| transform = {'name': 'CAST', 'fieldRef': {'name': 's', 'type': 'STRING'}, 'type': type_} | |
| rule = {'kind': 'LEAF', 'transform': transform, 'function': 'IS_NULL', 'literals': []} | |
| self.assertEqual(parse_predicate_to_batch_filter(json.dumps(rule))(batch).to_pylist(), [False, True]) | |
| rule.update(function='EQUAL', literals=[None]) | |
| self.assertEqual(parse_predicate_to_batch_filter(json.dumps(rule))(batch).to_pylist(), [False, False]) | |
| for function in ['EQUAL', 'TRUE', 'IS_NULL']: | |
| rule.update(function=function, literals=['YQAA']) | |
| with self.subTest(type=type_, function=function), self.assertRaises(ValueError): | |
| parse_predicate_to_batch_filter(json.dumps(rule))(batch) | |
| 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]) | |
| def test_cast_enforces_java_character_and_byte_widths(self): | |
| batch = pa.RecordBatch.from_pydict({'s': ['🙂中', 'a', None]}) | |
| for type_, expected in [ | |
| ('VARCHAR(1)', ['🙂', 'a', None]), ('CHAR(3)', ['🙂中 ', 'a ', None]), | |
| ('BINARY(6)', ['🙂中'.encode('utf8')[:6], b'a' + b'\x00' * 5, None]), | |
| ('VARBINARY(1)', [b'\xf0', b'a', None]), | |
| ('VARCHAR', ['🙂', 'a', None]), ('CHAR', ['🙂', 'a', None]), | |
| ('BINARY', [b'\xf0', b'a', None]), ('VARBINARY', [b'\xf0', b'a', None]), | |
| ('CHAR(3) NULL', ['🙂中 ', 'a ', None]), ('CHAR(3)NOT NULL', ['🙂中 ', 'a ', None]), | |
| ]: | |
| transform = {'name': 'CAST', 'fieldRef': {'name': 's', 'type': 'STRING'}, 'type': type_} | |
| with self.subTest(type=type_): | |
| self.assertEqual(_apply_predicate_transform(transform, batch).to_pylist(), expected) | |
| def test_rebound_cast_keeps_original_java_runtime_class(self): | |
| transform = {'name': 'CAST', 'fieldRef': {'name': 's', 'type': 'STRING'}, 'type': 'VARCHAR(5)'} | |
| with self.assertRaises(ValueError): | |
| _apply_predicate_transform(transform, pa.RecordBatch.from_pydict({'s': [42]})) | |
| nulls = pa.RecordBatch.from_pydict({'s': pa.array([None], type=pa.int32())}) | |
| self.assertEqual(_apply_predicate_transform(transform, nulls).to_pylist(), [None]) | |
| # NumericToString accepts Object, including a rebound string column. | |
| transform['fieldRef']['type'] = 'INT' | |
| strings = pa.RecordBatch.from_pydict({'s': ['abc', None]}) | |
| self.assertEqual(_apply_predicate_transform(transform, strings).to_pylist(), ['abc', None]) | |
| def test_rebound_numeric_to_string_uses_java_internal_temporal_values(self): | |
| transform = {'name': 'CAST', 'fieldRef': {'name': 's', 'type': 'INT'}, 'type': 'STRING'} | |
| for type_, values, expected in [ | |
| (pa.date32(), [0, -1, None], ['0', '-1', None]), | |
| (pa.time32('ms'), [0, 3600001, None], ['0', '3600001', None]), | |
| (pa.time64('us'), [0, 3600001123, None], ['0', '3600001', None]), | |
| ]: | |
| batch = pa.RecordBatch.from_pydict({'s': pa.array(values, type=type_)}) | |
| with self.subTest(type=type_): | |
| self.assertEqual(_apply_predicate_transform(transform, batch).to_pylist(), expected) | |
| def test_rebound_numeric_to_string_uses_decimal_plain_text(self): | |
| from decimal import Decimal | |
| transform = {'name': 'CAST', 'fieldRef': {'name': 's', 'type': 'INT'}, 'type': 'STRING'} | |
| values = ['0.0000000123', '0.0000000000', '-0.0000000123', None] | |
| array = pa.array([Decimal(v) if v is not None else None for v in values], type=pa.decimal128(38, 10)) | |
| batch = pa.RecordBatch.from_pydict({'s': array}) | |
| self.assertEqual(_apply_predicate_transform(transform, batch).to_pylist(), values) | |
| transform['type'] = 'VARCHAR(3)' | |
| self.assertEqual(_apply_predicate_transform(transform, batch).to_pylist(), ['0.0', '0.0', '-0.', None]) | |
| def test_rebound_numeric_to_string_refuses_other_object_formatting(self): | |
| transform = {'name': 'CAST', 'fieldRef': {'name': 's', 'type': 'INT'}, 'type': 'STRING'} | |
| for array in [pa.array([0, None], type=pa.timestamp('ms')), pa.array([b'a', None]), | |
| pa.array([1.0, None], type=pa.float64())]: | |
| with self.subTest(type=array.type), self.assertRaises(ValueError): | |
| _apply_predicate_transform(transform, pa.RecordBatch.from_pydict({'s': array})) | |
| nulls = pa.RecordBatch.from_pydict({'s': pa.nulls(2, type=array.type)}) | |
| self.assertEqual(_apply_predicate_transform(transform, nulls).to_pylist(), [None, None]) | |
| def test_empty_concat_is_the_empty_string(self): | |
| self.assertEqual(_apply_predicate_transform({'name': 'CONCAT', 'inputs': []}, self.batch).to_pylist(), | |
| [''] * len(self.batch)) | |
| 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_calendar_parts_match_java_golden_for_every_storage_unit(self): | |
| parts = ['YEAR', 'MONTH', 'DAY', 'HOUR', 'MINUTE', 'SECOND', 'QUARTER', 'DAY_OF_YEAR', | |
| 'WEEKDAY', 'ISO_DAY_OF_WEEK', 'DAY_OF_WEEK', 'WEEK', 'YEAR_OF_WEEK'] | |
| golden = Path(__file__).parent / 'resources' / 'calendar_extract_java.jsonl' | |
| for line in golden.read_text().splitlines(): | |
| case = json.loads(line) | |
| unit = case['unit'] | |
| type_ = pa.date32() if unit == 'date' else pa.timestamp(unit) | |
| wire = 'DATE' if unit == 'date' else 'TIMESTAMP(9)' | |
| batch = pa.RecordBatch.from_pydict({'t': pa.array([case['value'], None], type=type_)}) | |
| for part, expected in zip(parts, case['parts']): | |
| with self.subTest(unit=unit, value=case['value'], part=part): | |
| self.assertEqual(self._apply(part, batch, 't', wire), [expected, None]) | |
| def test_empty_and_all_null_calendar_batches(self): | |
| for unit in ['date', 's', 'ms', 'us', 'ns']: | |
| type_ = pa.date32() if unit == 'date' else pa.timestamp(unit) | |
| wire = 'DATE' if unit == 'date' else 'TIMESTAMP(9)' | |
| for values in [[], [None, None]]: | |
| batch = pa.RecordBatch.from_pydict({'t': pa.array(values, type=type_)}) | |
| for part in ['YEAR', 'MONTH', 'DAY', 'HOUR', 'MINUTE', 'SECOND', 'QUARTER', 'DAY_OF_YEAR', | |
| 'WEEKDAY', 'ISO_DAY_OF_WEEK', 'DAY_OF_WEEK', 'WEEK', 'YEAR_OF_WEEK']: | |
| with self.subTest(unit=unit, values=values, part=part): | |
| self.assertEqual(self._apply(part, batch, 't', wire), values) | |
| def test_calendar_operand_is_a_dependency_even_if_hidden(self): | |
| for part in ['YEAR', 'MONTH', 'DAY', 'HOUR', 'MINUTE', 'SECOND', 'QUARTER', 'DAY_OF_YEAR', | |
| 'WEEKDAY', 'ISO_DAY_OF_WEEK', 'DAY_OF_WEEK', 'WEEK', 'YEAR_OF_WEEK']: | |
| rule = {'name': part, 'fieldRef': {'name': 'hidden', 'index': 99, 'type': 'DATE'}} | |
| predicate = {'kind': 'LEAF', 'function': 'IS_NULL', 'transform': rule} | |
| self.assertEqual(extract_referenced_fields(json.dumps(predicate)), {'hidden'}) | |
| def test_stored_and_rebound_calendar_types_are_validated(self): | |
| for stored in ['INT', 'BIGINT', 'TIME(3)', 'TIMESTAMP_LTZ(3)', 'TIMESTAMP(10)', 'DATENULL', | |
| 'TIMESTAMP(٣)', 'TIMESTAMP(３)', | |
| 'TIMESTAMP(3) WITHOUT TIME ZONENULL', | |
| 'TIMESTAMP(3) WITHOUT TIME ZONENOT NULL', | |
| {'type': 'DATE'}, {'type': 'TIMESTAMP(6)'}, None]: | |
| with self.subTest(stored=stored), self.assertRaises(ValueError): | |
| self._apply('YEAR', self._ts_batch(), 't', stored) | |
| for type_, value in [(pa.int32(), 1), (pa.time32('ms'), 1), (pa.timestamp('ms', tz='UTC'), 1)]: | |
| batch = pa.RecordBatch.from_pydict({'t': pa.array([value], type=type_)}) | |
| with self.subTest(type=type_), self.assertRaises(ValueError): | |
| self._apply('YEAR', batch, 't', 'DATE') | |
| # As in Java copyWithNewInputs, DATE and local TIMESTAMP may rebind to | |
| # one another; the actual column decides the executed calendar value. | |
| for stored in ['DATE', 'DATE NULL', 'TIMESTAMP(3)NOT NULL', 'TIMESTAMP WITHOUT TIME ZONE', | |
| 'TIMESTAMP(3) WITHOUT TIME ZONE', | |
| 'TIMESTAMP(3)WITHOUT TIME ZONE']: | |
| with self.subTest(stored=stored): | |
| self.assertEqual(self._apply('YEAR', self._ts_batch(), 't', stored), [2024, 2024, None]) | |
| def test_full_date32_domain_does_not_require_a_timestamp_cast(self): | |
| batch = pa.RecordBatch.from_pydict({"d": pa.array([2147483647, -2147483648, None], type=pa.date32())}) | |
| self.assertEqual(self._apply("HOUR", batch, "d", "DATE"), [0, 0, None]) | |
| def test_timestamp_epoch_day_narrows_as_java_int_before_calendar_extraction(self): | |
| # Java Timestamp narrows its millisecond epoch day to a signed int. | |
| ticks = (2147483648 * 86400000) | |
| batch = pa.RecordBatch.from_pydict({"t": pa.array([ticks, -1, None], type=pa.timestamp("ms"))}) | |
| days = pa.RecordBatch.from_pydict({"d": pa.array([-2147483648, -1, None], type=pa.date32())}) | |
| self.assertEqual(self._apply("YEAR", batch, "t", "TIMESTAMP(3)"), | |
| self._apply("YEAR", days, "d", "DATE")) | |
| 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_integer_overflow_is_eager_even_with_null_string_inputs(self): | |
| for rule in [ | |
| {"name": "PAD", "direction": "LEFT", "inputs": [None, 2147483648, "x"]}, | |
| {"name": "OVERLAY", "inputs": [None, "x", 2147483648, None]}, | |
| {"name": "OVERLAY", "inputs": [None, "x", None, 2147483648]}, | |
| ]: | |
| with self.subTest(rule=rule), self.assertRaises(ValueError): | |
| self._apply(rule, [None]) | |
| def test_field_types_are_checked_before_rebinding(self): | |
| batch = pa.RecordBatch.from_pydict({"s": ["abc"], "n": [2]}) | |
| for name, inputs in [ | |
| ("PAD", [self._field(), {"name": "n", "type": "STRING"}, "x"]), | |
| ("OVERLAY", [self._field(), "x", {"name": "n", "type": "STRING"}]), | |
| ("TRANSLATE", [{"name": "s", "type": "INT"}, "a", "b"]), | |
| ("LENGTH", [{"name": "s", "type": "INT"}]), | |
| ("BIT_LENGTH", [{"name": "s", "type": "INT"}]), | |
| ]: | |
| rule = {"name": name, "inputs": inputs, "direction": "LEFT"} | |
| with self.subTest(rule=rule), self.assertRaises(ValueError): | |
| _apply_predicate_transform(rule, batch) | |
| for value in ["2", True, 1.5]: | |
| with self.subTest(value=value), self.assertRaises(ValueError): | |
| self._apply({"name": "PAD", "inputs": [None, value, "x"], "direction": "LEFT"}, [None]) | |
| 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"]) |