From b951cc596bddd946cc2cb6fdd0512538aed09667 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Tue, 13 Oct 2026 00:56:57 -0700 Subject: [PATCH 1/5] [python] Apply query filters after column masking Predicates on masked columns could match raw values or drop rows before masking. Evaluate them on masked values after the row filter and before LIMIT, authorize filter columns that are not projected, and reject chunk shuffling that authorization filters would resize. Conjuncts that reference no masked column are still pushed down for pruning. Only the masks the filter reads run before it; all masks are applied after LIMIT, so BLOB payloads are fetched only for the rows that survive. Comparisons and string predicates use the query-schema types, and NOT IN with a NULL literal skips that conversion. --- .../pypaimon/read/query_auth_split.py | 12 +- .../read/reader/auth_masking_reader.py | 18 +- .../read/reader/filter_record_batch_reader.py | 70 ++- .../pypaimon/read/scanner/file_scanner.py | 58 +-- paimon-python/pypaimon/read/split_read.py | 15 +- paimon-python/pypaimon/read/table_read.py | 82 ++- paimon-python/pypaimon/read/table_scan.py | 55 +- .../tests/deferred_blob_resolve_test.py | 118 ++++- .../pypaimon/tests/query_auth_filter_test.py | 472 ++++++++++++++++++ .../pypaimon/tests/table_query_auth_test.py | 1 + 10 files changed, 785 insertions(+), 116 deletions(-) create mode 100644 paimon-python/pypaimon/tests/query_auth_filter_test.py diff --git a/paimon-python/pypaimon/read/query_auth_split.py b/paimon-python/pypaimon/read/query_auth_split.py index bfbb1a2799a2..b172223558b5 100644 --- a/paimon-python/pypaimon/read/query_auth_split.py +++ b/paimon-python/pypaimon/read/query_auth_split.py @@ -68,10 +68,20 @@ def __getattr__(self, name): return getattr(self._split, name) -def resolve_auth_result(query_auth_fn, read_type): +def resolve_auth_result(query_auth_fn, read_type, predicate=None, partition_predicate=None, + columns=None): if query_auth_fn is None: return None select = [f.name for f in read_type] if read_type else None + if select is not None: + from pypaimon.read.push_down_utils import predicate_field_names + for query_filter in (predicate, partition_predicate): + if query_filter is not None: + names = predicate_field_names(query_filter) - set(select) + if columns is not None: + # other names are leaf aliases of projected columns, which are already selected + names &= columns + select.extend(sorted(names)) result = query_auth_fn(select) if result is None or not result.has_restrictions: return None diff --git a/paimon-python/pypaimon/read/reader/auth_masking_reader.py b/paimon-python/pypaimon/read/reader/auth_masking_reader.py index 5c6873052afe..c40711ad6002 100644 --- a/paimon-python/pypaimon/read/reader/auth_masking_reader.py +++ b/paimon-python/pypaimon/read/reader/auth_masking_reader.py @@ -217,13 +217,15 @@ def apply_auth_masking(batch: pa.RecordBatch, parsed_rules: Dict[str, dict]): target_col_type = original_batch.schema.field(col_idx).type masked_columns[col_idx] = _apply_predicate_transform( transform, original_batch, null_type=target_col_type) + if not masked_columns: + return batch + columns = list(batch.columns) + fields = list(batch.schema) for col_idx, masked_array in masked_columns.items(): original_field = original_batch.schema.field(col_idx) - batch = batch.set_column( - col_idx, - original_field.with_type(masked_array.type).with_nullable(True), - masked_array) - return batch + columns[col_idx] = masked_array + fields[col_idx] = original_field.with_type(masked_array.type).with_nullable(True) + return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields, metadata=batch.schema.metadata)) class AuthMaskingReader(RecordBatchReader): @@ -262,7 +264,11 @@ def read_arrow_batch(self) -> Optional[pa.RecordBatch]: columns = self._columns if "_row_kind" in batch.schema.names and "_row_kind" not in columns: columns = ["_row_kind"] + list(columns) - return batch.select(columns) + if not columns: + return zero_column_batch(batch.num_rows, batch.schema.metadata) + return pa.RecordBatch.from_arrays( + [batch.column(name) for name in columns], + schema=pa.schema([batch.schema.field(name) for name in columns], metadata=batch.schema.metadata)) def close(self): self._inner.close() diff --git a/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py b/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py index 381536e7c0da..7529bf0b160c 100644 --- a/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py +++ b/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py @@ -16,7 +16,7 @@ # under the License. import logging -from typing import List, Optional +from typing import Callable, List, Optional import pyarrow as pa import pyarrow.dataset as ds @@ -27,12 +27,28 @@ from pypaimon.read.reader.iface.record_batch_reader import ( InternalRowWrapperIterator, RecordBatchReader, + _polars_arrow_table, ) -from pypaimon.schema.data_types import DataField +from pypaimon.schema.data_types import DataField, PyarrowFieldParser logger = logging.getLogger(__name__) +def _comparison_fields(predicate): + if predicate.method in ('and', 'or'): + return set().union(*(_comparison_fields(child) for child in predicate.literals)) + literals = predicate.literals or [] + # NOT IN with a NULL literal never matches, so skip the cast. + if predicate.method == 'notIn' and any(value is None for value in literals): + return set() + if (predicate.method in ('equal', 'notEqual', 'lessThan', 'lessOrEqual', 'greaterThan', + 'greaterOrEqual', 'in', 'notIn', 'between', 'notBetween', + 'startsWith', 'endsWith', 'contains', 'like') + and any(value is not None for value in literals)): + return {predicate.field} + return set() + + class FilterRecordBatchReader(RecordBatchReader): """ Wraps a RecordBatchReader and filters each batch by predicate. @@ -46,11 +62,18 @@ def __init__( predicate: Predicate, field_names: Optional[List[str]] = None, schema_fields: Optional[List[DataField]] = None, + filter_view: Optional[Callable[[pa.RecordBatch], pa.RecordBatch]] = None, ): self.reader = reader self.predicate = predicate self.field_names = field_names self.schema_fields = schema_fields + # The predicate tests this view; rows are taken from the original batch. + self._filter_view = filter_view + comparison_fields = _comparison_fields(predicate) + self._comparison_types = { + field.name: PyarrowFieldParser.from_paimon_type(field.type) + for field in schema_fields or [] if field.name in comparison_fields} self._use_arrow_filter = predicate_supports_arrow_filter(predicate) self._adopt_metadata(reader) @@ -68,11 +91,21 @@ def read_arrow_batch(self) -> Optional[pa.RecordBatch]: continue def _filter_batch(self, batch: pa.RecordBatch) -> Optional[pa.RecordBatch]: + view = batch if self._filter_view is None else self._filter_view(batch) + filter_batch = self._coerce_comparison_types(view) if not self._use_arrow_filter: - return self._filter_batch_by_row(batch) - expr = self.predicate.to_arrow(batch.schema) + return self._filter_batch_by_row(batch, filter_batch) + expr = self.predicate.to_arrow(filter_batch.schema) if expr is None: - return self._filter_batch_by_row(batch) + return self._filter_batch_by_row(batch, filter_batch) + if filter_batch is not batch: + position = '_paimon_filter_position' + while position in filter_batch.schema.names: + position += '_' + data = pa.Table.from_batches([filter_batch]).append_column( + position, pa.array(range(batch.num_rows), type=pa.int64())) + selected = ds.InMemoryDataset(data).scanner(columns=[position], filter=expr).to_table() + return batch.take(selected.column(0).combine_chunks()) if selected.num_rows else None result = ds.InMemoryDataset(pa.Table.from_batches([batch])).scanner( filter=expr ).to_table() @@ -91,14 +124,29 @@ def _filter_batch(self, batch: pa.RecordBatch) -> Optional[pa.RecordBatch]: schema=result.schema, ) - def _filter_batch_by_row(self, batch: pa.RecordBatch) -> Optional[pa.RecordBatch]: + def _coerce_comparison_types(self, batch): + # Cast for comparison only; the output is unchanged. + columns = None + for name, data_type in self._comparison_types.items(): + column = batch.column(name) + if column.type != data_type: + if columns is None: + columns, fields = list(batch.columns), list(batch.schema) + index = batch.schema.get_field_index(name) + columns[index] = column.cast(data_type) + fields[index] = fields[index].with_type(data_type) + if columns is None: + return batch + return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields, metadata=batch.schema.metadata)) + + def _filter_batch_by_row(self, batch: pa.RecordBatch, filter_batch=None) -> Optional[pa.RecordBatch]: + filter_batch = batch if filter_batch is None else filter_batch # Predicate indices address logical fields. A record-reader adapter # may prepend _row_kind, which is metadata rather than an operand. - columns = ([field.name for field in self.schema_fields] - if self.schema_fields is not None else - [name for name in batch.schema.names if name != "_row_kind"]) - # PyArrow 6 has Table.select but not RecordBatch.select. - df = polars.from_arrow(pa.Table.from_batches([batch]).select(columns)) + columns = (self.field_names if self.field_names is not None else + [field.name for field in self.schema_fields] if self.schema_fields is not None else + [name for name in filter_batch.schema.names if name != "_row_kind"]) + df = polars.from_arrow(_polars_arrow_table(filter_batch).select(columns)) iterator = InternalRowWrapperIterator( self._iter_df_rows(df), df.width, diff --git a/paimon-python/pypaimon/read/scanner/file_scanner.py b/paimon-python/pypaimon/read/scanner/file_scanner.py index b938a7f7d262..e32c1388f665 100755 --- a/paimon-python/pypaimon/read/scanner/file_scanner.py +++ b/paimon-python/pypaimon/read/scanner/file_scanner.py @@ -230,39 +230,15 @@ def __init__( self.table: FileStoreTable = table self.manifest_scanner = manifest_scanner - self.predicate = predicate + self._user_predicate = predicate + self._partition_predicate = partition_predicate self._explicit_partition_filter = partition_predicate is not None - row_ranges = ( - _row_ranges_from_predicate(predicate) if predicate else None - ) - if predicate and row_ranges is not None: - self.predicate_for_stats = remove_row_id_filter(predicate) - else: - self.predicate_for_stats = predicate - self.predicate_for_stats = exclude_predicate_with_fields( - self.predicate_for_stats, SpecialFields.SYSTEM_FIELD_NAMES) - # Partition columns aren't in data files, so skip them for value-stats pruning. - self.predicate_for_stats = exclude_predicate_with_fields( - self.predicate_for_stats, set(self.table.partition_keys)) self.limit = limit self.snapshot_manager = table.snapshot_manager() self.manifest_list_manager = ManifestListManager(table) self.manifest_file_manager = ManifestFileManager(table) - self.primary_key_predicate = trim_and_transform_predicate( - self.predicate, self.table.field_names, self.table.trimmed_primary_keys) - - if partition_predicate is None: - self.partition_key_predicate = trim_and_transform_predicate( - self.predicate, self.table.field_names, self.table.partition_keys) - else: - # External predicate may carry full-schema indices and non-partition - # leaves; drop non-partition leaves and rebind the rest by name to the - # partition-row layout (matches derived path / Java), else IndexError. - self.partition_key_predicate = rewrite_predicate_indices( - trim_predicate_by_fields(partition_predicate, self.table.partition_keys), - self.table.partition_keys_fields) options = self.table.options # Get split target size and open file cost from table options self.target_split_size = options.source_split_target_size() @@ -290,17 +266,31 @@ def __init__( self.auth_partition_predicate = None self.auth_has_non_partition_filter = False - # Predicate-driven bucket pruning (HASH_FIXED only). Mirrors Java - # BucketSelectConverter. Set on demand and reused across all - # _filter_manifest_entry calls; the inner _Selector caches the - # bucket set per ``total_buckets`` value. - self._bucket_selector = self._init_bucket_selector() + self._set_predicate(predicate) self.simple_stats_evolutions = SimpleStatsEvolutions( self._schema_fields, self.table.table_schema.id ) + def _set_predicate(self, predicate): + self.predicate = predicate + row_ranges = _row_ranges_from_predicate(predicate) if predicate else None + self.predicate_for_stats = ( + remove_row_id_filter(predicate) if row_ranges is not None else predicate) + self.predicate_for_stats = exclude_predicate_with_fields( + self.predicate_for_stats, SpecialFields.SYSTEM_FIELD_NAMES | set(self.table.partition_keys)) + self.primary_key_predicate = trim_and_transform_predicate( + predicate, self.table.field_names, self.table.trimmed_primary_keys) + if self._partition_predicate is None: + self.partition_key_predicate = trim_and_transform_predicate( + predicate, self.table.field_names, self.table.partition_keys) + else: + self.partition_key_predicate = rewrite_predicate_indices( + trim_predicate_by_fields(self._partition_predicate, self.table.partition_keys), + self.table.partition_keys_fields) + self._bucket_selector = self._init_bucket_selector() + def _schema_fields(self, schema_id: int): """Resolve schema fields, short-circuiting current table schema id to avoid filesystem access (REST catalog would get 403). @@ -727,12 +717,14 @@ def _validate_chunk_shuffle_compat(self) -> None: raise ValueError("chunk_shuffle cannot combine with global index") if self._row_ranges is not None: raise ValueError("chunk_shuffle cannot combine with row ranges") + if self.auth_has_non_partition_filter: + raise ValueError("chunk_shuffle cannot combine with authorization filters evaluated during reading") # Only partition predicates are allowed: row-level / column-level # predicates would silently shrink each chunk's effective row count, # breaking the chunk_size contract DataLoader callers expect. - if self.predicate is not None: + if self._user_predicate is not None: partition_keys = set(self.table.partition_keys or []) - non_partition_fields = _get_all_fields(self.predicate) - partition_keys + non_partition_fields = _get_all_fields(self._user_predicate) - partition_keys if non_partition_fields: raise ValueError( "chunk_shuffle predicate must reference only partition keys; " diff --git a/paimon-python/pypaimon/read/split_read.py b/paimon-python/pypaimon/read/split_read.py index e23db8092d7f..596c661cf715 100644 --- a/paimon-python/pypaimon/read/split_read.py +++ b/paimon-python/pypaimon/read/split_read.py @@ -1393,7 +1393,8 @@ def __init__( outer_flat_read_type: Optional[List[DataField]] = None, post_merge_filter=None, eager_blob_fields=None, - post_filter_after_inline=False): + post_filter_after_inline=False, + defer_blob_resolution=False): self.row_ranges = None actual_split = split if isinstance(split, IndexedSplit): @@ -1409,6 +1410,7 @@ def __init__( self._post_merge_filter = post_merge_filter # Apply the auth filter after inline BLOB resolution, so scalar BLOBs still defer. self._post_filter_after_inline = post_filter_after_inline + self._defer_blob_resolution = defer_blob_resolution self._eager_blob_fields = set(eager_blob_fields or []) self._deferred_blob_fields = self._deferred_blob_field_names() @@ -1427,11 +1429,18 @@ def create_reader(self) -> RecordReader: if self.limit is not None: reader = LimitedRecordBatchReader(reader, self.limit) + if not self._defer_blob_resolution: + reader = self._resolve_deferred_blobs(reader) + return reader + + def _resolve_deferred_blobs(self, reader, excluded_fields=()): if self._deferred_blob_fields: blob_names = [ field.name for field in self.read_fields - if field.name in self._deferred_blob_fields + if field.name in self._deferred_blob_fields and field.name not in excluded_fields ] + if not blob_names: + return reader reader = DeferredBlobResolveReader( reader, self.table.file_io, @@ -1447,7 +1456,7 @@ def _deferred_blob_field_names(self) -> set: self.read_fields, self.predicate_for_reader, self.limit, - has_post_filter=self._post_merge_filter is not None, + has_post_filter=self._post_merge_filter is not None or self._defer_blob_resolution, ) - self._eager_blob_fields def _create_raw_reader(self) -> RecordReader: diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index 6e0b26362f16..29545e4cba13 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -29,14 +29,16 @@ from pypaimon.common.options.core_options import CoreOptions from pypaimon.common.predicate import Predicate -from pypaimon.common.predicate_json_parser import extract_referenced_fields +from pypaimon.common.predicate_json_parser import ( + _collect_all_field_refs_from_transform, extract_referenced_fields) from pypaimon.data.map_shared_shredding import map_selected_keys_field from pypaimon.read.push_down_utils import ( exclude_predicate_with_fields, predicate_field_names, rewrite_predicate_indices) from pypaimon.read.query_auth_split import QueryAuthSplit from pypaimon.read.reader.auth_masking_reader import ( AuthFilterReader, AuthMaskingReader, ColumnProjectReader, - RecordReaderToBatchAdapter, BatchToRecordReaderAdapter) + RecordReaderToBatchAdapter, BatchToRecordReaderAdapter, + apply_auth_masking, compile_auth_masking_rules) from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader from pypaimon.read.reader.limited_record_reader import LimitedRecordBatchReader from pypaimon.read.reader.filter_record_batch_reader import FilterRecordBatchReader @@ -1487,7 +1489,8 @@ def _create_split_read(self, split: Split, blob_parallelism: int = 1, post_merge_filter=None, eager_blob_fields=None, post_filter_after_inline: bool = False, - predicate=_DEFAULT_PREDICATE) -> SplitRead: + predicate=_DEFAULT_PREDICATE, + defer_blob_resolution: bool = False) -> SplitRead: sr = self._build_split_read( split, read_type, @@ -1497,6 +1500,7 @@ def _create_split_read(self, split: Split, blob_parallelism: int = 1, eager_blob_fields, post_filter_after_inline, predicate, + defer_blob_resolution, ) sr._blob_parallelism = blob_parallelism sr._parquet_row_group_cache = self._parquet_row_group_cache @@ -1508,7 +1512,8 @@ def _build_split_read(self, split: Split, read_type=None, post_merge_filter=None, eager_blob_fields=None, post_filter_after_inline: bool = False, - predicate=_DEFAULT_PREDICATE) -> SplitRead: + predicate=_DEFAULT_PREDICATE, + defer_blob_resolution: bool = False) -> SplitRead: predicate = self.predicate if predicate is _DEFAULT_PREDICATE else predicate effective_limit = ( self.limit if limit is None else limit @@ -1603,6 +1608,7 @@ def _build_split_read(self, split: Split, read_type=None, post_merge_filter=post_merge_filter, eager_blob_fields=eager_blob_fields, post_filter_after_inline=post_filter_after_inline, + defer_blob_resolution=defer_blob_resolution, ) else: inner_read_type = scan_read_type @@ -1679,7 +1685,7 @@ def _project_batch_to_output(self, batch: pyarrow.RecordBatch) -> pyarrow.Record arrays.insert(0, batch.column(ROW_KIND_COLUMN)) fields.insert(0, batch.schema.field(ROW_KIND_COLUMN)) if not arrays: - return batch.select([]) + return zero_column_batch(batch.num_rows, batch.schema.metadata) return pyarrow.RecordBatch.from_arrays( arrays, schema=pyarrow.schema(fields, metadata=batch.schema.metadata)) if not self._needs_output_projection(): @@ -1693,7 +1699,7 @@ def _project_batch_to_output(self, batch: pyarrow.RecordBatch) -> pyarrow.Record arrays = [batch.column(name_to_pos[name]) for name in output_names] fields = [batch.schema.field(name_to_pos[name]) for name in output_names] if not arrays: - return batch.select([]) + return zero_column_batch(batch.num_rows, batch.schema.metadata) return pyarrow.RecordBatch.from_arrays( arrays, schema=pyarrow.schema(fields)) @@ -1829,6 +1835,10 @@ def __authed_reader(self, split, auth_result, blob_parallelism=1, "MAP-key projection with query authorization is not supported") table_fields = self.table.fields read_fields = self._adapter_read_type + masked_fields = set(auth_result.column_masking or {}) + post_filter = (self.predicate if self.predicate is not None + and predicate_field_names(self.predicate) & masked_fields else None) + filter_read_fields = self._with_predicate_extra_fields(read_fields) if post_filter is not None else read_fields from pypaimon.read.table_scan import latest_auth_fields, validate_auth_rules # a split can carry rules without having been planned by a validating scan snapshot = getattr(self, "_auth_schema", None) @@ -1837,24 +1847,27 @@ def __authed_reader(self, split, auth_result, blob_parallelism=1, latest_fields = validate_auth_rules(self.table, auth_result, snapshot) # bind to the latest schema: a generated alias may since have become a real column auth_result.validate_read_type( - latest_fields, read_fields, self.nested_name_paths, table_fields) + latest_fields, filter_read_fields, self.nested_name_paths, table_fields) - # Query filters can reference masked columns absent from the output. - # Include them and their mask dependencies before the final projection. - filter_fields = self._predicate_extra_fields - needed_fields = read_fields + filter_fields - auth_extra_fields = auth_result.get_extra_fields(needed_fields, table_fields) - extra_fields = filter_fields + auth_extra_fields - effective_read_type = read_fields + extra_fields + extra_fields = auth_result.get_extra_fields(filter_read_fields, table_fields) + effective_read_type = filter_read_fields + extra_fields filter_fn = auth_result.extract_row_filter() - raw_predicate = exclude_predicate_with_fields( - self.predicate, set(auth_result.column_masking or {})) effective_limit = self.limit if limit is None else limit auth_fields = ( self._auth_filter_field_names(auth_result, effective_read_type) if filter_fn is not None else set() ) + eager_blob_fields = set(auth_fields) + if post_filter is not None: + # Only inputs of the filter and its masks are read eagerly. + filter_names = predicate_field_names(post_filter) + filter_masks = compile_auth_masking_rules( + {name: rule for name, rule in auth_result.column_masking.items() if name in filter_names}, + effective_read_type) + eager_blob_fields.update(filter_names - masked_fields) + for transform in filter_masks.values(): + eager_blob_fields.update(_collect_all_field_refs_from_transform(transform)) inline_blob_fields = ( self.table.options.blob_descriptor_fields() | self.table.options.blob_view_fields() @@ -1871,11 +1884,14 @@ def __authed_reader(self, split, auth_result, blob_parallelism=1, blob_parallelism=blob_parallelism, read_type=effective_read_type, limit=limit, - push_down_limit=(filter_fn is None or embed_filter) and self.predicate is None, - predicate=raw_predicate, + push_down_limit=post_filter is None and (filter_fn is None or embed_filter), post_merge_filter=filter_fn if embed_filter else None, - eager_blob_fields=auth_fields if embed_filter else None, + eager_blob_fields=eager_blob_fields, post_filter_after_inline=post_filter_after_inline, + # push down conjuncts without masked columns + predicate=(self.predicate if post_filter is None + else exclude_predicate_with_fields(post_filter, masked_fields)), + defer_blob_resolution=post_filter is not None, ) reader = split_read.create_reader() @@ -1897,17 +1913,29 @@ def __authed_reader(self, split, auth_result, blob_parallelism=1, if filter_fn and not embed_filter: reader = AuthFilterReader(reader, filter_fn) - if auth_result.column_masking: - reader = AuthMaskingReader(reader, auth_result.column_masking, effective_read_type) - - if self.predicate is not None: - predicate = rewrite_predicate_indices(self.predicate, effective_read_type) + if post_filter is not None: + # Filter on masked values but keep raw rows; masks apply after LIMIT. reader = FilterRecordBatchReader( - reader, predicate, schema_fields=effective_read_type) - if effective_limit is not None: + reader, rewrite_predicate_indices(post_filter, effective_read_type), + field_names=[field.name for field in effective_read_type], + schema_fields=effective_read_type, + filter_view=lambda batch: apply_auth_masking(batch, filter_masks)) + if effective_limit is not None and (post_filter is not None or (filter_fn and not embed_filter)): reader = LimitedRecordBatchReader(reader, effective_limit) - if extra_fields: + if post_filter is not None and isinstance(split_read, DataEvolutionSplitRead): + # Fetch only payloads output unmasked or read by a mask. + needed = {f.name for f in read_fields} - masked_fields + for transform in compile_auth_masking_rules( + auth_result.column_masking, effective_read_type).values(): + needed.update(_collect_all_field_refs_from_transform(transform)) + reader = split_read._resolve_deferred_blobs( + reader, {f.name for f in effective_read_type} - needed) + + if auth_result.column_masking: + reader = AuthMaskingReader(reader, auth_result.column_masking, effective_read_type) + + if len(effective_read_type) > len(read_fields): original_columns = [f.name for f in read_fields] reader = ColumnProjectReader(reader, original_columns) diff --git a/paimon-python/pypaimon/read/table_scan.py b/paimon-python/pypaimon/read/table_scan.py index 40f0383d30fc..fedc8c1441bc 100755 --- a/paimon-python/pypaimon/read/table_scan.py +++ b/paimon-python/pypaimon/read/table_scan.py @@ -27,9 +27,11 @@ from pypaimon.common.predicate_builder import PredicateBuilder from pypaimon.manifest.manifest_list_manager import ManifestListManager from pypaimon.read.plan import Plan +from pypaimon.read.push_down_utils import exclude_predicate_with_fields, predicate_field_names from pypaimon.read.query_auth_split import QueryAuthSplit, resolve_auth_result, wrap_plan_with_auth from pypaimon.read.scan_stats import ScanStats from pypaimon.read.scanner.file_scanner import FileScanner +from pypaimon.table.special_fields import SpecialFields logger = logging.getLogger(__name__) @@ -74,6 +76,8 @@ def plan(self) -> Plan: if native is not None: return native auth_result = self.__auth_query() + if self._query_auth_fn is not None: + prune_scanner_by_auth(self.table, self.file_scanner, auth_result) # The native planner bypasses the auth-aware file scanner. Resolve auth # before selecting a planning backend. if (not native_auth and auth_result is None and self._native_requested() @@ -81,8 +85,6 @@ def plan(self) -> Plan: native = self._try_native_plan() if native is not None: return native - if auth_result is not None: - prune_scanner_by_auth(self.table, self.file_scanner, auth_result) plan = self.file_scanner.scan() return wrap_plan_with_auth(auth_result, plan) @@ -280,6 +282,8 @@ def _try_native_plan(self) -> Optional[Plan]: def plan_for_write(self) -> Plan: if self.__auth_query() is not None: raise TableNoPermissionException(self.table.identifier) + if self._query_auth_fn is not None: + prune_scanner_by_auth(self.table, self.file_scanner, None) # Writer restore/overwrite must see L0 even when batch reads hide it. skip_level0 = self.file_scanner.skip_level0 try: @@ -306,13 +310,13 @@ def scan_with_stats(self) -> Tuple[Plan, Optional[ScanStats]]: if native is not None: return native, None auth_result = self.__auth_query() + if self._query_auth_fn is not None: + prune_scanner_by_auth(self.table, self.file_scanner, auth_result) if (not native_auth and auth_result is None and self._native_requested() and self._native_plan_supported()): native = self._try_native_plan() if native is not None: return native, None - if auth_result is not None: - prune_scanner_by_auth(self.table, self.file_scanner, auth_result) plan, stats = self.file_scanner.scan_with_stats() return wrap_plan_with_auth(auth_result, plan), stats @@ -606,48 +610,35 @@ def validate_auth_rules(table, auth_result, latest_fields=None): def authorize(table, query_auth_fn, read_type, predicate=None, partition_predicate=None): if query_auth_fn is None: return None - requested_type = read_type - if read_type: - from pypaimon.read.push_down_utils import predicate_field_names - selected = {field.name for field in read_type} - names = set(selected) - for query_filter in (predicate, partition_predicate): - if query_filter is not None: - names.update(predicate_field_names(query_filter)) - requested_type = list(read_type) + [ - field for field in table.fields if field.name in names and field.name not in selected] - auth_result = resolve_auth_result(query_auth_fn, requested_type) + fields = table.fields + if table.options.row_tracking_enabled(): + fields = SpecialFields.row_type_with_row_tracking(fields) + auth_result = resolve_auth_result(query_auth_fn, read_type, predicate, partition_predicate, + {field.name for field in fields}) if auth_result is not None: validate_auth_rules(table, auth_result) if partition_predicate is not None and auth_result.column_masking: - from pypaimon.read.push_down_utils import predicate_field_names if predicate_field_names(partition_predicate) & set(auth_result.column_masking): raise ValueError("Explicit partition filters cannot reference masked columns") return auth_result def prune_scanner_by_auth(table, scanner, auth_result): - # Java excludes masked columns from all raw-data pruning. The complete - # query predicate is executed after masks have produced visible values. - if auth_result.column_masking: - from pypaimon.read.push_down_utils import exclude_predicate_with_fields - masked = set(auth_result.column_masking) - query_predicate = getattr(scanner, 'predicate', None) - for name in ('predicate', 'predicate_for_stats', 'primary_key_predicate', 'partition_key_predicate'): - setattr(scanner, name, exclude_predicate_with_fields(getattr(scanner, name, None), masked)) - init_selector = getattr(scanner, '_init_bucket_selector', None) - if init_selector is not None: - scanner._bucket_selector = init_selector() - if query_predicate is not None: - scanner.auth_has_non_partition_filter = True - if not auth_result.filter: + masked = set(auth_result.column_masking or {}) if auth_result is not None else set() + predicate = scanner._user_predicate + effective = exclude_predicate_with_fields(predicate, masked) + if effective != scanner.predicate: + scanner._set_predicate(effective) + scanner.auth_partition_predicate = None + scanner.auth_has_non_partition_filter = bool( + predicate is not None and masked & predicate_field_names(predicate)) + if auth_result is None or not auth_result.filter: return partition_preds, has_non_partition = __split_auth_filter(table, auth_result) if partition_preds: combined = PredicateBuilder.and_predicates(partition_preds) scanner.auth_partition_predicate = combined - if has_non_partition: - scanner.auth_has_non_partition_filter = True + scanner.auth_has_non_partition_filter |= has_non_partition def __split_auth_filter(table, auth_result): diff --git a/paimon-python/pypaimon/tests/deferred_blob_resolve_test.py b/paimon-python/pypaimon/tests/deferred_blob_resolve_test.py index 3c8a1a071a6c..51ec83c1970b 100644 --- a/paimon-python/pypaimon/tests/deferred_blob_resolve_test.py +++ b/paimon-python/pypaimon/tests/deferred_blob_resolve_test.py @@ -29,6 +29,7 @@ from pypaimon import CatalogFactory, Schema from pypaimon.catalog.table_query_auth import TableQueryAuthResult from pypaimon.read.query_auth_split import QueryAuthSplit +from pypaimon.read.reader.format_blob_reader import BlobRecordIterator from pypaimon.table.row.blob import BlobRef @@ -122,6 +123,15 @@ def counted_to_data(blob): cls.blob_ref_to_data_patch = patch.object( BlobRef, "to_data", counted_to_data) cls.blob_ref_to_data_patch.start() + original_inline_read = BlobRecordIterator._read_inline_blob + + def counted_inline_read(reader, position, length): + if isinstance(reader.file_io, _BlobCountingFileIO): + reader.file_io.blobs_fetched += 1 + return original_inline_read(reader, position, length) + + cls.inline_blob_read_patch = patch.object(BlobRecordIterator, '_read_inline_blob', counted_inline_read) + cls.inline_blob_read_patch.start() cls.catalog = CatalogFactory.create({ "warehouse": os.path.join(cls.tempdir, "warehouse") }) @@ -134,6 +144,7 @@ def counted_to_data(blob): @classmethod def tearDownClass(cls): + cls.inline_blob_read_patch.stop() cls.blob_ref_to_data_patch.stop() shutil.rmtree(cls.tempdir, ignore_errors=True) @@ -172,14 +183,14 @@ def _create_table(self, name, extra_options=None, payloads=None, commit.close() return self.catalog.get_table(identifier) - def _read(self, table, predicate, limit=None, blob_parallelism=None): + def _read(self, table, predicate, limit=None, blob_parallelism=None, + projection=("sample_id", "payload", "score")): counting_file_io = _BlobCountingFileIO(table.file_io) table.file_io = counting_file_io read_builder = table.new_read_builder() if predicate is not None: read_builder = read_builder.with_filter(predicate) - read_builder = read_builder.with_projection( - ["sample_id", "payload", "score"]) + read_builder = read_builder.with_projection(list(projection)) if limit is not None: read_builder = read_builder.with_limit(limit) splits = read_builder.new_scan().plan().splits() @@ -367,6 +378,107 @@ def test_auth_only_defers_non_auth_payloads(self): self.assertEqual(_ROW_COUNT - 1, result.num_rows) self.assertEqual(_ROW_COUNT - 1, counting_file_io.blobs_fetched) + def test_post_mask_filter_defers_unrelated_payloads(self): + table = self._create_table('defer_post_mask') + auth = TableQueryAuthResult(None, {'sample_id': json.dumps({ + 'name': 'UPPER', 'inputs': [{'index': 0, 'name': 'sample_id', 'type': 'STRING'}]})}) + table.catalog_environment.table_query_auth = lambda options, identifier: lambda select: auth + predicate = table.new_read_builder().new_predicate_builder().is_not_null('sample_id') + result, counting = self._read(table, predicate, limit=2) + self.assertEqual(['SAMPLE_0', 'SAMPLE_1'], result.column('sample_id').to_pylist()) + self.assertEqual(2, counting.blobs_fetched) + + def test_post_mask_filter_keeps_blob_mask_inputs_eager(self): + table = self._create_table('defer_post_mask_inputs', payloads=[b'keep'] * _ROW_COUNT) + auth = TableQueryAuthResult(None, {'sample_id': json.dumps({ + 'name': 'FIELD_REF', 'fieldRef': {'index': 1, 'name': 'payload', 'type': 'BLOB'}})}) + table.catalog_environment.table_query_auth = lambda options, identifier: lambda select: auth + predicate = table.new_read_builder().new_predicate_builder().startswith('sample_id', 'keep') + result, counting = self._read(table, predicate, limit=2) + self.assertEqual(['keep', 'keep'], result.column('sample_id').to_pylist()) + self.assertEqual(_ROW_COUNT, counting.blobs_fetched) + + def test_post_mask_filter_defers_output_mask_inputs(self): + for projection in (('sample_id', 'payload', 'score'), ('sample_id', 'score')): + with self.subTest(projection=projection): + table = self._create_table( + 'defer_post_mask_output_inputs_%d' % len(projection), + payloads=[str(100 + index).encode() for index in range(_ROW_COUNT)]) + auth = TableQueryAuthResult(None, { + 'sample_id': json.dumps({'name': 'UPPER', 'inputs': [ + {'index': 0, 'name': 'sample_id', 'type': 'STRING'}]}), + 'score': json.dumps({'name': 'FIELD_REF', 'fieldRef': { + 'index': 1, 'name': 'payload', 'type': 'BLOB'}}), + }) + table.catalog_environment.table_query_auth = lambda options, identifier: lambda select: auth + predicate = table.new_read_builder().new_predicate_builder().is_not_null('sample_id') + result, counting = self._read(table, predicate, limit=2, projection=projection) + self.assertEqual(list(projection), result.column_names) + self.assertEqual(['SAMPLE_0', 'SAMPLE_1'], result.column('sample_id').to_pylist()) + self.assertEqual([100, 101], result.column('score').to_pylist()) + if 'payload' in projection: + self.assertEqual([b'100', b'101'], result.column('payload').to_pylist()) + self.assertEqual(2, counting.blobs_fetched) + + def test_post_mask_filter_skips_null_masked_payloads(self): + table = self._create_table('defer_post_mask_null_payloads') + auth = TableQueryAuthResult(None, { + 'sample_id': json.dumps({'name': 'UPPER', 'inputs': [ + {'index': 0, 'name': 'sample_id', 'type': 'STRING'}]}), + 'payload': json.dumps({'name': 'NULL'}), + }) + table.catalog_environment.table_query_auth = lambda options, identifier: lambda select: auth + predicate = table.new_read_builder().new_predicate_builder().is_not_null('sample_id') + result, counting = self._read(table, predicate, limit=2) + self.assertEqual([None, None], result.column('payload').to_pylist()) + self.assertEqual(0, counting.blobs_fetched) + + def test_post_mask_filter_defers_payloads_without_limit(self): + table = self._create_table('defer_post_mask_without_limit') + auth = TableQueryAuthResult(None, {'sample_id': json.dumps({ + 'name': 'UPPER', 'inputs': [{'index': 0, 'name': 'sample_id', 'type': 'STRING'}]})}) + table.catalog_environment.table_query_auth = lambda options, identifier: lambda select: auth + predicate = table.new_read_builder().new_predicate_builder().startswith('sample_id', 'SAMPLE_1') + result, counting = self._read(table, predicate, blob_parallelism=4) + self.assertEqual([1], result.column('score').to_pylist()) + self.assertEqual(1, counting.blobs_fetched) + + def test_post_mask_limit_follows_row_authorization(self): + table = self._create_table('defer_post_mask_row_auth') + auth = TableQueryAuthResult(_RejectScoreOneAuthResult.filter, {'sample_id': json.dumps({ + 'name': 'UPPER', 'inputs': [{'index': 0, 'name': 'sample_id', 'type': 'STRING'}]})}) + table.catalog_environment.table_query_auth = lambda options, identifier: lambda select: auth + predicate = table.new_read_builder().new_predicate_builder().is_not_null('sample_id') + result, counting = self._read(table, predicate, limit=2) + self.assertEqual([0, 2], result.column('score').to_pylist()) + self.assertEqual(2, counting.blobs_fetched) + + def test_post_mask_limit_defers_payloads_across_splits(self): + for method in ('batches', 'iterator'): + with self.subTest(method=method): + table = self._create_table( + 'defer_post_mask_splits_' + method, extra_options={'source.split.target-size': '1b'}, + partition_keys=['sample_id'], sample_ids=['a'] + ['b'] * (_ROW_COUNT - 1)) + table.file_io = counting = _BlobCountingFileIO(table.file_io) + auth = TableQueryAuthResult(None, {'sample_id': json.dumps({ + 'name': 'UPPER', 'inputs': [{'index': 0, 'name': 'sample_id', 'type': 'STRING'}]})}) + table.catalog_environment.table_query_auth = lambda options, identifier: lambda select: auth + predicate = table.new_read_builder().new_predicate_builder().is_not_null('sample_id') + builder = table.new_read_builder().with_filter(predicate).with_limit(2) + splits = builder.new_scan().plan().splits() + self.assertEqual(2, len(splits)) + reader = builder.new_read() + if method == 'batches': + result = pa.Table.from_batches(reader.to_arrow_batch_reader(splits)) + values = zip(result.column('score').to_pylist(), result.column('payload').to_pylist()) + else: + values = [(row.get_field(2), row.get_field(1)) for row in reader.to_iterator(splits)] + values = list(values) + self.assertEqual(2, len(values)) + self.assertEqual(2, counting.blobs_fetched) + for score, payload in values: + self.assertEqual(bytes([score]) * 1024, payload) + def test_auth_filter_preserves_file_io_for_blob_descriptors(self): table = self._create_table( "auth_blob_descriptor", diff --git a/paimon-python/pypaimon/tests/query_auth_filter_test.py b/paimon-python/pypaimon/tests/query_auth_filter_test.py new file mode 100644 index 000000000000..4e0c98357638 --- /dev/null +++ b/paimon-python/pypaimon/tests/query_auth_filter_test.py @@ -0,0 +1,472 @@ +# 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 pyarrow as pa +import pytest + +from pypaimon import CatalogFactory, Schema +from pypaimon.catalog.table_query_auth import TableQueryAuthResult +from pypaimon.common.predicate_builder import PredicateBuilder +from pypaimon.read.table_read import TableRead +from pypaimon.table.row.generic_row import GenericRow +from pypaimon.table.row.row_kind import RowKind + + +def field_ref(name, index): + return {'name': 'FIELD_REF', 'fieldRef': {'name': name, 'index': index, 'type': 'STRING'}} + + +def rules(mask=None, row_filter=None): + return TableQueryAuthResult( + [row_filter] if row_filter is not None else None, + {'secret': json.dumps(mask or {'name': 'NULL'})}) + + +def arrow_rows(data): + return [dict(zip(data.column_names, values)) + for values in zip(*(column.to_pylist() for column in data.columns))] + + +@pytest.fixture(params=['append', 'primary-key', 'data-evolution', + 'partitioned-append', 'partitioned-primary-key']) +def table(request, tmp_path): + catalog = CatalogFactory.create({'warehouse': str(tmp_path / 'warehouse')}) + catalog.create_database('db', True) + fields = pa.schema([('id', pa.int32()), ('secret', pa.string()), ('salt', pa.string())]) + options = {'file.format': 'parquet', 'write.native.enabled': 'false', + 'source.split.target-size': '1 b', 'source.split.open-file-cost': '1 b'} + if request.param == 'data-evolution': + options.update({'data-evolution.enabled': 'true', 'row-tracking.enabled': 'true'}) + primary_key = 'primary-key' in request.param + partitions = ['secret'] if request.param.startswith('partitioned') else [] + if primary_key: + options['bucket'] = '1' + catalog.create_table('db.t', Schema.from_pyarrow_schema( + fields, primary_keys=['id'] + partitions if primary_key else [], + partition_keys=partitions, options=options), False) + raw = catalog.get_table('db.t') + rows = [{'id': 1, 'secret': 'alpha', 'salt': 'hidden'}, + {'id': 2, 'secret': 'beta', 'salt': 'visible'}, + {'id': 3, 'secret': 'gamma', 'salt': 'hidden'}] + if primary_key: + rows.append(rows[0]) + for row in rows: + builder = raw.new_batch_write_builder() + writer = builder.new_write() + try: + writer.write_arrow(pa.Table.from_pydict({name: [value] for name, value in row.items()}, schema=fields)) + builder.new_commit().commit(writer.prepare_commit()) + finally: + writer.close() + return catalog.get_table('db.t').copy({'query-auth.enabled': 'true'}) + + +def authorize(table, result): + table.catalog_environment.table_query_auth = lambda options, identifier: lambda select: result + + +def read(table, predicate, projection=None, limit=None, method='arrow'): + builder = table.new_read_builder().with_filter(predicate) + if projection is not None: + builder.with_projection(projection) + if limit is not None: + builder.with_limit(limit) + splits = builder.new_scan().plan().splits() + reader = builder.new_read() + if method == 'iterator': + return [tuple(row.get_field(i) for i in range(len(reader.read_type))) + for row in reader.to_iterator(splits)] + if method == 'batches': + return arrow_rows(reader.to_arrow_batch_reader(splits).read_all()) + return arrow_rows(reader.to_arrow(splits)) + + +def stream_read(table, builder, phase, follow_up_snapshot=2): + scan = builder.new_streaming_scan() + snapshots = table.snapshot_manager() + try: + if phase == 'initial': + plan = scan._create_initial_plan(snapshots.get_latest_snapshot()) + elif phase == 'follow_up': + plan = scan._create_follow_up_plan(snapshots.get_snapshot_by_id(follow_up_snapshot)) + else: + plan = scan._create_catch_up_plan(follow_up_snapshot, snapshots.get_latest_snapshot()) + return arrow_rows(builder.new_read().to_arrow(plan.splits())) + finally: + scan._prefetch_executor.shutdown(wait=True) + + +def selection_sensitive(table, masked_rules): + selected = [] + + def auth(select): + selected.append(select) + return masked_rules if 'secret' in select else TableQueryAuthResult(None, None) + + table.catalog_environment.table_query_auth = lambda options, identifier: auth + return selected + + +@pytest.mark.parametrize('condition,expected', [('is_null', [1, 2, 3]), ('raw_equal', [])]) +def test_user_filter_sees_masked_values(table, condition, expected): + authorize(table, rules()) + predicates = PredicateBuilder(table.fields) + predicate = predicates.is_null('secret') if condition == 'is_null' else predicates.equal('secret', 'alpha') + rows = read(table, predicate) + assert sorted(row['id'] for row in rows) == expected + assert all(row['secret'] is None for row in rows) + + +@pytest.mark.parametrize('method', ['arrow', 'batches', 'iterator']) +def test_unprojected_mask_target_and_input(table, method): + authorize(table, rules(field_ref('salt', 2))) + predicate = PredicateBuilder(table.fields).equal('secret', 'visible') + rows = read(table, predicate, projection=['id'], method=method) + assert rows == ([(2,)] if method == 'iterator' else [{'id': 2}]) + + +@pytest.mark.parametrize('compound,expected', [('and', [2, 3]), ('or', [1, 2, 3])]) +def test_compound_predicate_keeps_masked_condition_intact(table, compound, expected): + authorize(table, rules()) + predicates = PredicateBuilder(table.fields) + combine = PredicateBuilder.and_predicates if compound == 'and' else PredicateBuilder.or_predicates + predicate = combine([predicates.is_null('secret'), predicates.greater_than('id', 1)]) + assert sorted(row['id'] for row in read(table, predicate, projection=['id'])) == expected + + +def test_limit_is_applied_after_masked_filter(table): + authorize(table, rules(field_ref('salt', 2))) + predicate = PredicateBuilder(table.fields).equal('secret', 'visible') + assert read(table, predicate, projection=['id'], limit=1) == [{'id': 2}] + + +def test_filter_columns_are_authorized(table): + selected = [] + + def auth(select): + selected.append(select) + return rules(field_ref('salt', 2)) if 'secret' in select else TableQueryAuthResult(None, None) + + table.catalog_environment.table_query_auth = lambda options, identifier: auth + predicate = PredicateBuilder(table.fields).equal('secret', 'visible') + assert read(table, predicate, projection=['id']) == [{'id': 2}] + assert selected == [['id', 'secret']] + + +@pytest.mark.parametrize('table', ['partitioned-append', 'partitioned-primary-key'], indirect=True) +def test_partition_filter_columns_are_authorized(table): + selected = selection_sensitive(table, rules()) + predicate = PredicateBuilder(table.fields).equal('secret', 'alpha') + scan = table.new_read_builder().with_projection(['id']).with_partition_filter(predicate).new_scan() + with pytest.raises(ValueError, match='partition filters cannot reference masked columns'): + scan.plan() + assert selected == [['id', 'secret']] + + +@pytest.mark.parametrize('phase', ['initial', 'follow_up', 'catch_up']) +def test_streaming_filter_columns_are_authorized(table, phase): + selected = selection_sensitive(table, rules(field_ref('salt', 2))) + predicate = PredicateBuilder(table.fields).equal('secret', 'visible') + builder = (table.new_stream_read_builder().with_filter(predicate) + .with_projection(['id']).with_include_row_kind(False)) + assert stream_read(table, builder, phase) == [{'id': 2}] + assert selected and all(select == ['id', 'secret'] for select in selected) + + +@pytest.mark.parametrize('masked_value,expected', [('visible', [2, 3]), ('beta', [3])]) +def test_or_arms_match_independently(table, masked_value, expected): + authorize(table, rules(field_ref('salt', 2))) + predicates = PredicateBuilder(table.fields) + # id 2 matches only the masked arm, id 3 only the other; 'beta' is id 2's raw value + predicate = PredicateBuilder.or_predicates([ + predicates.equal('secret', masked_value), predicates.equal('id', 3)]) + assert sorted(row['id'] for row in read(table, predicate, projection=['id'])) == expected + + +def test_unmasked_conjunct_is_pushed_to_split_reads(table, monkeypatch): + authorize(table, rules()) + pushed = [] + create_split_read = TableRead._create_split_read + + def recording(self, *args, **kwargs): + split_read = create_split_read(self, *args, **kwargs) + pushed.append(split_read.predicate) + return split_read + + monkeypatch.setattr(TableRead, '_create_split_read', recording) + predicates = PredicateBuilder(table.fields) + unmasked = predicates.greater_than('id', 1) + for combine, expected_push, expected_ids in [ + (PredicateBuilder.and_predicates, unmasked, [2, 3]), + (PredicateBuilder.or_predicates, None, [1, 2, 3])]: + del pushed[:] + predicate = combine([unmasked, predicates.is_null('secret')]) + assert sorted(row['id'] for row in read(table, predicate, projection=['id'])) == expected_ids + assert pushed and all(p == expected_push for p in pushed) + + +def test_row_authorization_precedes_mask_and_user_filter(table): + row_filter = json.dumps({'kind': 'LEAF', 'transform': field_ref('secret', 1), + 'function': 'EQUAL', 'literals': ['beta']}) + authorize(table, rules(row_filter=row_filter)) + predicate = PredicateBuilder(table.fields).is_null('secret') + assert read(table, predicate, projection=['id', 'secret']) == [{'id': 2, 'secret': None}] + + +def test_unmasked_conjunct_keeps_its_pushdown(table): + authorize(table, rules()) + predicates = PredicateBuilder(table.fields) + unmasked = predicates.greater_than('id', 1) + predicate = PredicateBuilder.and_predicates([unmasked, predicates.is_null('secret')]) + scan = table.new_read_builder().with_filter(predicate).new_scan() + scan.plan() + assert scan.file_scanner.predicate == unmasked + assert scan.file_scanner.auth_has_non_partition_filter + + +def test_same_scan_refreshes_masked_filter_rules(table): + state = [rules()] + table.catalog_environment.table_query_auth = lambda options, identifier: lambda select: state[0] + predicate = PredicateBuilder(table.fields).is_null('secret') + builder = table.new_read_builder().with_filter(predicate) + scan = builder.new_scan() + reader = builder.new_read() + for auth, ids in [(rules(), [1, 2, 3]), (TableQueryAuthResult(None, None), []), (rules(), [1, 2, 3])]: + state[0] = auth + rows = arrow_rows(reader.to_arrow(scan.plan().splits())) + assert sorted(row['id'] for row in rows) == ids + + +def test_explain_keeps_masked_filter_candidates(table): + authorize(table, rules()) + predicate = PredicateBuilder(table.fields).is_null('secret') + builder = table.new_read_builder().with_filter(predicate) + plan, _ = builder.new_scan().scan_with_stats() + rows = arrow_rows(builder.new_read().to_arrow(plan.splits())) + assert sorted(row['id'] for row in rows) == [1, 2, 3] + + +@pytest.mark.parametrize('table', ['partitioned-append', 'partitioned-primary-key'], indirect=True) +def test_explicit_partition_filter_on_masked_key_is_rejected(table): + authorize(table, rules()) + predicate = PredicateBuilder(table.fields).equal('secret', 'alpha') + scan = table.new_read_builder().with_partition_filter(predicate).new_scan() + with pytest.raises(ValueError, match='partition filters cannot reference masked columns'): + scan.plan() + + +@pytest.mark.parametrize('phase', ['initial', 'follow_up', 'catch_up']) +def test_streaming_user_filter_runs_after_masking(table, phase): + authorize(table, rules(field_ref('salt', 2))) + predicate = PredicateBuilder(table.fields).equal('secret', 'visible') + builder = (table.new_stream_read_builder().with_filter(predicate) + .with_projection(['id']).with_include_row_kind(False)) + assert stream_read(table, builder, phase) == [{'id': 2}] + + +def test_row_predicate_indices_follow_the_widened_read_type(table): + authorize(table, rules(field_ref('salt', 2))) + predicate = PredicateBuilder(table.fields).startswith('secret', 'vis') + assert read(table, predicate, projection=['salt', 'id']) == [{'salt': 'visible', 'id': 2}] + + +def test_write_scan_restores_filter_after_rules_are_removed(table): + state = [rules()] + table.catalog_environment.table_query_auth = lambda options, identifier: lambda select: state[0] + predicate = PredicateBuilder(table.fields).is_null('secret') + builder = table.new_read_builder().with_filter(predicate) + scan = builder.new_scan() + scan.plan() + state[0] = None + splits = scan.plan_for_write().splits() + assert scan.file_scanner.predicate == predicate + assert not scan.file_scanner.auth_has_non_partition_filter + assert arrow_rows(builder.new_read().to_arrow(splits)) == [] + + +def test_zero_column_projection_preserves_filtered_row_count(table): + authorize(table, rules()) + predicate = PredicateBuilder(table.fields).is_null('secret') + builder = table.new_read_builder().with_filter(predicate).with_projection([]) + result = builder.new_read().to_arrow(builder.new_scan().plan().splits()) + assert result.num_columns == 0 + assert result.num_rows == 3 + + +@pytest.mark.parametrize('table', ['primary-key', 'partitioned-primary-key'], indirect=True) +@pytest.mark.parametrize('include_row_kind', [False, True]) +@pytest.mark.parametrize('phase', ['initial', 'follow_up']) +def test_streaming_string_filter_preserves_row_kind(table, include_row_kind, phase): + authorize(table, rules(field_ref('salt', 2))) + predicates = PredicateBuilder(table.fields) + predicate = PredicateBuilder.and_predicates([ + predicates.startswith('secret', 'vis'), predicates.equal('id', 2)]) + builder = (table.new_stream_read_builder().with_filter(predicate) + .with_projection(['id']).with_include_row_kind(include_row_kind)) + expected = {'id': 2} + if include_row_kind: + expected['_row_kind'] = '+I' + assert stream_read(table, builder, phase) == [expected] + + +@pytest.mark.parametrize('table', ['partitioned-append'], indirect=True) +def test_chunk_shuffle_rejects_post_mask_partition_filter(table): + authorize(table, rules(field_ref('salt', 2))) + predicate = PredicateBuilder(table.fields).equal('secret', 'visible') + scan = table.new_read_builder().with_filter(predicate).new_scan().with_chunk_shuffle(1, 1) + with pytest.raises(ValueError, match='chunk_shuffle.*authorization'): + scan.plan() + + +@pytest.mark.parametrize('table', ['partitioned-append'], indirect=True) +@pytest.mark.parametrize('with_partition_filter', [False, True]) +def test_chunk_shuffle_keeps_safe_authorized_chunks(table, with_partition_filter): + authorize(table, TableQueryAuthResult(None, {'salt': json.dumps({'name': 'NULL'})})) + builder = table.new_read_builder() + if with_partition_filter: + builder.with_filter(PredicateBuilder(table.fields).equal('secret', 'alpha')) + splits = builder.new_scan().with_chunk_shuffle(1, 1).plan().splits() + chunks = [arrow_rows(builder.new_read().to_arrow([split])) for split in splits] + assert all(len(chunk) == 1 for chunk in chunks) + assert sorted(row['id'] for chunk in chunks for row in chunk) == ([1] if with_partition_filter else [1, 2, 3]) + assert all(row['salt'] is None for chunk in chunks for row in chunk) + + +def table_from_rows(tmp_path, fields, rows, primary_keys=None): + catalog = CatalogFactory.create({'warehouse': str(tmp_path / 'warehouse')}) + catalog.create_database('db', True) + options = {'file.format': 'parquet', 'write.native.enabled': 'false'} + if primary_keys: + options['bucket'] = '1' + catalog.create_table('db.t', Schema.from_pyarrow_schema( + fields, primary_keys=primary_keys or [], options=options), False) + table = catalog.get_table('db.t') + for values in [rows, rows[:1]] if primary_keys else [rows]: + builder = table.new_batch_write_builder() + writer = builder.new_write() + try: + writer.write_arrow(pa.Table.from_pydict( + {name: [row[name] for row in values] for name in fields.names}, schema=fields)) + builder.new_commit().commit(writer.prepare_commit()) + finally: + writer.close() + return catalog.get_table('db.t').copy({'query-auth.enabled': 'true'}) + + +def test_post_mask_row_filter_preserves_null_maps(tmp_path): + fields = pa.schema([('id', pa.int32()), ('s', pa.string()), ('m', pa.map_(pa.string(), pa.int64()))]) + table = table_from_rows(tmp_path, fields, [ + {'id': 1, 's': 'a', 'm': None}, {'id': 2, 's': 'a', 'm': []}, + {'id': 3, 's': 'b', 'm': None}, {'id': 4, 's': 'a', 'm': [('key', 1)]}, + ], primary_keys=['id']) + authorize(table, TableQueryAuthResult(None, {'s': json.dumps({ + 'name': 'UPPER', 'inputs': [{'index': 1, 'name': 's', 'type': 'STRING'}]})})) + predicates = PredicateBuilder(table.fields) + predicate = PredicateBuilder.and_predicates([predicates.startswith('s', 'A'), predicates.is_null('m')]) + assert read(table, predicate, projection=['id', 'm']) == [{'id': 1, 'm': None}] + + +@pytest.mark.parametrize('method', ['arrow', 'iterator']) +@pytest.mark.parametrize('compound', [None, 'or-null', 'and-string']) +def test_post_mask_filter_uses_declared_predicate_types(tmp_path, method, compound): + fields = pa.schema([('id', pa.int32()), ('score', pa.int32()), ('label', pa.string())]) + table = table_from_rows(tmp_path, fields, [ + {'id': 1, 'score': 100, 'label': 'keep'}, {'id': 2, 'score': 10, 'label': 'keep'}, + {'id': 10, 'score': 2, 'label': 'keep'}, {'id': 100, 'score': 1, 'label': 'keep'}]) + authorize(table, TableQueryAuthResult(None, {'id': json.dumps({ + 'name': 'CAST', 'fieldRef': {'index': 1, 'name': 'score', 'type': 'INT'}, 'type': 'STRING'})})) + predicates = PredicateBuilder(table.fields) + predicate = predicates.greater_than('id', 2) + if compound == 'or-null': + predicate = PredicateBuilder.or_predicates([predicate, predicates.not_equal('score', None)]) + elif compound == 'and-string': + predicate = PredicateBuilder.and_predicates([predicate, predicates.startswith('label', 'keep')]) + result = read(table, predicate, projection=['id'], method=method) + assert result == ([('100',), ('10',)] if method == 'iterator' else [{'id': 100}, {'id': 10}]) + + +@pytest.mark.parametrize('condition,pattern', [ + ('equal', '100'), ('startswith', '1'), ('endswith', '00'), ('contains', '0'), ('like', '1_0')]) +@pytest.mark.parametrize('with_equality', [False, True]) +@pytest.mark.parametrize('method', ['arrow', 'batches']) +def test_post_mask_string_filter_uses_declared_type(tmp_path, condition, pattern, with_equality, method): + fields = pa.schema([('id', pa.int32()), ('label', pa.string()), ('score', pa.int32())]) + table = table_from_rows(tmp_path, fields, [ + {'id': 1, 'label': 'hidden', 'score': 100}, {'id': 2, 'label': '100', 'score': 234}, + {'id': 3, 'label': '100', 'score': None}]) + authorize(table, TableQueryAuthResult(None, {'label': json.dumps({ + 'name': 'FIELD_REF', 'fieldRef': {'index': 2, 'name': 'score', 'type': 'INT'}})})) + predicates = PredicateBuilder(table.fields) + predicate = getattr(predicates, condition)('label', pattern) + if with_equality: + predicate = PredicateBuilder.and_predicates([predicate, predicates.equal('label', '100')]) + assert read(table, predicate, projection=['id', 'label'], method=method) == [{'id': 1, 'label': '100'}] + + +@pytest.mark.parametrize('table', ['primary-key', 'partitioned-primary-key'], indirect=True) +def test_streaming_post_mask_filter_keeps_delete_row_kind(table): + builder = table.copy({'query-auth.enabled': 'false'}).new_batch_write_builder() + writer = builder.new_write() + try: + writer.write_row(GenericRow([2, 'beta', 'visible'], table.fields, RowKind.DELETE)) + builder.new_commit().commit(writer.prepare_commit()) + finally: + writer.close() + authorize(table, rules(field_ref('salt', 2))) + predicate = PredicateBuilder(table.fields).equal('secret', 'visible') + read_builder = (table.new_stream_read_builder().with_filter(predicate) + .with_projection(['id']).with_include_row_kind(True)) + delete_snapshot = table.snapshot_manager().get_latest_snapshot().id + rows = stream_read(table, read_builder, 'follow_up', follow_up_snapshot=delete_snapshot) + assert rows == [{'_row_kind': '-D', 'id': 2}] + + +@pytest.mark.parametrize('method', ['arrow', 'iterator']) +@pytest.mark.parametrize('with_other_arm', [False, True]) +def test_not_in_with_null_skips_masked_value_conversion(tmp_path, method, with_other_arm): + fields = pa.schema([('id', pa.int32()), ('salt', pa.string())]) + table = table_from_rows(tmp_path, fields, [{'id': 1, 'salt': 'hidden'}, {'id': 2, 'salt': 'kept'}]) + # masked INT values are strings that cannot be cast back + authorize(table, TableQueryAuthResult(None, {'id': json.dumps(field_ref('salt', 1))})) + predicates = PredicateBuilder(table.fields) + predicate = predicates.is_not_in('id', [2, None]) + if with_other_arm: + predicate = PredicateBuilder.or_predicates([predicate, predicates.equal('salt', 'kept')]) + rows = read(table, predicate, projection=['salt'], method=method) + if not with_other_arm: + assert rows == [] + else: + assert rows == ([('kept',)] if method == 'iterator' else [{'salt': 'kept'}]) + + +def test_nested_filter_alias_authorizes_its_source_column(tmp_path): + fields = pa.schema([('id', pa.int32()), ('s', pa.struct([('a', pa.string()), ('b', pa.int32())]))]) + table = table_from_rows(tmp_path, fields, [ + {'id': 1, 's': {'a': 'keep', 'b': 1}}, {'id': 2, 's': {'a': 'drop', 'b': 2}}]) + selected = [] + + def auth(select): + selected.append(select) + return TableQueryAuthResult(None, None) + + table.catalog_environment.table_query_auth = lambda options, identifier: auth + builder = table.new_read_builder().with_projection(['id', 's.a']) + builder.with_filter(builder.new_predicate_builder().equal('s_a', 'keep')) + assert arrow_rows(builder.new_read().to_arrow(builder.new_scan().plan().splits())) == [{'id': 1, 's_a': 'keep'}] + assert selected == [['id', 's']] diff --git a/paimon-python/pypaimon/tests/table_query_auth_test.py b/paimon-python/pypaimon/tests/table_query_auth_test.py index 1c39f6c2cf34..56ec0ddb942b 100644 --- a/paimon-python/pypaimon/tests/table_query_auth_test.py +++ b/paimon-python/pypaimon/tests/table_query_auth_test.py @@ -77,6 +77,7 @@ def __init__(self, identifier=None): class _FakeScan: def __init__(self, plan, identifier=None, auth_result=None): self._plan = plan + self._query_auth_fn = None self.table = _FakeTable(identifier) self.file_scanner = _FakeFileScanner(plan) self._auth_result = auth_result From 00ac6b2b0b7fa07cdce7eb3ba76cdd5b121e96a6 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Tue, 13 Oct 2026 20:48:19 -0700 Subject: [PATCH 2/5] [python] Filter authorized reads exactly before LIMIT Split reads can apply a query predicate loosely: a nested projection drops OR arms on unprojected columns, and merged primary-key rows compare NULL literals as Python values. Evaluate every query predicate again before LIMIT, keeping its columns until the final projection, and accept authorization results without column masks. --- paimon-python/pypaimon/read/table_read.py | 15 ++++----- .../pypaimon/tests/query_auth_filter_test.py | 32 +++++++++++++++++++ 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index 29545e4cba13..c041b49be492 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -1835,9 +1835,10 @@ def __authed_reader(self, split, auth_result, blob_parallelism=1, "MAP-key projection with query authorization is not supported") table_fields = self.table.fields read_fields = self._adapter_read_type - masked_fields = set(auth_result.column_masking or {}) - post_filter = (self.predicate if self.predicate is not None - and predicate_field_names(self.predicate) & masked_fields else None) + masks = auth_result.column_masking or {} + masked_fields = set(masks) + # Split reads may filter loosely (merged rows, nested leaves), so filter exactly before LIMIT. + post_filter = self.predicate filter_read_fields = self._with_predicate_extra_fields(read_fields) if post_filter is not None else read_fields from pypaimon.read.table_scan import latest_auth_fields, validate_auth_rules # a split can carry rules without having been planned by a validating scan @@ -1863,7 +1864,7 @@ def __authed_reader(self, split, auth_result, blob_parallelism=1, # Only inputs of the filter and its masks are read eagerly. filter_names = predicate_field_names(post_filter) filter_masks = compile_auth_masking_rules( - {name: rule for name, rule in auth_result.column_masking.items() if name in filter_names}, + {name: rule for name, rule in masks.items() if name in filter_names}, effective_read_type) eager_blob_fields.update(filter_names - masked_fields) for transform in filter_masks.values(): @@ -1889,8 +1890,7 @@ def __authed_reader(self, split, auth_result, blob_parallelism=1, eager_blob_fields=eager_blob_fields, post_filter_after_inline=post_filter_after_inline, # push down conjuncts without masked columns - predicate=(self.predicate if post_filter is None - else exclude_predicate_with_fields(post_filter, masked_fields)), + predicate=exclude_predicate_with_fields(post_filter, masked_fields), defer_blob_resolution=post_filter is not None, ) reader = split_read.create_reader() @@ -1926,8 +1926,7 @@ def __authed_reader(self, split, auth_result, blob_parallelism=1, if post_filter is not None and isinstance(split_read, DataEvolutionSplitRead): # Fetch only payloads output unmasked or read by a mask. needed = {f.name for f in read_fields} - masked_fields - for transform in compile_auth_masking_rules( - auth_result.column_masking, effective_read_type).values(): + for transform in compile_auth_masking_rules(masks, effective_read_type).values(): needed.update(_collect_all_field_refs_from_transform(transform)) reader = split_read._resolve_deferred_blobs( reader, {f.name for f in effective_read_type} - needed) diff --git a/paimon-python/pypaimon/tests/query_auth_filter_test.py b/paimon-python/pypaimon/tests/query_auth_filter_test.py index 4e0c98357638..21c30d38822f 100644 --- a/paimon-python/pypaimon/tests/query_auth_filter_test.py +++ b/paimon-python/pypaimon/tests/query_auth_filter_test.py @@ -229,6 +229,14 @@ def test_row_authorization_precedes_mask_and_user_filter(table): assert read(table, predicate, projection=['id', 'secret']) == [{'id': 2, 'secret': None}] +def test_row_only_authorization_with_user_filter(table): + row_filter = json.dumps({'kind': 'LEAF', 'transform': field_ref('secret', 1), + 'function': 'EQUAL', 'literals': ['beta']}) + authorize(table, TableQueryAuthResult([row_filter], None)) + predicate = PredicateBuilder(table.fields).greater_than('id', 1) + assert read(table, predicate, projection=['id', 'secret']) == [{'id': 2, 'secret': 'beta'}] + + def test_unmasked_conjunct_keeps_its_pushdown(table): authorize(table, rules()) predicates = PredicateBuilder(table.fields) @@ -470,3 +478,27 @@ def auth(select): builder.with_filter(builder.new_predicate_builder().equal('s_a', 'keep')) assert arrow_rows(builder.new_read().to_arrow(builder.new_scan().plan().splits())) == [{'id': 1, 's_a': 'keep'}] assert selected == [['id', 's']] + + +def test_unmasked_filter_keeps_unprojected_columns_with_nested_projection(tmp_path): + fields = pa.schema([('id', pa.int32()), ('s', pa.struct([('a', pa.string()), ('b', pa.int32())])), + ('secret', pa.string())]) + table = table_from_rows(tmp_path, fields, [ + {'id': 1, 's': {'a': 'keep', 'b': 1}, 'secret': 'x'}, {'id': 2, 's': {'a': 'drop', 'b': 2}, 'secret': 'y'}, + {'id': 3, 's': {'a': 'drop', 'b': 3}, 'secret': 'z'}]) + authorize(table, rules()) + builder = table.new_read_builder().with_projection(['s.a', 'secret']) + # the nested split read cannot evaluate this OR, since 'id' is not projected + builder.with_filter(PredicateBuilder.or_predicates([ + builder.new_predicate_builder().equal('s_a', 'keep'), PredicateBuilder(table.fields).equal('id', 2)])) + rows = arrow_rows(builder.new_read().to_arrow(builder.new_scan().plan().splits())) + assert sorted(row['s_a'] for row in rows) == ['drop', 'keep'] + + +def test_unmasked_null_comparison_after_merge_matches_no_rows(tmp_path): + fields = pa.schema([pa.field('id', pa.int32(), nullable=False), ('score', pa.int32()), ('secret', pa.string())]) + table = table_from_rows(tmp_path, fields, [ + {'id': 1, 'score': 10, 'secret': 'x'}, {'id': 2, 'score': None, 'secret': 'y'}], primary_keys=['id']) + authorize(table, rules()) + predicate = PredicateBuilder(table.fields).not_equal('score', None) + assert read(table, predicate, projection=['id', 'score']) == [] From a7c8f299ed892a151529d250bb7afa5e5f614da1 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Tue, 13 Oct 2026 22:20:38 -0700 Subject: [PATCH 3/5] [python] Compare masked values per row when some miss the declared type Casting a masked comparison column to its declared type fails for values such as 'REDACTED' on an INT column, which broke OR predicates whose other arm already matched. Evaluate such batches row by row instead: values that fit the declared type compare as it, the others keep their masked value. --- .../read/reader/filter_record_batch_reader.py | 28 +++++++++++++++++-- .../pypaimon/tests/query_auth_filter_test.py | 16 +++++++++++ 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py b/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py index 7529bf0b160c..49c223335c9c 100644 --- a/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py +++ b/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py @@ -49,6 +49,18 @@ def _comparison_fields(predicate): return set() +def _cast_values(rows, types): + # Values that fit the declared type compare as it; the others keep their masked value. + for row in rows: + row = list(row) + for index, data_type in types.items(): + try: + row[index] = pa.array([row[index]]).cast(data_type)[0].as_py() + except (pa.ArrowInvalid, pa.ArrowNotImplementedError, pa.ArrowTypeError): + pass + yield tuple(row) + + class FilterRecordBatchReader(RecordBatchReader): """ Wraps a RecordBatchReader and filters each batch by predicate. @@ -92,7 +104,11 @@ def read_arrow_batch(self) -> Optional[pa.RecordBatch]: def _filter_batch(self, batch: pa.RecordBatch) -> Optional[pa.RecordBatch]: view = batch if self._filter_view is None else self._filter_view(batch) - filter_batch = self._coerce_comparison_types(view) + try: + filter_batch = self._coerce_comparison_types(view) + except (pa.ArrowInvalid, pa.ArrowNotImplementedError): + # Some masked values do not fit the declared type; convert the others per row. + return self._filter_batch_by_row(batch, view, cast_values=True) if not self._use_arrow_filter: return self._filter_batch_by_row(batch, filter_batch) expr = self.predicate.to_arrow(filter_batch.schema) @@ -139,7 +155,8 @@ def _coerce_comparison_types(self, batch): return batch return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields, metadata=batch.schema.metadata)) - def _filter_batch_by_row(self, batch: pa.RecordBatch, filter_batch=None) -> Optional[pa.RecordBatch]: + def _filter_batch_by_row(self, batch: pa.RecordBatch, filter_batch=None, + cast_values=False) -> Optional[pa.RecordBatch]: filter_batch = batch if filter_batch is None else filter_batch # Predicate indices address logical fields. A record-reader adapter # may prepend _row_kind, which is metadata rather than an operand. @@ -147,8 +164,13 @@ def _filter_batch_by_row(self, batch: pa.RecordBatch, filter_batch=None) -> Opti [field.name for field in self.schema_fields] if self.schema_fields is not None else [name for name in filter_batch.schema.names if name != "_row_kind"]) df = polars.from_arrow(_polars_arrow_table(filter_batch).select(columns)) + rows = self._iter_df_rows(df) + if cast_values: + types = {columns.index(name): data_type for name, data_type in self._comparison_types.items() + if name in columns and filter_batch.column(name).type != data_type} + rows = _cast_values(rows, types) iterator = InternalRowWrapperIterator( - self._iter_df_rows(df), + rows, df.width, self.file_io, self.blob_field_indices, diff --git a/paimon-python/pypaimon/tests/query_auth_filter_test.py b/paimon-python/pypaimon/tests/query_auth_filter_test.py index 21c30d38822f..00b3756ca1e5 100644 --- a/paimon-python/pypaimon/tests/query_auth_filter_test.py +++ b/paimon-python/pypaimon/tests/query_auth_filter_test.py @@ -463,6 +463,22 @@ def test_not_in_with_null_skips_masked_value_conversion(tmp_path, method, with_o assert rows == ([('kept',)] if method == 'iterator' else [{'salt': 'kept'}]) +@pytest.mark.parametrize('method', ['arrow', 'iterator']) +@pytest.mark.parametrize('first_arm', ['startswith', 'equal']) +def test_uncastable_masked_value_keeps_or_short_circuit(tmp_path, method, first_arm): + fields = pa.schema([('id', pa.int32()), ('secret', pa.int32()), ('public', pa.string())]) + # one batch: '100' still compares as INT, 'REDACTED' only matches the first arm + table = table_from_rows(tmp_path, fields, [ + {'id': 1, 'secret': 1, 'public': '100'}, {'id': 2, 'secret': 2, 'public': 'REDACTED'}]) + authorize(table, TableQueryAuthResult(None, {'secret': json.dumps(field_ref('public', 2))})) + predicates = PredicateBuilder(table.fields) + public = (predicates.startswith('public', 'R') if first_arm == 'startswith' + else predicates.equal('public', 'REDACTED')) + predicate = PredicateBuilder.or_predicates([public, predicates.greater_than('secret', 50)]) + rows = read(table, predicate, projection=['id'], method=method) + assert rows == ([(1,), (2,)] if method == 'iterator' else [{'id': 1}, {'id': 2}]) + + def test_nested_filter_alias_authorizes_its_source_column(tmp_path): fields = pa.schema([('id', pa.int32()), ('s', pa.struct([('a', pa.string()), ('b', pa.int32())]))]) table = table_from_rows(tmp_path, fields, [ From c64689f81a5878b4b6c99288deefe9d1ace02ca9 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Tue, 13 Oct 2026 23:23:08 -0700 Subject: [PATCH 4/5] [python] Treat masked values outside the declared type as NULL in filters Falling back to the Python row predicate for such batches treated `x != NULL` as true and raised on ordering comparisons. Cast them like SQL TRY_CAST instead and keep the evaluation path; null checks see the same values. Only columns whose type a mask changed are cast. --- .../read/reader/filter_record_batch_reader.py | 61 +++++++++---------- .../pypaimon/tests/query_auth_filter_test.py | 31 ++++++---- 2 files changed, 50 insertions(+), 42 deletions(-) diff --git a/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py b/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py index 49c223335c9c..878ccaf8a4c0 100644 --- a/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py +++ b/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py @@ -19,6 +19,7 @@ from typing import Callable, List, Optional import pyarrow as pa +import pyarrow.compute as pc import pyarrow.dataset as ds import polars @@ -37,6 +38,9 @@ def _comparison_fields(predicate): if predicate.method in ('and', 'or'): return set().union(*(_comparison_fields(child) for child in predicate.literals)) + if predicate.method in ('isNull', 'isNotNull'): + # Null checks see the cast values, so they agree with comparisons. + return {predicate.field} literals = predicate.literals or [] # NOT IN with a NULL literal never matches, so skip the cast. if predicate.method == 'notIn' and any(value is None for value in literals): @@ -49,16 +53,21 @@ def _comparison_fields(predicate): return set() -def _cast_values(rows, types): - # Values that fit the declared type compare as it; the others keep their masked value. - for row in rows: - row = list(row) - for index, data_type in types.items(): +def _try_cast(column, data_type): + # Like SQL TRY_CAST: values that do not fit the declared type compare as NULL. + try: + return column.cast(data_type) + except pa.ArrowNotImplementedError: + return pa.nulls(len(column), data_type) + except pa.ArrowInvalid: + values = column.unique() + cast = [] + for index in range(len(values)): try: - row[index] = pa.array([row[index]]).cast(data_type)[0].as_py() - except (pa.ArrowInvalid, pa.ArrowNotImplementedError, pa.ArrowTypeError): - pass - yield tuple(row) + cast.append(values.slice(index, 1).cast(data_type)) + except pa.ArrowInvalid: + cast.append(pa.nulls(1, data_type)) + return pa.concat_arrays(cast).take(pc.index_in(column, value_set=values)) class FilterRecordBatchReader(RecordBatchReader): @@ -104,11 +113,7 @@ def read_arrow_batch(self) -> Optional[pa.RecordBatch]: def _filter_batch(self, batch: pa.RecordBatch) -> Optional[pa.RecordBatch]: view = batch if self._filter_view is None else self._filter_view(batch) - try: - filter_batch = self._coerce_comparison_types(view) - except (pa.ArrowInvalid, pa.ArrowNotImplementedError): - # Some masked values do not fit the declared type; convert the others per row. - return self._filter_batch_by_row(batch, view, cast_values=True) + filter_batch = self._coerce_comparison_types(view, batch) if not self._use_arrow_filter: return self._filter_batch_by_row(batch, filter_batch) expr = self.predicate.to_arrow(filter_batch.schema) @@ -140,23 +145,22 @@ def _filter_batch(self, batch: pa.RecordBatch) -> Optional[pa.RecordBatch]: schema=result.schema, ) - def _coerce_comparison_types(self, batch): - # Cast for comparison only; the output is unchanged. + def _coerce_comparison_types(self, view, batch): + # Cast columns a mask retyped, for the filter only; the output is unchanged. columns = None for name, data_type in self._comparison_types.items(): - column = batch.column(name) - if column.type != data_type: + column = view.column(name) + if column.type != data_type and column.type != batch.column(name).type: if columns is None: - columns, fields = list(batch.columns), list(batch.schema) - index = batch.schema.get_field_index(name) - columns[index] = column.cast(data_type) + columns, fields = list(view.columns), list(view.schema) + index = view.schema.get_field_index(name) + columns[index] = _try_cast(column, data_type) fields[index] = fields[index].with_type(data_type) if columns is None: - return batch - return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields, metadata=batch.schema.metadata)) + return view + return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields, metadata=view.schema.metadata)) - def _filter_batch_by_row(self, batch: pa.RecordBatch, filter_batch=None, - cast_values=False) -> Optional[pa.RecordBatch]: + def _filter_batch_by_row(self, batch: pa.RecordBatch, filter_batch=None) -> Optional[pa.RecordBatch]: filter_batch = batch if filter_batch is None else filter_batch # Predicate indices address logical fields. A record-reader adapter # may prepend _row_kind, which is metadata rather than an operand. @@ -164,13 +168,8 @@ def _filter_batch_by_row(self, batch: pa.RecordBatch, filter_batch=None, [field.name for field in self.schema_fields] if self.schema_fields is not None else [name for name in filter_batch.schema.names if name != "_row_kind"]) df = polars.from_arrow(_polars_arrow_table(filter_batch).select(columns)) - rows = self._iter_df_rows(df) - if cast_values: - types = {columns.index(name): data_type for name, data_type in self._comparison_types.items() - if name in columns and filter_batch.column(name).type != data_type} - rows = _cast_values(rows, types) iterator = InternalRowWrapperIterator( - rows, + self._iter_df_rows(df), df.width, self.file_io, self.blob_field_indices, diff --git a/paimon-python/pypaimon/tests/query_auth_filter_test.py b/paimon-python/pypaimon/tests/query_auth_filter_test.py index 00b3756ca1e5..626f878c5227 100644 --- a/paimon-python/pypaimon/tests/query_auth_filter_test.py +++ b/paimon-python/pypaimon/tests/query_auth_filter_test.py @@ -464,19 +464,28 @@ def test_not_in_with_null_skips_masked_value_conversion(tmp_path, method, with_o @pytest.mark.parametrize('method', ['arrow', 'iterator']) -@pytest.mark.parametrize('first_arm', ['startswith', 'equal']) -def test_uncastable_masked_value_keeps_or_short_circuit(tmp_path, method, first_arm): - fields = pa.schema([('id', pa.int32()), ('secret', pa.int32()), ('public', pa.string())]) - # one batch: '100' still compares as INT, 'REDACTED' only matches the first arm +@pytest.mark.parametrize('case,expected', [ + ('startswith', [1, 2]), ('equal', [1, 2]), ('null-literal', [3]), ('is-null', [2, 4])]) +def test_masked_values_outside_declared_type_compare_as_null(tmp_path, method, case, expected): + fields = pa.schema([('id', pa.int32()), ('score', pa.int32()), ('secret', pa.int32()), ('public', pa.string())]) + # one batch: secret reads '100' and '7' as INT, 'REDACTED' and 'ALLOW' as NULL table = table_from_rows(tmp_path, fields, [ - {'id': 1, 'secret': 1, 'public': '100'}, {'id': 2, 'secret': 2, 'public': 'REDACTED'}]) - authorize(table, TableQueryAuthResult(None, {'secret': json.dumps(field_ref('public', 2))})) - predicates = PredicateBuilder(table.fields) - public = (predicates.startswith('public', 'R') if first_arm == 'startswith' - else predicates.equal('public', 'REDACTED')) - predicate = PredicateBuilder.or_predicates([public, predicates.greater_than('secret', 50)]) + {'id': 1, 'score': None, 'secret': 1, 'public': '100'}, + {'id': 2, 'score': 1, 'secret': 2, 'public': 'REDACTED'}, + {'id': 3, 'score': None, 'secret': 3, 'public': '7'}, + {'id': 4, 'score': 4, 'secret': 4, 'public': 'ALLOW'}]) + authorize(table, TableQueryAuthResult(None, {'secret': json.dumps(field_ref('public', 3))})) + p = PredicateBuilder(table.fields) + predicate = { + 'startswith': PredicateBuilder.or_predicates([p.startswith('public', 'R'), p.greater_than('secret', 50)]), + 'equal': PredicateBuilder.or_predicates([p.equal('public', 'REDACTED'), p.greater_than('secret', 50)]), + 'null-literal': PredicateBuilder.or_predicates([ + p.not_equal('score', None), p.less_than('secret', 50), + PredicateBuilder.and_predicates([p.equal('public', 'ALLOW'), p.greater_than('secret', 50)])]), + 'is-null': p.is_null('secret'), + }[case] rows = read(table, predicate, projection=['id'], method=method) - assert rows == ([(1,), (2,)] if method == 'iterator' else [{'id': 1}, {'id': 2}]) + assert rows == ([(i,) for i in expected] if method == 'iterator' else [{'id': i} for i in expected]) def test_nested_filter_alias_authorizes_its_source_column(tmp_path): From 65c747350dfdfe6d768ef641c83cf5ea41b1594f Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Wed, 14 Oct 2026 00:19:06 -0700 Subject: [PATCH 5/5] [python] Cast masked filter values only across type families The TRY_CAST cast turned valid but lossy values into NULL, such as 1.5 on an INT column, a BIGINT outside the INT range or sub-millisecond timestamps, which master compares exactly. It also failed on ARRAY, ROW and MAP masks, even for null checks, because Arrow cannot hash nested values. Let Arrow compare numbers and times of one kind directly, leave nested columns to null checks as masked, and filter on the masked values when Arrow has no cast between the two types. --- .../read/reader/filter_record_batch_reader.py | 17 ++++++++-- .../pypaimon/tests/query_auth_filter_test.py | 33 +++++++++++++++++++ 2 files changed, 47 insertions(+), 3 deletions(-) diff --git a/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py b/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py index 878ccaf8a4c0..2dccc1c995bc 100644 --- a/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py +++ b/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py @@ -58,7 +58,7 @@ def _try_cast(column, data_type): try: return column.cast(data_type) except pa.ArrowNotImplementedError: - return pa.nulls(len(column), data_type) + return column # Arrow cannot convert these types; filter on the masked values except pa.ArrowInvalid: values = column.unique() cast = [] @@ -70,6 +70,16 @@ def _try_cast(column, data_type): return pa.concat_arrays(cast).take(pc.index_in(column, value_set=values)) +def _needs_cast(source, target): + # Arrow compares numbers, and times of one kind, directly, and nested columns only take + # null checks; casting those could only lose precision or nullness. + numeric = (pa.types.is_integer, pa.types.is_floating, pa.types.is_decimal) + return not (pa.types.is_nested(target) + or any(check(source) for check in numeric) and any(check(target) for check in numeric) + or any(check(source) and check(target) + for check in (pa.types.is_timestamp, pa.types.is_date, pa.types.is_time))) + + class FilterRecordBatchReader(RecordBatchReader): """ Wraps a RecordBatchReader and filters each batch by predicate. @@ -150,12 +160,13 @@ def _coerce_comparison_types(self, view, batch): columns = None for name, data_type in self._comparison_types.items(): column = view.column(name) - if column.type != data_type and column.type != batch.column(name).type: + if (column.type != data_type and column.type != batch.column(name).type + and _needs_cast(column.type, data_type)): if columns is None: columns, fields = list(view.columns), list(view.schema) index = view.schema.get_field_index(name) columns[index] = _try_cast(column, data_type) - fields[index] = fields[index].with_type(data_type) + fields[index] = fields[index].with_type(columns[index].type) if columns is None: return view return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields, metadata=view.schema.metadata)) diff --git a/paimon-python/pypaimon/tests/query_auth_filter_test.py b/paimon-python/pypaimon/tests/query_auth_filter_test.py index 626f878c5227..28cd2e627918 100644 --- a/paimon-python/pypaimon/tests/query_auth_filter_test.py +++ b/paimon-python/pypaimon/tests/query_auth_filter_test.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +import datetime import json import pyarrow as pa @@ -488,6 +489,38 @@ def test_masked_values_outside_declared_type_compare_as_null(tmp_path, method, c assert rows == ([(i,) for i in expected] if method == 'iterator' else [{'id': i} for i in expected]) +@pytest.mark.parametrize('kind', ['ARRAY', 'ROW', 'MAP']) +def test_null_checks_on_retyped_nested_masks(tmp_path, kind): + target, source, values, type_name = { + 'ARRAY': (pa.list_(pa.int32()), pa.list_(pa.string()), [['1'], ['redacted'], None, ['3']], 'ARRAY'), + 'ROW': (pa.struct([('a', pa.int32())]), pa.struct([('a', pa.string())]), + [{'a': '1'}, {'a': 'redacted'}, None, {'a': '3'}], 'ROW'), + 'MAP': (pa.map_(pa.string(), pa.int32()), pa.map_(pa.string(), pa.string()), + [[('k', '1')], [('k', 'redacted')], None, [('k', '3')]], 'MAP'), + }[kind] + fields = pa.schema([('id', pa.int32()), ('v', target), ('src', source)]) + table = table_from_rows(tmp_path, fields, [{'id': i + 1, 'v': None, 'src': v} for i, v in enumerate(values)]) + authorize(table, TableQueryAuthResult(None, {'v': json.dumps( + {'name': 'FIELD_REF', 'fieldRef': {'name': 'src', 'index': 2, 'type': type_name}})})) + predicates = PredicateBuilder(table.fields) + assert read(table, predicates.is_null('v'), projection=['id']) == [{'id': 3}] + assert read(table, predicates.is_not_null('v'), projection=['id']) == [{'id': 1}, {'id': 2}, {'id': 4}] + + +@pytest.mark.parametrize('target,source,type_name,values,literal', [ + (pa.int32(), pa.float64(), 'DOUBLE', [1.5, 0.5], 1), + (pa.int32(), pa.int64(), 'BIGINT', [2 ** 40, 1], 5), + (pa.timestamp('ms'), pa.timestamp('us'), 'TIMESTAMP(6)', [1500, 500], + datetime.datetime(1970, 1, 1, microsecond=1000))]) +def test_numeric_and_timestamp_masks_compare_exactly(tmp_path, target, source, type_name, values, literal): + fields = pa.schema([('id', pa.int32()), ('v', target), ('src', source)]) + table = table_from_rows(tmp_path, fields, [{'id': i + 1, 'v': None, 'src': v} for i, v in enumerate(values)]) + authorize(table, TableQueryAuthResult(None, {'v': json.dumps( + {'name': 'FIELD_REF', 'fieldRef': {'name': 'src', 'index': 2, 'type': type_name}})})) + predicate = PredicateBuilder(table.fields).greater_than('v', literal) + assert read(table, predicate, projection=['id']) == [{'id': 1}] + + def test_nested_filter_alias_authorizes_its_source_column(tmp_path): fields = pa.schema([('id', pa.int32()), ('s', pa.struct([('a', pa.string()), ('b', pa.int32())]))]) table = table_from_rows(tmp_path, fields, [