[table] Support query-auth row filters and column masking at read time

crates/integrations/datafusion/Cargo.toml

@@ -50,6 +50,8 @@ uuid = { version = "1", features = ["v4"] }
 [dev-dependencies]
 arrow-array = { workspace = true }
 arrow-schema = { workspace = true }
+# for the shared REST catalog mock (crates/paimon/tests/mock_server.rs)
+axum = { version = "0.7", features = ["macros", "tokio", "http1", "http2"] }
 bytes = "1.7.1"
 flate2 = "1"
 paimon-ftindex-core = "0.1.0"

crates/integrations/datafusion/src/merge_into.rs

@@ -1233,12 +1233,14 @@ pub(crate) async fn register_cow_target_table(
         return Ok((false, table_name));
     }
 
-    // Read all files in parallel
-    let read_futures: Vec<_> = file_index
-        .iter()
-        .enumerate()
-        .map(|(file_idx, file_info)| async move {
-            let single_split = DataSplitBuilder::new()
+    // This read rewrites the target files, so it must see raw rows: authorize an
+    // unrestricted grant once and stamp it on every split. Stamping the scan
+    // plan's grant instead would shift the positional row offsets the writer
+    // replays under a row filter, rewriting the wrong rows.
+    let mut splits = Vec::with_capacity(file_index.len());
+    for file_info in file_index.iter() {
+        splits.push(
+            DataSplitBuilder::new()
                 .with_snapshot(file_info.snapshot_id)
                 .with_partition(
                     paimon::spec::BinaryRow::from_serialized_bytes(&file_info.partition)
@@ -1249,8 +1251,19 @@ pub(crate) async fn register_cow_target_table(
                 .with_total_buckets(file_info.total_buckets)
                 .with_data_files(vec![file_info.file_meta.clone()])
                 .build()
-                .map_err(to_datafusion_error)?;
+                .map_err(to_datafusion_error)?,
+        );
+    }
+    let splits = table
+        .authorize_rewrite_splits(splits)
+        .await
+        .map_err(to_datafusion_error)?;
 
+    // Read all files in parallel
+    let read_futures: Vec<_> = splits
+        .into_iter()
+        .enumerate()
+        .map(|(file_idx, single_split)| async move {
             let read = table
                 .new_read_builder()
                 .new_read()

crates/integrations/datafusion/src/physical_plan/scan.rs

@@ -851,6 +851,18 @@ impl PaimonTableScan {
             return Statistics::unknown_column(&self.schema());
         }
 
+        // Manifest stats describe the data BEFORE query-auth enforcement: the
+        // bounds of a masked column are the raw values the mask hides, and null
+        // counts cover rows the row filter drops. Publishing them (e.g. via
+        // EXPLAIN, or to the optimizer) would leak exactly what enforcement hides.
+        if partitions
+            .iter()
+            .flat_map(|splits| splits.iter())
+            .any(DataSplit::has_restricted_query_auth_grant)
+        {
+            return Statistics::unknown_column(&self.schema());
+        }
+
         let Ok(merge_engine) = self.table.schema().core_options().merge_engine() else {
             return Statistics::unknown_column(&self.schema());
         };
@@ -1066,6 +1078,18 @@ impl ExecutionPlan for PaimonTableScan {
             None => &self.planned_partitions,
         };
 
+        // Under a restricted grant the manifest row count is the count BEFORE
+        // the row filter runs, so publishing it (EXPLAIN, the optimizer) would
+        // disclose how much data the filter hides — the same reason the column
+        // statistics are suppressed. Report the whole thing as unknown.
+        if partitions
+            .iter()
+            .flat_map(|splits| splits.iter())
+            .any(DataSplit::has_restricted_query_auth_grant)
+        {
+            return Ok(Arc::new(Statistics::new_unknown(&self.schema())));
+        }
+
         let mut total_rows: usize = 0;
         let mut all_row_counts_known = true;
         for splits in partitions {

crates/integrations/datafusion/src/table/mod.rs

@@ -450,7 +450,11 @@ impl TableProvider for PaimonTableProvider {
             .map_err(to_datafusion_error)?;
 
         let target = state.config_options().execution.target_partitions;
+        // Inexact plan row counts (a query-auth row filter drops rows inside
+        // `TableRead`) would let DataFusion's aggregate-statistics rule answer
+        // COUNT(*) with the unfiltered count without ever invoking the read.
         let filter_exact = !filter_analysis.requires_residual
+            && plan.row_counts_exact()
             && filter_analysis
                 .pushed_predicate
                 .as_ref()

crates/integrations/datafusion/src/variant_pushdown.rs

@@ -243,6 +243,7 @@ impl ExtensionPlanner for VariantExtractionExtensionPlanner {
                 .collect()
         };
         let filter_exact = !filter_analysis.requires_residual
+            && plan.row_counts_exact()
             && filter_analysis
                 .pushed_predicate
                 .as_ref()

crates/integrations/datafusion/tests/query_auth.rs

@@ -0,0 +1,212 @@
+// 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.
+
+//! Query-auth enforcement through the DataFusion provider: SQL over a REST
+//! catalog table with `query-auth.enabled` must apply the per-user grant
+//! (row filter + column masking) fetched at scan-plan time, and COUNT(*)
+//! must not shortcut to unfiltered statistics.
+//!
+//! This is the same provider path the Python binding
+//! (`pypaimon_rust.datafusion`) drives via FFI.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use datafusion::arrow::array::{Array, Int32Array, Int64Array, StringArray};
+use datafusion::arrow::compute::cast;
+use datafusion::arrow::datatypes::{
+    DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema,
+};
+use datafusion::arrow::record_batch::RecordBatch;
+use paimon::api::{AuthTableQueryResponse, ConfigResponse};
+use paimon::catalog::{Identifier, RESTCatalog};
+use paimon::spec::{BigIntType, DataType, IntType, Schema, VarCharType};
+use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options, Table};
+use paimon_datafusion::SQLContext;
+
+// Shared REST catalog mock (same source the paimon crate's integration tests
+// use); only a subset of its helpers is exercised here.
+#[allow(dead_code)]
+#[path = "../../../paimon/tests/mock_server.rs"]
+mod mock_server;
+use mock_server::start_mock_server;
+
+async fn write_batch(table: &Table, batch: RecordBatch, commit_user: &str) {
+    let write_builder = table
+        .new_write_builder()
+        .with_commit_user(commit_user)
+        .expect("valid commit user");
+    let mut write = write_builder.new_write().expect("create writer");
+    write.write_arrow_batch(&batch).await.expect("write batch");
+    let messages = write.prepare_commit().await.expect("prepare commit");
+    write_builder
+        .new_commit()
+        .commit(messages)
+        .await
+        .expect("commit batch");
+}
+
+// Multi-threaded: the provider's `block_on_with_runtime` bridges park the
+// current thread, which must not be the only thread serving the mock server.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn test_query_auth_grant_enforced_via_sql() {
+    let tmp = tempfile::tempdir().unwrap();
+    let warehouse = format!("file://{}", tmp.path().display());
+
+    // Write demo_employees (id, name, salary) through a plain FileSystemCatalog.
+    let mut fs_options = Options::new();
+    fs_options.set(CatalogOptions::WAREHOUSE, &warehouse);
+    let fs_catalog = FileSystemCatalog::new(fs_options).expect("create filesystem catalog");
+    fs_catalog
+        .create_database("default", true, HashMap::new())
+        .await
+        .unwrap();
+    let columns = || {
+        Schema::builder()
+            .column("id", DataType::Int(IntType::new()))
+            .column("name", DataType::VarChar(VarCharType::new(255).unwrap()))
+            .column("salary", DataType::BigInt(BigIntType::new()))
+            .option("bucket", "1")
+            .option("bucket-key", "id")
+    };
+    let identifier = Identifier::new("default", "qa_emp");
+    fs_catalog
+        .create_table(&identifier, columns().build().unwrap(), false)
+        .await
+        .unwrap();
+    let writer = fs_catalog.get_table(&identifier).await.unwrap();
+    let arrow_schema = Arc::new(ArrowSchema::new(vec![
+        ArrowField::new("id", ArrowDataType::Int32, true),
+        ArrowField::new("name", ArrowDataType::Utf8, true),
+        ArrowField::new("salary", ArrowDataType::Int64, true),
+    ]));
+    let batch = RecordBatch::try_new(
+        arrow_schema,
+        vec![
+            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])),
+            Arc::new(StringArray::from(vec![
+                "alice", "bob", "charlie", "diana", "eve",
+            ])),
+            Arc::new(Int64Array::from(vec![120000, 85000, 95000, 70000, 99000])),
+        ],
+    )
+    .unwrap();
+    write_batch(&writer, batch, "u1").await;
+
+    // Serve the same files through a mock REST catalog whose grant applies a
+    // row filter (salary >= 90000) and masks name -> UPPER(name).
+    let mut defaults = HashMap::new();
+    defaults.insert("prefix".to_string(), "mock-test".to_string());
+    let server = start_mock_server(
+        "test_warehouse".to_string(),
+        "/tmp/test_warehouse".to_string(),
+        ConfigResponse::new(defaults),
+        vec!["default".to_string()],
+    )
+    .await;
+    server.add_table_with_schema(
+        "default",
+        "qa_emp",
+        columns()
+            .option("query-auth.enabled", "true")
+            .build()
+            .unwrap(),
+        &format!("{warehouse}/default.db/qa_emp"),
+    );
+    server.set_auth_response(
+        "default",
+        "qa_emp",
+        AuthTableQueryResponse {
+            filter: Some(vec![
+                r#"{"kind":"LEAF","transform":{"name":"FIELD_REF","fieldRef":{"index":2,"name":"salary","type":"BIGINT"}},"function":"GREATER_OR_EQUAL","literals":[90000]}"#
+                    .to_string(),
+            ]),
+            column_masking: Some(HashMap::from([(
+                "name".to_string(),
+                r#"{"name":"UPPER","inputs":[{"index":1,"name":"name","type":"STRING"}]}"#
+                    .to_string(),
+            )])),
+        },
+    );
+
+    let mut rest_options = Options::new();
+    rest_options.set("uri", server.url().expect("server url"));
+    rest_options.set("warehouse", "test_warehouse");
+    rest_options.set("token.provider", "bear");
+    rest_options.set("token", "test_token");
+    let rest_catalog = RESTCatalog::new(rest_options, true)
+        .await
+        .expect("create REST catalog");
+
+    // Sanity: the mock-served table must resolve through the Catalog trait.
+    rest_catalog
+        .get_table(&identifier)
+        .await
+        .expect("REST get_table(default.qa_emp)");
+
+    let mut ctx = SQLContext::new();
+    ctx.register_catalog("paimon", Arc::new(rest_catalog))
+        .await
+        .expect("register catalog");
+
+    // Row filter + masking must both be applied on the SQL result.
+    let batches = ctx
+        .sql("SELECT id, name, salary FROM paimon.default.qa_emp ORDER BY id")
+        .await
+        .expect("plan select")
+        .collect()
+        .await
+        .expect("execute select");
+    let mut rows = Vec::new();
+    for b in &batches {
+        let ids = b.column(0).as_any().downcast_ref::<Int32Array>().unwrap();
+        // DataFusion may hand back Utf8View for string columns; normalize.
+        let names = cast(b.column(1), &ArrowDataType::Utf8).expect("cast name to Utf8");
+        let names = names.as_any().downcast_ref::<StringArray>().unwrap();
+        let sal = b.column(2).as_any().downcast_ref::<Int64Array>().unwrap();
+        for r in 0..b.num_rows() {
+            rows.push((ids.value(r), names.value(r).to_string(), sal.value(r)));
+        }
+    }
+    assert_eq!(
+        rows,
+        vec![
+            (1, "ALICE".to_string(), 120000),
+            (3, "CHARLIE".to_string(), 95000),
+            (5, "EVE".to_string(), 99000),
+        ],
+        "grant must drop salary<90000 rows and uppercase name"
+    );
+
+    // COUNT(*) must reflect the filtered row count, not unfiltered statistics
+    // (the aggregate-statistics optimization must be disabled by the inexact
+    // plan row counts).
+    let batches = ctx
+        .sql("SELECT COUNT(*) FROM paimon.default.qa_emp")
+        .await
+        .expect("plan count")
+        .collect()
+        .await
+        .expect("execute count");
+    let count = batches[0]
+        .column(0)
+        .as_any()
+        .downcast_ref::<Int64Array>()
+        .unwrap()
+        .value(0);
+    assert_eq!(count, 3, "COUNT(*) must not use unfiltered statistics");
+}

crates/paimon/src/arrow/residual.rs

@@ -566,7 +566,7 @@ fn evaluate_set_membership_predicate(
     Ok(combined)
 }
 
-fn evaluate_column_predicate(
+pub(crate) fn evaluate_column_predicate(
     column: &ArrayRef,
     scalar: &Scalar<ArrayRef>,
     op: PredicateOperator,
@@ -725,7 +725,7 @@ fn combine_filter_masks(left: &BooleanArray, right: &BooleanArray, use_or: bool)
     BooleanArray::new(values, None)
 }
 
-fn boolean_mask_from_predicate(
+pub(crate) fn boolean_mask_from_predicate(
     len: usize,
     mut predicate: impl FnMut(usize) -> bool,
 ) -> BooleanArray {

crates/paimon/src/spec/schema.rs

@@ -164,6 +164,19 @@ impl TableSchema {
         new_schema
     }
 
+    /// Force `query-auth.enabled = true` on this schema.
+    ///
+    /// Used when a copy must inherit the flag from a table whose schema carries
+    /// it (the REST catalog delivers it on the table response, so an on-disk
+    /// schema — e.g. a branch's — need not have it).
+    pub(crate) fn copy_with_query_auth_enabled(&self) -> Self {
+        let mut new_schema = self.clone();
+        new_schema
+            .options
+            .insert(QUERY_AUTH_ENABLED_OPTION.to_string(), "true".to_string());
+        new_schema
+    }
+
     /// Apply a list of schema changes and return a new schema with incremented ID.
     ///
     /// Column-level changes operate on **top-level** columns only: a

crates/paimon/src/table/btree_global_index_build_builder.rs

@@ -614,7 +614,15 @@ async fn extract_index_rows(
     index_field: &DataField,
     serialize_key: SerializeKeyFn,
 ) -> Result<Vec<BTreeKeyRow>> {
-    let splits = build_read_splits_for_shard(shard)?;
+    // Building the global index reads the indexed column across the shard. Under
+    // a restricted query-auth grant that read would drop/mask rows, so the index
+    // would be built over a filtered view. Require an unrestricted grant and read
+    // raw (stamp the returned grant on each split).
+    let write_grant = table.authorize_unrestricted_read().await?;
+    let splits: Vec<crate::table::DataSplit> = build_read_splits_for_shard(shard)?
+        .into_iter()
+        .map(|s| s.with_query_auth_grant(write_grant.clone()))
+        .collect();
 
     let mut read_builder = table.new_read_builder();
     read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?;

crates/paimon/src/table/bucket_assigner_cross.rs

@@ -85,6 +85,13 @@ impl GlobalPartitionIndex {
             .collect();
         let projected_pk_indices: Vec<usize> = (0..pk_fields.len()).collect();
 
+        // Building the cross-partition PK index reads every primary key. Under a
+        // restricted query-auth grant the planned splits would filter rows, so
+        // the index would miss hidden keys and produce duplicate PKs / lost
+        // updates on upsert. Require a fully unrestricted grant (the plan then
+        // stamps that unrestricted grant on its splits, so the reads run raw).
+        table.authorize_unrestricted_read().await?;
+
         let mut rb = table.new_read_builder();
         rb.with_projection(&pk_field_names)?;
         let scan = rb.new_scan().with_scan_all_files();