perf(vindex): size native batches by active indexes

crates/paimon/src/table/vector_search_builder.rs

@@ -106,7 +106,7 @@ async fn execute_vindex_searches<S: SeekRead + 'static, G: Send + 'static>(
     vector_searches: Vec<VectorSearch>,
     source: S,
     file_name: String,
-    shard_concurrency: usize,
+    index_parallelism: usize,
     guard: G,
 ) -> crate::Result<Vec<Option<HashMap<u64, f32>>>> {
     let panic_context = if vector_searches.len() > 1 {
@@ -116,7 +116,7 @@ async fn execute_vindex_searches<S: SeekRead + 'static, G: Send + 'static>(
     };
     execute_global_index_with_guard(panic_context, guard, move || {
         let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options)
-            .with_batch_shard_concurrency(shard_concurrency);
+            .with_batch_index_parallelism(index_parallelism);
         reader
             .visit_batch_vector_search(&vector_searches, |_| Ok(source))
             .map_err(|e| crate::Error::DataInvalid {
@@ -134,6 +134,10 @@ fn current_tokio_runtime_handle() -> crate::Result<tokio::runtime::Handle> {
     })
 }
 
+fn vindex_index_parallelism(entry_count: usize, max_concurrency: usize) -> usize {
+    entry_count.min(max_concurrency).max(1)
+}
+
 pub struct VectorSearchBuilder<'a> {
     table: &'a Table,
     vector_column: Option<String>,
@@ -821,6 +825,16 @@ async fn plan_and_search_pk_candidates_batch(
             source: None,
         }
     })?;
+    let batch_index_parallelism = match backend {
+        VectorIndexBackend::Vindex => vindex_index_parallelism(
+            plan.splits
+                .iter()
+                .map(|split| split.ann_segments.len())
+                .sum(),
+            concurrency,
+        ),
+        VectorIndexBackend::Lumina => 1,
+    };
 
     // Production data-file reader, mirroring `table_read.rs::new_data_file_reader`
     // but projecting only the vector column with no predicates.
@@ -906,7 +920,7 @@ async fn plan_and_search_pk_candidates_batch(
                 }
                 (VectorIndexBackend::Vindex, AnnSegmentSource::Vindex(source)) => {
                     let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options.clone())
-                        .with_batch_shard_concurrency(concurrency);
+                        .with_batch_index_parallelism(batch_index_parallelism);
                     reader.load_validated(
                         |_| Ok(source),
                         |metadata| {
@@ -1523,6 +1537,13 @@ async fn evaluate_batch_vector_search(
         }
         ensure_global_index_executor_capacity(concurrency);
         let range_read_permits = Arc::new(tokio::sync::Semaphore::new(concurrency));
+        let batch_index_parallelism = vindex_index_parallelism(
+            vector_entries
+                .iter()
+                .filter(|entry| is_vindex_index_type(&entry.index_file.index_type))
+                .count(),
+            concurrency,
+        );
         let futures: Vec<_> = vector_entries
             .into_iter()
             .map(|entry| {
@@ -1607,7 +1628,7 @@ async fn evaluate_batch_vector_search(
                                         vector_searches,
                                         source,
                                         file_name,
-                                        concurrency,
+                                        batch_index_parallelism,
                                         permit,
                                     )
                                     .await?
@@ -1628,7 +1649,7 @@ async fn evaluate_batch_vector_search(
                                         vector_searches,
                                         Cursor::new(data),
                                         file_name,
-                                        concurrency,
+                                        batch_index_parallelism,
                                         permit,
                                     )
                                     .await?
@@ -3084,6 +3105,14 @@ mod tests {
         VectorSearchMetric::L2.distance_to_score(distance)
     }
 
+    #[test]
+    fn vindex_batch_parallelism_tracks_active_entries() {
+        assert_eq!(vindex_index_parallelism(1, 1), 1);
+        assert_eq!(vindex_index_parallelism(1, 64), 1);
+        assert_eq!(vindex_index_parallelism(8, 4), 4);
+        assert_eq!(vindex_index_parallelism(4, 8), 4);
+    }
+
     fn make_field(id: i32, name: &str) -> DataField {
         DataField::new(id, name.to_string(), DataType::Int(IntType::default()))
     }

crates/paimon/src/vindex/range_reader.rs

@@ -633,6 +633,42 @@ mod tests {
         assert_eq!(tracking.max_active.load(Ordering::SeqCst), 1);
     }
 
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn read_many_caps_each_batch_at_32_ranges() {
+        let range_count = RANGE_READ_CONCURRENCY + 1;
+        let stride = RANGE_COALESCE_GAP + 2;
+        let data = Bytes::from(vec![8u8; range_count * stride as usize]);
+        let tracking = Arc::new(ConcurrencyTrackingRead {
+            data: data.clone(),
+            active: AtomicUsize::new(0),
+            max_active: AtomicUsize::new(0),
+        });
+        let source: Arc<dyn FileRead> = tracking.clone();
+        let mut reader = VindexFileReader::new(
+            source,
+            tokio::runtime::Handle::current(),
+            data.len() as u64,
+            "index".to_string(),
+        );
+
+        tokio::task::spawn_blocking(move || {
+            let mut buffers = vec![[0u8; 1]; range_count];
+            let mut requests = buffers
+                .iter_mut()
+                .enumerate()
+                .map(|(index, buffer)| ReadRequest::new(index as u64 * stride, buffer))
+                .collect::<Vec<_>>();
+            reader.pread(&mut requests).unwrap();
+        })
+        .await
+        .unwrap();
+
+        assert_eq!(
+            tracking.max_active.load(Ordering::SeqCst),
+            RANGE_READ_CONCURRENCY
+        );
+    }
+
     #[test]
     fn local_fs_read_completes_with_one_host_blocking_thread() {
         let temp_dir = tempfile::tempdir().unwrap();