fix(auth): stop BLOB view predicate lookups at LIMIT under query-auth

crates/paimon/src/table/table_read.rs

@@ -42,6 +42,7 @@ use arrow_select::concat::concat as arrow_concat;
 use arrow_select::take::take;
 use futures::{stream, StreamExt};
 use std::cmp::Ordering;
+use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
 use std::sync::Arc;
 
 #[path = "audit_log_table/read.rs"]
@@ -1302,6 +1303,15 @@ impl<'a> PaimonTableRead<'a> {
         if resolve_predicate_views {
             output_views.retain(|field| !predicate_views.contains(field));
         }
+        // Never look up views past the rows LIMIT still needs.
+        let quota = self
+            .limit
+            .filter(|_| resolve_predicate_views)
+            .map(|limit| Arc::new(AtomicUsize::new(limit)));
+        let stream = match &quota {
+            Some(quota) => take_within_quota(stream, Arc::clone(quota)),
+            None => stream,
+        };
         // User predicates see resolved descriptors or payloads, after authorization.
         let stream = match self.table.rest_env().filter(|_| resolve_predicate_views) {
             Some(env) => super::data_evolution_reader::resolve_blob_view_stream(
@@ -1359,10 +1369,11 @@ impl<'a> PaimonTableRead<'a> {
                     )
                 });
                 let stream = project_authorized_stream(Box::pin(stream), self.read_type.len());
-                match self.limit {
+                let stream = match self.limit {
                     Some(limit) => limit_rows(stream, limit),
                     None => stream,
-                }
+                };
+                spend_quota(stream, quota.clone())
             }
         };
         let stream = match self
@@ -1380,13 +1391,16 @@ impl<'a> PaimonTableRead<'a> {
             None => stream,
         };
         Ok(match blob_plan {
-            Some(plan) => plan.finish(
-                stream,
-                core_options,
-                self.table.file_io.clone(),
-                self.blob_parallelism,
-                self.limit,
-            ),
+            Some(plan) => {
+                let stream = plan.finish(
+                    stream,
+                    core_options,
+                    self.table.file_io.clone(),
+                    self.blob_parallelism,
+                    self.limit,
+                );
+                spend_quota(stream, quota)
+            }
             None => resolve_primary_key_blob_stream(
                 stream,
                 &self.read_type,
@@ -1641,6 +1655,47 @@ fn project_authorized_stream(
     }))
 }
 
