| ################################################################################ |
| # 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 json |
| import re |
| from typing import Callable |
| |
| import pyarrow as pa |
| import pyarrow.compute as pc |
| |
| |
| def parse_predicate_to_batch_filter(json_str: str) -> Callable[[pa.RecordBatch], pa.Array]: |
| data = json.loads(json_str) |
| return _build_filter(data) |
| |
| |
| def _build_filter(data: dict) -> Callable[[pa.RecordBatch], pa.Array]: |
| kind = data["kind"] |
| if kind == "LEAF": |
| return _build_leaf_filter(data) |
| elif kind == "COMPOUND": |
| return _build_compound_filter(data) |
| raise ValueError(f"Unknown predicate kind: {kind}") |
| |
| |
| def _build_leaf_filter(data: dict) -> Callable: |
| transform = data["transform"] |
| function = data["function"] |
| literals = data.get("literals", []) |
| |
| def filter_fn(batch: pa.RecordBatch) -> pa.Array: |
| value_array = _apply_predicate_transform(transform, batch) |
| return _apply_leaf_function(function, value_array, literals, len(batch)) |
| |
| return filter_fn |
| |
| |
| def _build_compound_filter(data: dict) -> Callable: |
| function = data["function"] |
| child_filters = [_build_filter(child) for child in data["children"]] |
| |
| def filter_fn(batch: pa.RecordBatch) -> pa.Array: |
| if function == "AND": |
| result = child_filters[0](batch) |
| for cf in child_filters[1:]: |
| result = pc.and_(result, cf(batch)) |
| return result |
| elif function == "OR": |
| result = child_filters[0](batch) |
| for cf in child_filters[1:]: |
| result = pc.or_(result, cf(batch)) |
| return result |
| raise ValueError(f"Unknown compound function: {function}") |
| |
| return filter_fn |
| |
| |
| def _apply_predicate_transform(transform: dict, batch: pa.RecordBatch, |
| null_type: pa.DataType = pa.bool_()) -> pa.Array: |
| name = transform["name"] |
| |
| if name == "FIELD_REF": |
| return batch.column(transform["fieldRef"]["name"]) |
| |
| elif name == "CAST": |
| col = batch.column(transform["fieldRef"]["name"]) |
| target_type = _paimon_type_to_arrow(transform["type"]) |
| return pc.cast(col, target_type, safe=False) |
| |
| elif name == "UPPER": |
| input_col = _resolve_transform_input(transform["inputs"][0], batch) |
| return pc.utf8_upper(input_col) |
| |
| elif name == "LOWER": |
| input_col = _resolve_transform_input(transform["inputs"][0], batch) |
| return pc.utf8_lower(input_col) |
| |
| elif name == "CONCAT": |
| resolved = [_resolve_transform_input(inp, batch) for inp in transform["inputs"]] |
| if not resolved: |
| return pa.nulls(len(batch), type=pa.string()) |
| return pc.binary_join_element_wise(*resolved, "") |
| |
| elif name == "CONCAT_WS": |
| sep = _resolve_transform_input(transform["inputs"][0], batch) |
| values = [_resolve_transform_input(inp, batch) for inp in transform["inputs"][1:]] |
| if not values: |
| return pa.nulls(len(batch), type=pa.string()) |
| return _concat_ws(sep, values) |
| |
| elif name == "NULL": |
| return pa.nulls(len(batch), type=null_type) |
| |
| raise ValueError(f"Unknown transform type: {name}") |
| |
| |
| def _resolve_transform_input(inp, batch: pa.RecordBatch) -> pa.Array: |
| if isinstance(inp, dict): |
| return batch.column(inp["name"]) |
| elif isinstance(inp, str): |
| return pa.array([inp] * len(batch), type=pa.string()) |
| elif inp is None: |
| return pa.nulls(len(batch), type=pa.string()) |
| return pa.array([str(inp)] * len(batch), type=pa.string()) |
| |
| |
| def _concat_ws(sep: pa.Array, value_arrays: list) -> pa.Array: |
| sep_list = sep.to_pylist() |
| val_lists = [v.to_pylist() for v in value_arrays] |
| results = [] |
| for i in range(len(sep)): |
| s = sep_list[i] |
| if s is None: |
| results.append(None) |
| continue |
| parts = [vl[i] for vl in val_lists if vl[i] is not None] |
| results.append(s.join(parts)) |
| return pa.array(results, type=pa.string()) |
| |
| |
| def _null_as_false(arr: pa.Array) -> pa.Array: |
| """Replace nulls with False to match Java two-valued predicate semantics.""" |
| if arr.null_count == 0: |
| return arr |
| return pc.if_else(pc.is_valid(arr), arr, False) |
| |
| |
| def _apply_leaf_function(function: str, value_array: pa.Array, literals: list, batch_len: int) -> pa.Array: |
| """Null literal yields False to match Java LeafBinaryFunction/LeafTernaryFunction semantics. |
| |
| All comparison results are coerced from three-valued (PyArrow null) to |
| two-valued logic (null → False) so that compound AND/OR behaves |
| identically to Java. |
| """ |
| converted = [_convert_literal(lit, value_array.type) for lit in literals] |
| |
| if function == "EQUAL": |
| if converted[0] is None: |
| return pa.array([False] * batch_len, type=pa.bool_()) |
| return _null_as_false(pc.equal(value_array, converted[0])) |
| elif function == "NOT_EQUAL": |
| if converted[0] is None: |
| return pa.array([False] * batch_len, type=pa.bool_()) |
| return _null_as_false(pc.not_equal(value_array, converted[0])) |
| elif function == "LESS_THAN": |
| if converted[0] is None: |
| return pa.array([False] * batch_len, type=pa.bool_()) |
| return _null_as_false(pc.less(value_array, converted[0])) |
| elif function == "LESS_OR_EQUAL": |
| if converted[0] is None: |
| return pa.array([False] * batch_len, type=pa.bool_()) |
| return _null_as_false(pc.less_equal(value_array, converted[0])) |
| elif function == "GREATER_THAN": |
| if converted[0] is None: |
| return pa.array([False] * batch_len, type=pa.bool_()) |
| return _null_as_false(pc.greater(value_array, converted[0])) |
| elif function == "GREATER_OR_EQUAL": |
| if converted[0] is None: |
| return pa.array([False] * batch_len, type=pa.bool_()) |
| return _null_as_false(pc.greater_equal(value_array, converted[0])) |
| elif function == "IS_NULL": |
| return pc.is_null(value_array) |
| elif function == "IS_NOT_NULL": |
| return pc.is_valid(value_array) |
| elif function == "IN": |
| non_null = [v for v in converted if v is not None] |
| if not non_null: |
| return pa.array([False] * batch_len, type=pa.bool_()) |
| in_mask = pc.is_in(value_array, pa.array(non_null, type=value_array.type)) |
| return pc.if_else(pc.is_valid(value_array), in_mask, False) |
| elif function == "NOT_IN": |
| if any(lit is None for lit in literals): |
| return pa.array([False] * batch_len, type=pa.bool_()) |
| not_in_mask = pc.invert( |
| pc.is_in(value_array, pa.array(converted, type=value_array.type))) |
| return pc.if_else(pc.is_valid(value_array), not_in_mask, False) |
| elif function == "BETWEEN": |
| if converted[0] is None or converted[1] is None: |
| return pa.array([False] * batch_len, type=pa.bool_()) |
| return _null_as_false(pc.and_( |
| pc.greater_equal(value_array, converted[0]), |
| pc.less_equal(value_array, converted[1]))) |
| elif function == "NOT_BETWEEN": |
| if converted[0] is None or converted[1] is None: |
| return pa.array([False] * batch_len, type=pa.bool_()) |
| return _null_as_false(pc.or_( |
| pc.less(value_array, converted[0]), |
| pc.greater(value_array, converted[1]))) |
| elif function == "STARTS_WITH": |
| if converted[0] is None: |
| return pa.array([False] * batch_len, type=pa.bool_()) |
| return _null_as_false(pc.starts_with(value_array, converted[0])) |
| elif function == "ENDS_WITH": |
| if converted[0] is None: |
| return pa.array([False] * batch_len, type=pa.bool_()) |
| return _null_as_false(pc.ends_with(value_array, converted[0])) |
| elif function == "CONTAINS": |
| if converted[0] is None: |
| return pa.array([False] * batch_len, type=pa.bool_()) |
| return _null_as_false(pc.match_substring(value_array, converted[0])) |
| elif function == "LIKE": |
| raw = literals[0] |
| if raw is None: |
| return pa.array([False] * batch_len, type=pa.bool_()) |
| from pypaimon.common.predicate import Like |
| pattern = Like._sql_like_to_regex(raw) |
| return _null_as_false( |
| pc.match_substring_regex(value_array, f"^{pattern}$")) |
| elif function == "TRUE": |
| return pa.array([True] * batch_len, type=pa.bool_()) |
| elif function == "FALSE": |
| return pa.array([False] * batch_len, type=pa.bool_()) |
| elif function == "IS_NAN": |
| return _null_as_false(pc.is_nan(value_array)) |
| raise ValueError(f"Unknown leaf function: {function}") |
| |
| |
| def _convert_literal(literal, target_type: pa.DataType): |
| if literal is None: |
| return None |
| if pa.types.is_timestamp(target_type): |
| import datetime |
| if isinstance(literal, str): |
| dt = datetime.datetime.fromisoformat(literal.replace("Z", "+00:00")) |
| return pa.scalar(dt, type=target_type) |
| elif isinstance(literal, list): |
| dt = datetime.datetime(*literal[:6]) |
| if len(literal) > 6: |
| dt = dt.replace(microsecond=literal[6] // 1000) |
| return pa.scalar(dt, type=target_type) |
| elif isinstance(literal, (int, float)): |
| dt = datetime.datetime.fromtimestamp(literal, tz=datetime.timezone.utc) |
| return pa.scalar(dt, type=target_type) |
| elif pa.types.is_date(target_type): |
| import datetime |
| if isinstance(literal, str): |
| return pa.scalar(datetime.date.fromisoformat(literal), type=target_type) |
| elif isinstance(literal, list): |
| return pa.scalar(datetime.date(*literal[:3]), type=target_type) |
| elif pa.types.is_time(target_type): |
| import datetime |
| if isinstance(literal, str): |
| t = datetime.time.fromisoformat(literal) |
| return pa.scalar(t, type=target_type) |
| elif isinstance(literal, list): |
| t = datetime.time(*literal[:3]) |
| if len(literal) > 3: |
| t = t.replace(microsecond=literal[3] // 1000) |
| return pa.scalar(t, type=target_type) |
| elif pa.types.is_decimal(target_type): |
| import decimal |
| return pa.scalar(decimal.Decimal(str(literal)), type=target_type) |
| return literal |
| |
| |
| def _paimon_type_to_arrow(paimon_type: str) -> pa.DataType: |
| type_str = paimon_type.strip().upper() |
| |
| ltz_match = re.match( |
| r"^TIMESTAMP\s*\((\d+)\)\s+WITH\s+LOCAL\s+TIME\s+ZONE", type_str) |
| if ltz_match: |
| precision = int(ltz_match.group(1)) |
| return pa.timestamp(_timestamp_precision_to_unit(precision), tz="UTC") |
| |
| m = re.match(r"^([A-Z_ ]+?)(?:\((.+)\))?(?:\s+NOT\s+NULL)?$", type_str) |
| if not m: |
| raise ValueError(f"Cannot parse Paimon type: '{paimon_type}'") |
| base_type = m.group(1).strip() |
| params = m.group(2) |
| |
| simple_mapping = { |
| "INT": pa.int32(), |
| "BIGINT": pa.int64(), |
| "SMALLINT": pa.int16(), |
| "TINYINT": pa.int8(), |
| "FLOAT": pa.float32(), |
| "DOUBLE": pa.float64(), |
| "STRING": pa.string(), |
| "BOOLEAN": pa.bool_(), |
| "BYTES": pa.binary(), |
| "DATE": pa.date32(), |
| } |
| if base_type in simple_mapping: |
| return simple_mapping[base_type] |
| |
| if base_type in ("VARCHAR", "CHAR"): |
| return pa.string() |
| |
| if base_type in ("VARBINARY", "BINARY"): |
| return pa.binary() |
| |
| if base_type == "TIMESTAMP": |
| precision = int(params) if params else 6 |
| unit = _timestamp_precision_to_unit(precision) |
| return pa.timestamp(unit) |
| |
| if base_type in ("TIMESTAMP WITH LOCAL TIME ZONE", "TIMESTAMP_WITH_LOCAL_TIME_ZONE", "TIMESTAMP_LTZ"): |
| precision = int(params) if params else 6 |
| unit = _timestamp_precision_to_unit(precision) |
| return pa.timestamp(unit, tz="UTC") |
| |
| if base_type == "TIME": |
| precision = int(params) if params else 6 |
| unit = _timestamp_precision_to_unit(precision) |
| return pa.time64(unit) if unit in ("us", "ns") else pa.time32(unit) |
| |
| if base_type == "DECIMAL": |
| if params: |
| parts = [x.strip() for x in params.split(",")] |
| if len(parts) == 2: |
| return pa.decimal128(int(parts[0]), int(parts[1])) |
| raise ValueError(f"DECIMAL type requires (precision, scale): '{paimon_type}'") |
| |
| raise ValueError( |
| f"Unsupported Paimon type for PyArrow conversion: '{paimon_type}'. " |
| f"Supported: INT, BIGINT, SMALLINT, TINYINT, FLOAT, DOUBLE, STRING, VARCHAR, CHAR, " |
| f"BOOLEAN, BYTES, VARBINARY, DATE, TIME(p), TIMESTAMP(p), " |
| f"TIMESTAMP WITH LOCAL TIME ZONE(p), DECIMAL(p,s)." |
| ) |
| |
| |
| def _timestamp_precision_to_unit(precision: int) -> str: |
| if precision == 0: |
| return "s" |
| elif precision <= 3: |
| return "ms" |
| elif precision <= 6: |
| return "us" |
| else: |
| return "ns" |
| |
| |
| def extract_referenced_fields(json_str: str) -> set: |
| data = json.loads(json_str) |
| fields = set() |
| _collect_fields(data, fields) |
| return fields |
| |
| |
| def _collect_fields(data: dict, fields: set): |
| kind = data.get("kind") |
| if kind == "LEAF": |
| _collect_all_field_refs_from_transform(data["transform"], fields) |
| elif kind == "COMPOUND": |
| for child in data["children"]: |
| _collect_fields(child, fields) |
| |
| |
| def _collect_all_field_refs_from_transform(transform: dict, fields: set = None) -> set: |
| if fields is None: |
| fields = set() |
| name = transform.get("name") |
| if name == "FIELD_REF" and "fieldRef" in transform: |
| fields.add(transform["fieldRef"]["name"]) |
| elif name == "CAST" and "fieldRef" in transform: |
| fields.add(transform["fieldRef"]["name"]) |
| else: |
| for inp in transform.get("inputs", []): |
| if isinstance(inp, dict): |
| if "name" in inp and "index" in inp: |
| fields.add(inp["name"]) |
| elif "name" in inp: |
| _collect_all_field_refs_from_transform(inp, fields) |
| return fields |
