feat(table): execute a primary-key vector search over engine-planned splits

bindings/c/src/tests.rs

@@ -3175,6 +3175,89 @@ fn vector_search_pk_table_read_matches_rust() {
     }
 }
 
+/// The bucket-split terminal marshals an array of buffers, which the single-split
+/// terminal does not: a caller passing a null array, a zero count, or a null
+/// entry has made an input error, and it must be reported as one rather than
+/// reaching the decoder as corrupt data.
+#[test]
+fn vector_search_bucket_splits_reject_malformed_input() {
+    let path = "memory:/vsearch_pk_split_args";
+    let (query, vectors) = pk_fixture_smoke();
+    let table = build_pk_vector_table(path, &vectors);
+    let handle = unsafe { wrap_table(table) };
+
+    unsafe {
+        // No splits at all.
+        let builder = c_vector_builder(handle, VECTOR_COLUMN, &query, 3, ptr::null_mut());
+        let result = paimon_vector_search_builder_execute_read_for_bucket_splits(
+            builder,
+            ptr::null(),
+            ptr::null(),
+            0,
+        );
+        paimon_vector_search_builder_free(builder);
+        assert!(!result.error.is_null(), "a null split array must error");
+        paimon_error_free(result.error);
+
+        // A count that does not match the (absent) arrays.
+        let builder = c_vector_builder(handle, VECTOR_COLUMN, &query, 3, ptr::null_mut());
+        let result = paimon_vector_search_builder_execute_read_for_bucket_splits(
+            builder,
+            ptr::null(),
+            ptr::null(),
+            1,
+        );
+        paimon_vector_search_builder_free(builder);
+        assert!(!result.error.is_null(), "a null split array must error");
+        paimon_error_free(result.error);
+
+        // A null entry inside an otherwise valid array.
+        let bytes: Vec<u8> = vec![1, 2, 3, 4];
+        let ptrs: [*const u8; 2] = [bytes.as_ptr(), ptr::null()];
+        let lens: [usize; 2] = [bytes.len(), 0];
+        let builder = c_vector_builder(handle, VECTOR_COLUMN, &query, 3, ptr::null_mut());
+        let result = paimon_vector_search_builder_execute_read_for_bucket_splits(
+            builder,
+            ptrs.as_ptr(),
+            lens.as_ptr(),
+            2,
+        );
+        paimon_vector_search_builder_free(builder);
+        assert!(!result.error.is_null(), "a null split entry must error");
+        paimon_error_free(result.error);
+
+        unwrap_table(handle);
+    }
+}
+
+/// Split bytes come from outside the process, so a buffer that is not a split
+/// must surface as an error, not a panic across the ABI boundary.
+#[test]
+fn vector_search_bucket_splits_reject_corrupt_bytes() {
+    let path = "memory:/vsearch_pk_split_corrupt";
+    let (query, vectors) = pk_fixture_smoke();
+    let table = build_pk_vector_table(path, &vectors);
+    let handle = unsafe { wrap_table(table) };
+
+    unsafe {
+        let garbage: Vec<u8> = vec![0xAB; 64];
+        let ptrs: [*const u8; 1] = [garbage.as_ptr()];
+        let lens: [usize; 1] = [garbage.len()];
+        let builder = c_vector_builder(handle, VECTOR_COLUMN, &query, 3, ptr::null_mut());
+        let result = paimon_vector_search_builder_execute_read_for_bucket_splits(
+            builder,
+            ptrs.as_ptr(),
+            lens.as_ptr(),
+            1,
+        );
+        paimon_vector_search_builder_free(builder);
+        assert!(!result.error.is_null(), "garbage bytes must error");
+        assert!(result.reader.is_null());
+        paimon_error_free(result.error);
+        unwrap_table(handle);
+    }
+}
+
 #[test]
 fn vector_search_append_table_read_matches_rust() {
     let path = "memory:/vsearch_append_read";

bindings/c/src/vector_search.rs

@@ -344,6 +344,117 @@ pub unsafe extern "C" fn paimon_vector_search_builder_execute_read(
     }
 }
 
