[python] Refine post-mask query filtering for authorized reads

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

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()

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,

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; "

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: