perf(vindex): parallelize Parquet reads for index builds

crates/paimon/examples/ivfpq_build_benchmark.rs

@@ -0,0 +1,93 @@
+// 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.
+
+//! Build an IVF-PQ index through the production Paimon path.
+//!
+//! ```text
+//! PAIMON_CATALOG_OPTIONS='{"metastore":"filesystem","warehouse":"/tmp/warehouse"}' \
+//! PAIMON_LOG_VECTOR_INDEX_BUILD_TIMING=1 \
+//! cargo run --release -p paimon --example ivfpq_build_benchmark -- \
+//!   <database> <table> <vector-column> [--drop-existing]
+//! ```
+
+use std::collections::HashMap;
+use std::error::Error;
+use std::time::Instant;
+
+use paimon::catalog::Identifier;
+use paimon::{CatalogFactory, Options};
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn Error>> {
+    let mut args = std::env::args().skip(1);
+    let database = required_arg(&mut args, "database")?;
+    let table_name = required_arg(&mut args, "table")?;
+    let column = required_arg(&mut args, "vector-column")?;
+    let drop_existing = args.any(|arg| arg == "--drop-existing");
+
+    let catalog_options = std::env::var("PAIMON_CATALOG_OPTIONS")?;
+    let catalog =
+        CatalogFactory::create(Options::from_map(serde_json::from_str(&catalog_options)?)).await?;
+    let table = catalog
+        .get_table(&Identifier::new(&database, &table_name))
+        .await?;
+
+    let dropped_index_files = if drop_existing {
+        let mut builder = table.new_global_index_drop_builder();
+        builder.with_index_column(&column).with_index_type("ivf-pq");
+        builder.execute().await?
+    } else {
+        0
+    };
+
+    let options = HashMap::from([
+        ("dimension".to_string(), "768".to_string()),
+        ("metric".to_string(), "cosine".to_string()),
+        ("nlist".to_string(), "4096".to_string()),
+        ("pq.m".to_string(), "192".to_string()),
+    ]);
+    let started = Instant::now();
+    let built_shards = table
+        .new_vindex_index_build_builder("ivf-pq")
+        .with_index_column(&column)
+        .with_options(options.clone())
+        .execute()
+        .await?;
+
+    println!(
+        "{}",
+        serde_json::to_string_pretty(&serde_json::json!({
+            "database": database,
+            "table": table_name,
+            "column": column,
+            "index_type": "ivf-pq",
+            "build_options": options,
+            "dropped_index_files": dropped_index_files,
+            "built_shards": built_shards,
+            "duration_seconds": started.elapsed().as_secs_f64(),
+        }))?
+    );
+    Ok(())
+}
+
+fn required_arg(
+    args: &mut impl Iterator<Item = String>,
+    name: &str,
+) -> Result<String, Box<dyn Error>> {
+    args.next()
+        .ok_or_else(|| format!("missing <{name}> argument").into())
+}

crates/paimon/src/arrow/format/parquet.rs

@@ -465,8 +465,8 @@ impl FormatFileReader for ParquetFormatReader {
             combined_selection =
                 intersect_optional_row_selections(combined_selection, Some(range_selection));
         }
-        if let Some(sel) = combined_selection {
-            batch_stream_builder = batch_stream_builder.with_row_selection(sel);
+        if let Some(ref selection) = combined_selection {
+            batch_stream_builder = batch_stream_builder.with_row_selection(selection.clone());
         }
         if let Some(size) = batch_size {
             batch_stream_builder = batch_stream_builder.with_batch_size(size);
@@ -490,29 +490,46 @@ impl FormatFileReader for ParquetFormatReader {
         // preserving positional `_ROW_ID`, sort order, and batch backpressure. Reads
         // with predicates or an explicit row selection retain the original
         // single-stream path until their selections are split per row group.
-        let row_group_parallelism = self
-            .read_budget
-            .as_ref()
-            .filter(|_| preds.is_empty() && row_filter_factory.is_none() && row_selection.is_none())
+        let read_budget = self.read_budget.as_ref().filter(|_| {
+            preds.is_empty() && row_filter_factory.is_none() && row_selection.is_none()
+        });
+        let row_group_parallelism = read_budget
             .map(|budget| {
                 budget
                     .parallelism()
                     .min(batch_stream_builder.metadata().num_row_groups())
             })
             .unwrap_or(1);
+        let projected_bytes = self
+            .read_budget
+            .as_ref()
+            .filter(|budget| row_group_parallelism > 1 || budget.diagnostics_enabled())
+            .map(|budget| {
+                let mut diagnostic_selection = combined_selection;
+                let projected_bytes = batch_stream_builder
+                    .metadata()
+                    .row_groups()
+                    .iter()
+                    .filter(|row_group| {
+                        diagnostic_selection.as_mut().is_none_or(|selection| {
+                            selection
+                                .split_off(row_group.num_rows() as usize)
+                                .selects_any()
+                        })
+                    })
+                    .map(|row_group| projected_row_group_bytes(row_group, &mask))
+                    .collect::<Vec<_>>();
+                budget.record_projected_row_groups(&projected_bytes);
+                projected_bytes
+            });
         if row_group_parallelism > 1 {
             let row_group_count = batch_stream_builder.metadata().num_row_groups();
             let reader_metadata = ArrowReaderMetadata::try_new(
                 batch_stream_builder.metadata().clone(),
                 ArrowReaderOptions::new(),
             )?;
-            let projected_bytes = batch_stream_builder
-                .metadata()
-                .row_groups()
-                .iter()
-                .map(|row_group| projected_row_group_bytes(row_group, &mask))
-                .collect::<Vec<_>>();
-            let read_budget = Arc::clone(self.read_budget.as_ref().expect("checked above"));
+            let projected_bytes = projected_bytes.expect("parallel row-group reads need sizes");
+            let read_budget = Arc::clone(read_budget.expect("checked above"));
             let (row_group_tx, mut row_group_rx) = mpsc::channel(row_group_parallelism);
             tokio::spawn(async move {
                 for (row_group_index, projected_bytes) in projected_bytes.into_iter().enumerate() {
@@ -2885,6 +2902,35 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_parquet_diagnostics_include_reads_with_row_selection() {
+        let data = write_multi_row_group_parquet(32, 64, EnabledStatistics::Chunk).await;
+        let budget = Arc::new(ParquetReadBudget::new(8, 256 * 1024 * 1024).unwrap());
+        budget.enable_diagnostics();
+        let file_size = data.len() as u64;
+        let fields = vec![int_field("id")];
+        let batches = ParquetFormatReader::with_read_budget(Arc::clone(&budget))
+            .read_batch_stream(
+                Box::new(TrackingFileRead::new(Bytes::from(data))),
+                file_size,
+                &fields,
+                None,
+                Some(32),
+                Some(vec![RowRange::new(0, 9)]),
+            )
+            .await
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+
+        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 10);
+        let diagnostics = budget.diagnostics();
+        assert_eq!(diagnostics.row_group_count, 1);
+        assert!(diagnostics.projected_bytes_total > 0);
+        assert_eq!(diagnostics.peak_inflight, 0);
+    }
+
     #[tokio::test]
     async fn test_row_group_batch_forwarding_applies_backpressure() {
         let schema = Arc::new(ArrowSchema::empty());