+/// Run the vector search over bucket splits a Java planner produced, and stream
+/// the materialized rows.
+///
+/// This is the entry point for an engine that plans in Paimon Java and executes
+/// here: the planner emits one `BucketVectorSearchSplit` per bucket and ships its
+/// bytes to a worker, which calls this. `splits` points at `count` buffers and
+/// `split_lens` at their lengths; both arrays must hold `count` entries. The
+/// buffers are only read for the duration of the call.
+///
+/// The splits are the plan -- their payload files, per-file row ranges and pinned
+/// snapshot are used as given, and the table's index manifest is not read. Search,
+/// optional refine, Top-K and materialization are the same as
+/// `paimon_vector_search_builder_execute_read`, so the output is the projected
+/// user columns plus `__paimon_search_score`, best-first. The Top-K is local to
+/// the splits passed in; a caller distributing one call per bucket merges the
+/// per-bucket results itself.
+///
+/// Only a primary-key vector column can be read this way; a data-evolution table
+/// returns an error rather than an answer from a different plan. Consume via
+/// `paimon_record_batch_reader_next` and free with
+/// `paimon_record_batch_reader_free`.
+///
+/// # Safety
+/// `b` must be a valid pointer from `paimon_table_new_vector_search_builder`, or
+/// null (returns an error result). `splits` and `split_lens` must each point at
+/// `count` valid entries, and each `splits[i]` at `split_lens[i]` readable bytes.
+#[no_mangle]
+pub unsafe extern "C" fn paimon_vector_search_builder_execute_read_for_bucket_splits(
+    b: *mut paimon_vector_search_builder,
+    splits: *const *const u8,
+    split_lens: *const usize,
+    count: usize,
+) -> paimon_result_record_batch_reader {
+    if let Err(e) = check_non_null(b, "b") {
+        return paimon_result_record_batch_reader {
+            reader: std::ptr::null_mut(),
+            error: e,
+        };
+    }
+    if splits.is_null() || split_lens.is_null() || count == 0 {
+        return paimon_result_record_batch_reader {
+            reader: std::ptr::null_mut(),
+            error: paimon_error::new(
+                PaimonErrorCode::InvalidInput,
+                "paimon_vector_search_builder_execute_read_for_bucket_splits: null or empty splits"
+                    .to_string(),
+            ),
+        };
+    }
+
+    let ptrs = std::slice::from_raw_parts(splits, count);
+    let lens = std::slice::from_raw_parts(split_lens, count);
+    let mut buffers: Vec<&[u8]> = Vec::with_capacity(count);
+    for (i, (&ptr, &len)) in ptrs.iter().zip(lens).enumerate() {
+        // A null or empty buffer cannot be a split, and reaching the decoder with
+        // one would report it as corrupt data rather than as the caller's error.
+        if ptr.is_null() || len == 0 {
+            return paimon_result_record_batch_reader {
+                reader: std::ptr::null_mut(),
+                error: paimon_error::new(
+                    PaimonErrorCode::InvalidInput,
+                    format!(
+                        "paimon_vector_search_builder_execute_read_for_bucket_splits: \
+                         split {i} is null or empty"
+                    ),
+                ),
+            };
+        }
+        buffers.push(std::slice::from_raw_parts(ptr, len));
+    }
+
+    let state = &*((*b).inner as *const VectorSearchState);
+    let mut builder = state.table.new_vector_search_builder();
+    if let Some(col) = &state.vector_column {
+        builder.with_vector_column(col);
+    }
+    if let Some(v) = &state.query_vector {
+        builder.with_query_vector(v.clone());
+    }
+    if let Some(limit) = state.limit {
+        builder.with_limit(limit);
+    }
+    if !state.options.is_empty() {
+        builder.with_options(state.options.clone());
+    }
+    if let Some(f) = &state.filter {
+        builder.with_filter(f.clone());
+    }
+    if let Some(cols) = &state.projection {
+        let col_refs: Vec<&str> = cols.iter().map(String::as_str).collect();
+        builder.with_projection(&col_refs);
+    }
+
+    match runtime().block_on(builder.execute_read_for_bucket_splits(&buffers)) {
+        Ok(stream) => {
+            let reader = Box::new(stream);
+            let wrapper = Box::new(paimon_record_batch_reader {
+                inner: Box::into_raw(reader) as *mut c_void,
+            });
+            paimon_result_record_batch_reader {
+                reader: Box::into_raw(wrapper),
+                error: std::ptr::null_mut(),
+            }
+        }
+        Err(e) => paimon_result_record_batch_reader {
+            reader: std::ptr::null_mut(),
+            error: paimon_error::from_paimon(e),
+        },
+    }
+}
+
 // --- C ABI signature guards -------------------------------------------------
 //
 // These symbols are called across the FFI boundary with fixed argument counts:
@@ -387,3 +498,10 @@ const _: unsafe extern "C" fn(*mut paimon_vector_search_builder) =
 const _: unsafe extern "C" fn(
     *mut paimon_vector_search_builder,
 ) -> paimon_result_record_batch_reader = paimon_vector_search_builder_execute_read;
+const _: unsafe extern "C" fn(
+    *mut paimon_vector_search_builder,
+    *const *const u8,
+    *const usize,
+    usize,
+) -> paimon_result_record_batch_reader =
+    paimon_vector_search_builder_execute_read_for_bucket_splits;

crates/paimon/src/lumina/reader.rs

@@ -122,6 +122,27 @@ fn compare_scores(a: f32, b: f32) -> std::cmp::Ordering {
     }
 }
 
+/// Turn an include-set into the dense id array Lumina's filtered search takes.
+///
+/// Mirrors Java `LuminaVectorGlobalIndexReader.toScopedIds`: an include-set above
+/// `Integer.MAX_VALUE` is refused BEFORE the array is allocated, because that array
+/// is 8 bytes per id and the set can hold one id per live row of the segment.
+fn to_scoped_ids(include_row_ids: &roaring::RoaringTreemap) -> crate::Result<Vec<u64>> {
+    let cardinality = include_row_ids.len();
+    if cardinality > i32::MAX as u64 {
+        return Err(crate::Error::DataInvalid {
+            message: format!(
+                "include_row_ids cardinality ({cardinality}) exceeds {}",
+                i32::MAX
+            ),
+            source: None,
+        });
+    }
+    let mut ids = Vec::with_capacity(cardinality as usize);
+    ids.extend(include_row_ids.iter());
+    Ok(ids)
+}
+
 /// Allocate the label buffer for one native search, filled with [`SENTINEL`].
 ///
 /// [`SENTINEL`] is the "no result" marker (the C ABI's `-1`), and
