JingsongLi commented on code in PR #8136:
URL: https://github.com/apache/paimon/pull/8136#discussion_r3450555496
##########
paimon-python/pypaimon/catalog/catalog_environment.py:
##########
@@ -117,3 +117,12 @@ def empty() -> 'CatalogEnvironment':
catalog_loader=None,
supports_version_management=False
)
+
+ def table_query_auth(self, options, identifier):
+ if not options.query_auth_enabled or self.catalog_loader is None:
+ return None
+
+ def auth(select):
Review Comment:
This nested `auth` function is stored on `FileStoreTable` as
`_query_auth_fn`, which makes auth-enabled tables/read tasks unpickleable
(`AttributeError: Can't get local object ...`). Ray/Daft and the existing
serialization paths pickle tables or read tasks, so this breaks those paths
once `query-auth.enabled` is true. Please use a top-level callable/dataclass or
reconstruct the catalog auth function lazily instead of storing a local closure.
##########
paimon-python/pypaimon/catalog/table_query_auth.py:
##########
@@ -0,0 +1,87 @@
+################################################################################
+# 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.
+################################################################################
+
+from typing import Callable, Dict, List, Optional
+
+import pyarrow as pa
+import pyarrow.compute as pc
+
+from pypaimon.common.predicate_json_parser import (
+ extract_referenced_fields,
+ parse_predicate_to_batch_filter,
+)
+from pypaimon.schema.data_types import DataField
+
+
+class TableNoPermissionException(Exception):
Review Comment:
This defines a second `TableNoPermissionException` that is unrelated to
`pypaimon.catalog.catalog_exception.TableNoPermissionException`.
`RESTCatalog.auth_table_query` raises this new type on 403, so callers catching
the catalog exception will miss query-auth denials. Please reuse the existing
catalog exception or make this subclass it.
##########
paimon-python/pypaimon/common/predicate_json_parser.py:
##########
@@ -0,0 +1,344 @@
+################################################################################
+# 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) ->
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)
+
+ 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 pc.binary_join_element_wise(*values, sep, null_handling='skip')
+
+ elif name == "NULL":
+ return pa.nulls(len(batch), type=pa.bool_())
+
+ raise ValueError(f"Unknown transform type in predicate: {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 _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."""
+ converted = [_convert_literal(lit, value_array.type) for lit in literals]
+
+ if function == "EQUAL":
+ return pc.equal(value_array, converted[0])
+ elif function == "NOT_EQUAL":
+ return pc.not_equal(value_array, converted[0])
+ elif function == "LESS_THAN":
+ return pc.less(value_array, converted[0])
+ elif function == "LESS_OR_EQUAL":
+ return pc.less_equal(value_array, converted[0])
+ elif function == "GREATER_THAN":
+ return pc.greater(value_array, converted[0])
+ elif function == "GREATER_OR_EQUAL":
+ return 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":
+ return pc.is_in(value_array, pa.array(converted,
type=value_array.type))
Review Comment:
Java's `In.test` returns false when the field value is null and skips null
literals. `pc.is_in` treats `NULL IN (..., NULL)` as true, so an auth row
filter like `dept IN ('eng', NULL)` would leak rows where `dept` is null.
Please drop null literals and wrap the result with
`pc.if_else(pc.is_valid(value_array), ..., False)`; if no non-null literals
remain, the result should be all false.
##########
paimon-python/pypaimon/read/reader/auth_masking_reader.py:
##########
@@ -0,0 +1,224 @@
+################################################################################
+# 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
+from typing import Callable, Dict, List, Optional
+
+import pyarrow as pa
+import pyarrow.compute as pc
+
+from pypaimon.common.predicate_json_parser import (
+ _collect_all_field_refs_from_transform,
+ _paimon_type_to_arrow,
+)
+from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader
+
+
+class RecordReaderToBatchAdapter(RecordBatchReader):
+
+ def __init__(self, inner, schema: pa.Schema, chunk_size: int = 65536,
include_row_kind: bool = False):
+ self._inner = inner
+ self._schema = schema
+ self._chunk_size = chunk_size
+ self._exhausted = False
+ self._pending_iterator = None
+ self._include_row_kind = include_row_kind
+
+ def read_arrow_batch(self) -> Optional[pa.RecordBatch]:
+ if self._exhausted:
+ return None
+ row_tuples = []
+ row_kinds = []
+ while len(row_tuples) < self._chunk_size:
+ if self._pending_iterator is not None:
+ row = self._pending_iterator.next()
+ while row is not None:
+ row_tuples.append(
+ row.row_tuple[row.offset:row.offset + row.arity])
+ if self._include_row_kind:
+ row_kinds.append(row.get_row_kind().to_string())
+ if len(row_tuples) >= self._chunk_size:
+ return self._flush(row_tuples, row_kinds)
+ row = self._pending_iterator.next()
+ self._pending_iterator = None
+
+ row_iterator = self._inner.read_batch()
+ if row_iterator is None:
+ self._exhausted = True
+ break
+ self._pending_iterator = row_iterator
+
+ if not row_tuples:
+ return None
+ return self._flush(row_tuples, row_kinds)
+
+ def _flush(self, row_tuples, row_kinds=None):
+ columns_data = list(zip(*row_tuples))
+ pydict = {
+ name: list(col)
+ for name, col in zip(self._schema.names, columns_data)
+ }
+ batch = pa.RecordBatch.from_pydict(pydict, schema=self._schema)
+ if row_kinds:
+ row_kind_array = pa.array(row_kinds, type=pa.string())
+ row_kind_field = pa.field("_row_kind", pa.string())
+ new_schema = pa.schema([row_kind_field] + list(batch.schema))
+ columns = [row_kind_array] + [batch.column(i) for i in
range(batch.num_columns)]
+ batch = pa.RecordBatch.from_arrays(columns, schema=new_schema)
+ return batch
+
+ def close(self):
+ self._inner.close()
+
+
+class AuthFilterReader(RecordBatchReader):
+
+ def __init__(self, inner_reader: RecordBatchReader, filter_fn:
Callable[[pa.RecordBatch], pa.Array]):
+ self._inner = inner_reader
+ self._filter_fn = filter_fn
+
+ def read_arrow_batch(self) -> Optional[pa.RecordBatch]:
+ batch = self._inner.read_arrow_batch()
+ if batch is None:
+ return None
+ mask = self._filter_fn(batch)
+ return batch.filter(mask)
+
+ def close(self):
+ self._inner.close()
+
+
+class AuthMaskingReader(RecordBatchReader):
+
+ def __init__(self, inner_reader: RecordBatchReader, masking_rules:
Dict[str, str], read_fields: List):
+ self._inner = inner_reader
+ self._masking_rules = masking_rules
+ self._read_fields = read_fields
+ read_field_names = {f.name for f in read_fields}
+ # Filter to projected columns, then validate field references.
+ self._parsed_rules = {
+ col: json.loads(tj) for col, tj in masking_rules.items()
+ if col in read_field_names
+ }
+ for col_name, transform in self._parsed_rules.items():
+ for ref_name in _collect_all_field_refs_from_transform(transform):
+ if ref_name not in read_field_names:
+ raise RuntimeError(
+ f"Column masking refers to field '{ref_name}' which is
not present "
+ f"in output row type. Available fields:
{read_field_names}"
+ )
+
+ def read_arrow_batch(self) -> Optional[pa.RecordBatch]:
+ batch = self._inner.read_arrow_batch()
+ if batch is None:
+ return None
+ original_batch = batch
+ masked_columns = {}
+ for col_name, transform in self._parsed_rules.items():
+ if col_name in original_batch.schema.names:
+ col_idx = original_batch.schema.get_field_index(col_name)
+ target_col_type = original_batch.schema.field(col_idx).type
+ masked_columns[col_idx] = self._apply_transform(transform,
original_batch, target_col_type)
+ for col_idx, masked_array in masked_columns.items():
+ original_field = original_batch.schema.field(col_idx)
+ if masked_array.type != original_field.type:
+ masked_array = pc.cast(masked_array, original_field.type)
+ batch = batch.set_column(
+ col_idx, pa.field(original_field.name, original_field.type,
nullable=True), masked_array)
+ return batch
+
+ def close(self):
+ self._inner.close()
+
+ def _apply_transform(
+ self,
+ transform: dict,
+ original_batch: pa.RecordBatch,
+ target_col_type: pa.DataType,
+ ) -> pa.Array:
+ name = transform["name"]
+
+ if name == "NULL":
+ return pa.nulls(len(original_batch), type=target_col_type)
+
+ elif name == "FIELD_REF":
+ ref_name = transform["fieldRef"]["name"]
+ return original_batch.column(ref_name)
+
+ elif name == "CAST":
+ ref_name = transform["fieldRef"]["name"]
+ source_col = original_batch.column(ref_name)
+ target_type = _paimon_type_to_arrow(transform["type"])
+ return pc.cast(source_col, target_type)
+
+ elif name == "UPPER":
+ input_col = self._resolve_input(transform["inputs"][0],
original_batch)
+ return pc.utf8_upper(input_col)
+
+ elif name == "LOWER":
+ input_col = self._resolve_input(transform["inputs"][0],
original_batch)
+ return pc.utf8_lower(input_col)
+
+ elif name == "CONCAT":
+ return self._apply_concat(transform["inputs"], original_batch)
+
+ elif name == "CONCAT_WS":
+ return self._apply_concat_ws(transform["inputs"], original_batch)
+
+ raise ValueError(f"Unknown transform type: {name}")
+
+ def _resolve_input(self, inp, original_batch: pa.RecordBatch) -> pa.Array:
+ if isinstance(inp, dict):
+ return original_batch.column(inp["name"])
+ elif isinstance(inp, str):
+ return pa.array([inp] * len(original_batch), type=pa.string())
+ elif inp is None:
+ return pa.nulls(len(original_batch), type=pa.string())
+ return pa.array([str(inp)] * len(original_batch), type=pa.string())
+
+ def _apply_concat(self, inputs: list, original_batch: pa.RecordBatch) ->
pa.Array:
+ resolved = [self._resolve_input(inp, original_batch) for inp in inputs]
+ if not resolved:
+ return pa.nulls(len(original_batch), type=pa.string())
+ return pc.binary_join_element_wise(*resolved, "")
+
+ def _apply_concat_ws(self, inputs: list, original_batch: pa.RecordBatch)
-> pa.Array:
+ if len(inputs) < 2:
+ return pa.nulls(len(original_batch), type=pa.string())
+ sep = self._resolve_input(inputs[0], original_batch)
+ values = [self._resolve_input(inp, original_batch) for inp in
inputs[1:]]
+ return pc.binary_join_element_wise(*values, sep, null_handling='skip')
Review Comment:
`binary_join_element_wise(..., null_handling='skip')` does not match Java
`BinaryString.concatWs` when the separator is non-null and all value inputs for
a row are null. Java returns an empty string for that row, but Arrow returns a
shorter array, so `set_column` can fail with `ArrowInvalid: Added column's
length must match record batch's length`. Please normalize this case so the
result always has `len(original_batch)` rows; the predicate-side `CONCAT_WS`
path has the same issue.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
