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..2dccc1c995bc 100644 --- a/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py +++ b/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py @@ -16,9 +16,10 @@ # under the License. import logging -from typing import List, Optional +from typing import Callable, List, Optional import pyarrow as pa +import pyarrow.compute as pc import pyarrow.dataset as ds import polars @@ -27,12 +28,58 @@ 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)) + 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): + 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() + + +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 column # Arrow cannot convert these types; filter on the masked values + except pa.ArrowInvalid: + values = column.unique() + cast = [] + for index in range(len(values)): + try: + 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)) + + +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. @@ -46,11 +93,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 +122,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, batch) 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 +155,30 @@ 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, 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 = view.column(name) + 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(columns[index].type) + if columns is None: + 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) -> 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..c041b49be492 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,11 @@ 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 + 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 snapshot = getattr(self, "_auth_schema", None) @@ -1837,24 +1848,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 masks.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 +1885,13 @@ 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=exclude_predicate_with_fields(post_filter, masked_fields), + defer_blob_resolution=post_filter is not None, ) reader = split_read.create_reader() @@ -1897,17 +1913,28 @@ 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(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) + + 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..28cd2e627918 --- /dev/null +++ b/paimon-python/pypaimon/tests/query_auth_filter_test.py @@ -0,0 +1,562 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import datetime +import json + +import 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_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) + 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'}]) + + +@pytest.mark.parametrize('method', ['arrow', 'iterator']) +@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, '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 == ([(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, [ + {'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']] + + +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']) == [] 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