+/// Cuts batches to at most `quota` rows; stages up to `spend_quota` must not
+/// read ahead.
+fn take_within_quota(
+    stream: ArrowRecordBatchStream,
+    quota: Arc<AtomicUsize>,
+) -> ArrowRecordBatchStream {
+    Box::pin(async_stream::try_stream! {
+        let mut stream = stream;
+        while quota.load(AtomicOrdering::Relaxed) > 0 {
+            let Some(batch) = stream.next().await else { break };
+            let batch = batch?;
+            let mut offset = 0;
+            while offset < batch.num_rows() {
+                let width = quota.load(AtomicOrdering::Relaxed).min(batch.num_rows() - offset);
+                if width == 0 {
+                    break;
+                }
+                yield batch.slice(offset, width);
+                offset += width;
+            }
+        }
+    })
+}
+
+/// Subtracts the rows that passed the predicates from `quota`.
+fn spend_quota(
+    stream: ArrowRecordBatchStream,
+    quota: Option<Arc<AtomicUsize>>,
+) -> ArrowRecordBatchStream {
+    let Some(quota) = quota else {
+        return stream;
+    };
+    Box::pin(stream.inspect(move |batch| {
+        if let Ok(batch) = batch {
+            let left = quota.load(AtomicOrdering::Relaxed);
+            let left = left.saturating_sub(batch.num_rows());
+            quota.store(left, AtomicOrdering::Relaxed);
+        }
+    }))
+}
+
 /// Stops once `limit` rows are out, without polling for more.
 fn limit_rows(stream: ArrowRecordBatchStream, limit: usize) -> ArrowRecordBatchStream {
     Box::pin(stream::unfold(

crates/paimon/tests/mock_server.rs

@@ -54,6 +54,7 @@ type PartitionSpecPageResponse = (Vec<HashMap<String, String>>, Option<String>);
 struct MockState {
     databases: HashMap<String, GetDatabaseResponse>,
     tables: HashMap<String, GetTableResponse>,
+    get_table_call_counts: HashMap<String, usize>,
     views: HashMap<String, GetViewResponse>,
     functions: HashMap<String, GetFunctionResponse>,
     tags: HashMap<String, GetTagResponse>,
@@ -825,6 +826,7 @@ impl RESTServer {
         let mut s = state.inner.lock().unwrap();
 
         let key = format!("{db}.{table}");
+        *s.get_table_call_counts.entry(key.clone()).or_default() += 1;
         if let Some(expected) = s.table_read_via.get(&key) {
             if headers
                 .get("X-Paimon-Read-Via")
@@ -2159,6 +2161,17 @@ impl RESTServer {
             .unwrap_or_default()
     }
 
+    /// Return how many get-table requests the table received.
+    pub fn table_get_call_count(&self, database: &str, table: &str) -> usize {
+        self.inner
+            .lock()
+            .unwrap()
+            .get_table_call_counts
+            .get(&format!("{database}.{table}"))
+            .copied()
+            .unwrap_or_default()
+    }
+
     /// Return how many paged partition-list requests the table received.
     pub fn table_partition_list_call_count(&self, database: &str, table: &str) -> usize {
         self.inner

crates/paimon/tests/rest_catalog_test.rs

@@ -1259,6 +1259,27 @@ async fn test_blob_view_limit_only_resolves_selected_references() {
         );
     }
 
+    // LIMIT stops before the invalid fourth reference.
+    server.set_auth_response(
+        "default",
+        "blob_view_target",
+        rules(&[int_leaf(0, "id", "LESS_THAN", 5)], &[]),
+    );
+    let mut limited = authorized.new_read_builder();
+    limited.with_limit(2);
+    limited.with_filter(
+        PredicateBuilder::new(authorized.schema().fields())
+            .is_not_null("picture")
+            .unwrap(),
+    );
+    assert_eq!(
+        collect_blob_rows(&read_all(&limited).await.unwrap()),
+        vec![
+            (1, "Updated".to_string(), Some(first_payload.clone())),
+            (3, "Repeated again".to_string(), Some(b"bob".to_vec())),
+        ]
+    );
+
     // A view reference and its upstream BLOB can disagree on nullness.
     // Until rules can resolve that dependency safely, reject instead of
     // authorizing rows by the reference bytes.
@@ -3578,6 +3599,211 @@ async fn test_blob_view_mask_alias_resolves_the_source_view() {
     }
 }
 
+#[cfg(not(windows))]
+#[tokio::test]
+async fn test_query_auth_blob_view_predicate_limit_after_pk_merge() {
+    let source = written(
+        "view_limit_source",
+        blob_schema,
+        vec![blob_batch(
+            vec![1, 2],
+            vec!["Skipped", "Kept"],
+            vec![b"skip".to_vec(), b"keep".to_vec()],
+        )],
+    )
+    .await;
+    let picture_id = source.table.schema().fields()[2].id();
+    let reference = |row_id| {
+        BlobViewStruct::new(source.identifier.clone(), picture_id, row_id)
+            .serialize()
+            .unwrap()
+    };
+    let (skipped, kept, missing) = (reference(0), reference(1), reference(99));
+    let batch = |pictures: Vec<Option<&[u8]>>| {
+        RecordBatch::try_from_iter([
+            (
+                "id",
+                Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as arrow_array::ArrayRef,
+            ),
+            (
+                "picture",
+                Arc::new(LargeBinaryArray::from(pictures)) as arrow_array::ArrayRef,
+            ),
+            (
+                "alias",
+                Arc::new(LargeBinaryArray::from(vec![None::<&[u8]>; 4])) as arrow_array::ArrayRef,
+            ),
+        ])
+        .unwrap()
+    };
+    for engine in ["deduplicate", "partial-update"] {
+        let g = written(
+            "view_limit_target",
+            |options| {
+                let mut schema = Schema::builder()
+                    .column("id", DataType::Int(IntType::with_nullable(false)))
+                    .column("picture", DataType::Blob(BlobType::new()))
+                    .column("alias", DataType::Blob(BlobType::new()))
+                    .primary_key(["id"])
+                    .option("bucket", "1")
+                    .option("merge-engine", engine)
+                    .option("blob-view-field", "picture");
+                for (key, value) in options {
+                    schema = schema.option(*key, *value);
+                }
+                schema.build().unwrap()
+            },
+            // Overlapping commits force a merge.
+            vec![
+                batch(vec![Some(&missing); 4]),
+                batch(vec![
+                    Some(&missing),
+                    Some(&skipped),
+                    Some(&kept),
+                    Some(&missing),
+                ]),
+            ],
+        )
+        .await;
+        g.ctx.server.add_table_with_schema(
+            "default",
+            "view_limit_source",
+            blob_schema(&[]),
+            source.table.location(),
+        );
+        for target in ["picture", "alias"] {
+            let masks = if target == "alias" {
+                vec![("alias", field_ref_mask(1, "picture", "BLOB"))]
+            } else {
+                vec![]
+            };
+            g.set_rules(&[int_leaf(0, "id", "GREATER_THAN", 1)], &masks);
+            let plan = g.table.new_read_builder().new_scan().plan().await.unwrap();
+            assert!(plan.splits().iter().any(|split| !split.raw_convertible()));
+            for batch_size in [1, 1024] {
+                for descriptors in [false, true] {
+                    let table = g.table.copy_with_options(HashMap::from([
+                        ("read.batch-size".to_string(), batch_size.to_string()),
+                        ("blob-as-descriptor".to_string(), descriptors.to_string()),
+                    ]));
+                    let predicates = PredicateBuilder::new(table.schema().fields());
+                    let label = format!(
+                        "{engine}, {target}, batch={batch_size}, descriptors={descriptors}"
+                    );
+                    for projection in [vec!["id"], vec!["id", target]] {
+                        for limit in [Some(0), Some(1), Some(2), Some(3), None] {
+                            let mut read = table.new_read_builder();
+                            read.with_projection(&projection)
+                                .unwrap()
+                                .with_filter(predicates.is_not_null(target).unwrap());
+                            if let Some(limit) = limit {
+                                read.with_limit(limit);
+                            }
+                            let result = read_all(&read).await;
+                            if let Some(limit @ 0..=2) = limit {
+                                let batches = result
+                                    .unwrap_or_else(|e| panic!("{label}, limit={limit}: {e}"));
+                                assert_eq!(ids(&batches), [2, 3][..limit], "{label}");
+                            } else {
+                                let error = result.unwrap_err().to_string();
+                                assert!(error.contains("row_id=99"), "{label}: {error}");
+                            }
+                        }
+                    }
+                    if !descriptors {
+                        // A rejected candidate consumes no quota.
+                        let mut read = table.new_read_builder();
+                        read.with_projection(&["id"])
+                            .unwrap()
+                            .with_filter(
+                                predicates
+                                    .equal(target, Datum::Bytes(b"keep".to_vec()))
+                                    .unwrap(),
+                            )
+                            .with_limit(1);
+                        let batches = read_all(&read)
+                            .await
+                            .unwrap_or_else(|e| panic!("{label}: {e}"));
+                        assert_eq!(ids(&batches), vec![3], "{label}");
+                    }
+                }
+            }
+        }
+    }
+}
+
+#[cfg(not(windows))]
+#[tokio::test]
+async fn test_query_auth_blob_view_predicate_limit_groups_lookups() {
+    let names = vec!["a", "b", "c", "d"];
+    let source = written(
+        "lookup_source",
+        blob_schema,
+        vec![blob_batch(
+            vec![1, 2, 3, 4],
+            names.clone(),
+            names.iter().map(|name| name.as_bytes().to_vec()).collect(),
+        )],
+    )
+    .await;
+    let picture_id = source.table.schema().fields()[2].id();
+    let references: Vec<Vec<u8>> = (0..4)
+        .map(|row_id| {
+            BlobViewStruct::new(source.identifier.clone(), picture_id, row_id)
+                .serialize()
+                .unwrap()
+        })
+        .collect();
+    let g = written(
+        "lookup_target",
+        |options| {
+            let mut options = options.to_vec();
+            options.push(("blob-view-field", "picture"));
+            blob_schema(&options)
+        },
+        vec![blob_batch(vec![1, 2, 3, 4], names, references)],
+    )
+    .await;
+    let server = &g.ctx.server;
+    server.add_table_with_schema(
+        "default",
+        "lookup_source",
+        blob_schema(&[]),
+        source.table.location(),
+    );
+    g.set_rules(&[int_leaf(0, "id", "GREATER_THAN", 0)], &[]);
+    let predicates = PredicateBuilder::new(g.table.schema().fields());
+    // Distinct references load in one request, with and without a payload filter.
+    for descriptors in [false, true] {
+        let table = g.table.copy_with_options(HashMap::from([(
+            "blob-as-descriptor".to_string(),
+            descriptors.to_string(),
+        )]));
+        for (limit, expected) in [
+            (Some(2), vec![1, 2]),
+            (Some(4), vec![1, 2, 3, 4]),
+            (None, vec![1, 2, 3, 4]),
+        ] {
+            let label = format!("descriptors={descriptors}, limit={limit:?}");
+            let mut read = table.new_read_builder();
+            read.with_projection(&["id"])
+                .unwrap()
+                .with_filter(predicates.is_not_null("picture").unwrap());
+            if let Some(limit) = limit {
+                read.with_limit(limit);
+            }
+            let before = server.table_get_call_count("default", "lookup_source");
+            let batches = read_all(&read).await.unwrap();
+            assert_eq!(ids(&batches), expected, "{label}");
+            assert_eq!(
+                server.table_get_call_count("default", "lookup_source") - before,
+                1,
+                "{label}"
+            );
+        }
+    }
+}
+
 fn schema_of(columns: &[&str], options: &[(&str, &str)]) -> Schema {
     let mut builder = Schema::builder();
     for name in columns {