@@ -373,7 +394,7 @@ fn search_lumina<S: LuminaSearch + ?Sized>(
     let include_row_ids = vector_search.effective_include_row_ids();
 
     let (distances, labels) = if let Some(include_ids) = include_row_ids {
-        let filter_id_list: Vec<u64> = include_ids.iter().collect();
+        let filter_id_list: Vec<u64> = to_scoped_ids(include_ids)?;
         if filter_id_list.is_empty() {
             return Ok(None);
         }
@@ -460,14 +481,26 @@ fn search_lumina_batch<S: LuminaSearch + ?Sized>(
         }
     }
 
-    let filter_id_list =
-        shared_filter.map(|include_row_ids| include_row_ids.iter().collect::<Vec<_>>());
-    if filter_id_list.as_ref().is_some_and(Vec::is_empty) {
+    // An include-set that permits nothing can match nothing, whatever the index
+    // holds. Answered from the set itself, so no native call is made at all.
+    if shared_filter.is_some_and(|filter| filter.is_empty()) {
         return Ok(vec![None; vector_searches.len()]);
     }
 
-    let index_metric = index_meta.metric()?;
+    // Java's order, which decides what an empty or unreadable index reports:
+    // `index.size()` first, then the dense filter, and the metric only on the way
+    // out. An empty index has nothing to return whatever the filter says, so
+    // converting first would report an oversized filter instead, and reading the
+    // metric first would report a malformed one.
     let count = searcher.get_count()? as usize;
+    if count == 0 {
+        return Ok(vec![None; vector_searches.len()]);
+    }
+    let filter_id_list = shared_filter
+        .map(|include_row_ids| to_scoped_ids(include_row_ids))
+        .transpose()?;
+    let index_metric = index_meta.metric()?;
+
     let effective_k = filter_id_list.as_ref().map_or_else(
         || std::cmp::min(limit, count),
         |ids| std::cmp::min(std::cmp::min(limit, count), ids.len()),
@@ -589,6 +622,25 @@ impl Drop for LuminaVectorGlobalIndexReader {
 
 #[cfg(test)]
 mod tests {
+
+    #[test]
+    fn to_scoped_ids_refuses_a_set_larger_than_the_dense_array_can_index() {
+        // Java's own limit for this exact quantity at this exact seam
+        // (`LuminaVectorGlobalIndexReader.toScopedIds`). The array is 8 bytes per id,
+        // and the set can hold one id per live row of the segment, so the refusal has
+        // to come BEFORE the allocation.
+        let mut small = roaring::RoaringTreemap::new();
+        small.insert_range(0..=9u64);
+        assert_eq!(to_scoped_ids(&small).unwrap().len(), 10);
+
+        let mut over = roaring::RoaringTreemap::new();
+        over.insert_range(0..=(i32::MAX as u64 + 1));
+        let error = to_scoped_ids(&over)
+            .map(|_| ())
+            .expect_err("a set this size cannot be handed to the backend");
+        assert!(error.to_string().contains("exceeds"), "{error}");
+    }
+
     use super::*;
     use crate::lumina::{KEY_DIMENSION, KEY_DISTANCE_METRIC};
     use crate::vector_search::GlobalIndexIOMeta;
@@ -797,6 +849,46 @@ mod tests {
         assert!(calls.iter().all(|call| call.n == 1));
     }
 
+    #[test]
+    fn an_empty_index_returns_empty_before_the_filter_is_sized() {
+        // Java's order: `index.size()` decides first, so an index holding nothing
+        // returns nothing whatever the filter says. Sizing the dense filter first
+        // would turn this into an oversized-filter error instead.
+        let searcher = RecordingSearcher::new(0);
+        let mut oversized = roaring::RoaringTreemap::new();
+        oversized.insert_range(0..=(i32::MAX as u64 + 1));
+        let shared_filter = Arc::new(oversized);
+        let mut first = VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap();
+        first.set_shared_include_row_ids(Arc::clone(&shared_filter));
+        let mut second = VectorSearch::new(vec![0.0, 1.0], 2, "embedding".to_string()).unwrap();
+        second.set_shared_include_row_ids(Arc::clone(&shared_filter));
+
+        let results = search_lumina_batch(
+            &searcher,
+            &test_index_meta(2),
+            &HashMap::new(),
+            &[first, second],
+        )
+        .expect("an empty index reports no hits, not an oversized filter");
+
+        assert_eq!(results, vec![None, None]);
+        assert_eq!(
+            searcher.count_calls.load(Ordering::Relaxed),
+            1,
+            "the index size is what the answer came from"
+        );
+        assert!(searcher
+            .unfiltered_calls
+            .lock()
+            .expect("unfiltered call lock")
+            .is_empty());
+        assert!(searcher
+            .filtered_calls
+            .lock()
+            .expect("filtered call lock")
+            .is_empty());
+    }
+
     #[test]
     fn test_empty_shared_filter_skips_native_search() {
         let searcher = RecordingSearcher::new(10);