This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 6b84e6a989 [python] Refine post-mask query filtering for authorized
reads (#10523)
6b84e6a989 is described below
commit 6b84e6a989460be11d5b916db078e95e65d631e1
Author: Jiajia Li <[email protected]>
AuthorDate: Fri Oct 16 13:08:02 2026 +0800
[python] Refine post-mask query filtering for authorized reads (#10523)
---
paimon-python/pypaimon/read/query_auth_split.py | 12 +-
.../pypaimon/read/reader/auth_masking_reader.py | 18 +-
.../read/reader/filter_record_batch_reader.py | 102 +++-
.../pypaimon/read/scanner/file_scanner.py | 58 +--
paimon-python/pypaimon/read/split_read.py | 15 +-
paimon-python/pypaimon/read/table_read.py | 81 ++-
paimon-python/pypaimon/read/table_scan.py | 55 +-
.../pypaimon/tests/deferred_blob_resolve_test.py | 118 ++++-
.../pypaimon/tests/query_auth_filter_test.py | 562 +++++++++++++++++++++
.../pypaimon/tests/table_query_auth_test.py | 1 +
10 files changed, 906 insertions(+), 116 deletions(-)
diff --git a/paimon-python/pypaimon/read/query_auth_split.py
b/paimon-python/pypaimon/read/query_auth_split.py
index bfbb1a2799..b172223558 100644
--- a/paimon-python/pypaimon/read/query_auth_split.py
+++ b/paimon-python/pypaimon/read/query_auth_split.py
@@ -68,10 +68,20 @@ class QueryAuthSplit(Split):
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 5c6873052a..c40711ad60 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 @@ class ColumnProjectReader(RecordBatchReader):
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 381536e7c0..2dccc1c995 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.push_down_utils import
predicate_supports_arrow_filter
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 @@ class FilterRecordBatchReader(RecordBatchReader):
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 @@ class FilterRecordBatchReader(RecordBatchReader):
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 @@ class FilterRecordBatchReader(RecordBatchReader):
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 85c8061975..22779bdfc5 100755
--- a/paimon-python/pypaimon/read/scanner/file_scanner.py
+++ b/paimon-python/pypaimon/read/scanner/file_scanner.py
@@ -230,39 +230,15 @@ class FileScanner:
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 @@ class FileScanner:
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).
@@ -730,12 +720,14 @@ class FileScanner:
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 e23db8092d..596c661cf7 100644
--- a/paimon-python/pypaimon/read/split_read.py
+++ b/paimon-python/pypaimon/read/split_read.py
@@ -1393,7 +1393,8 @@ class DataEvolutionSplitRead(SplitRead):
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 @@ class DataEvolutionSplitRead(SplitRead):
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 @@ class DataEvolutionSplitRead(SplitRead):
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 @@ class DataEvolutionSplitRead(SplitRead):
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 6e0b26362f..c041b49be4 100644
--- a/paimon-python/pypaimon/read/table_read.py
+++ b/paimon-python/pypaimon/read/table_read.py
@@ -29,14 +29,16 @@ from pypaimon.casting.schema_evolution import
read_arrow_schema
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 @@ class TableRead:
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 @@ class TableRead:
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 @@ class TableRead:
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 @@ class TableRead:
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 @@ class TableRead:
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 @@ class TableRead:
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 @@ class TableRead:
"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 @@ class TableRead:
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 @@ class TableRead:
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 @@ class TableRead:
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 687d5c8690..c3d4edc504 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 import Predicate
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 @@ class TableScan:
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 @@ class TableScan:
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)
@@ -276,6 +278,8 @@ class TableScan:
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:
@@ -302,13 +306,13 @@ class TableScan:
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
@@ -602,48 +606,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 3c8a1a071a..51ec83c197 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 @@ import pytest
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 @@ class DeferredBlobResolveTest(unittest.TestCase):
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 @@ class DeferredBlobResolveTest(unittest.TestCase):
@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 @@ class DeferredBlobResolveTest(unittest.TestCase):
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 @@ class DeferredBlobResolveTest(unittest.TestCase):
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 0000000000..28cd2e6279
--- /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))]
+
+
[email protected](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
+
+
[email protected]('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)
+
+
[email protected]('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}])
+
+
[email protected]('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']]
+
+
[email protected]('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']]
+
+
[email protected]('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)
+
+
[email protected]('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]
+
+
[email protected]('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()
+
+
[email protected]('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
+
+
[email protected]('table', ['primary-key', 'partitioned-primary-key'],
indirect=True)
[email protected]('include_row_kind', [False, True])
[email protected]('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]
+
+
[email protected]('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()
+
+
[email protected]('table', ['partitioned-append'], indirect=True)
[email protected]('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}]
+
+
[email protected]('method', ['arrow', 'iterator'])
[email protected]('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}])
+
+
[email protected]('condition,pattern', [
+ ('equal', '100'), ('startswith', '1'), ('endswith', '00'), ('contains',
'0'), ('like', '1_0')])
[email protected]('with_equality', [False, True])
[email protected]('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'}]
+
+
[email protected]('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}]
+
+
[email protected]('method', ['arrow', 'iterator'])
[email protected]('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'}])
+
+
[email protected]('method', ['arrow', 'iterator'])
[email protected]('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])
+
+
[email protected]('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<STRING>'),
+ 'ROW': (pa.struct([('a', pa.int32())]), pa.struct([('a',
pa.string())]),
+ [{'a': '1'}, {'a': 'redacted'}, None, {'a': '3'}], 'ROW<a
STRING>'),
+ 'MAP': (pa.map_(pa.string(), pa.int32()), pa.map_(pa.string(),
pa.string()),
+ [[('k', '1')], [('k', 'redacted')], None, [('k', '3')]],
'MAP<STRING, STRING>'),
+ }[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}]
+
+
[email protected]('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 1c39f6c2cf..56ec0ddb94 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 @@ class _FakeTable:
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