perf(vindex): parallelize Parquet reads and multipart index uploads

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

@@ -490,29 +490,36 @@ 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 = read_budget
+            .filter(|budget| row_group_parallelism > 1 || budget.diagnostics_enabled())
+            .map(|budget| {
+                let projected_bytes = batch_stream_builder
+                    .metadata()
+                    .row_groups()
+                    .iter()
+                    .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() {

crates/paimon/src/arrow/parquet_read_budget.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
 use std::sync::Arc;
 
 use tokio::sync::{OwnedSemaphorePermit, Semaphore};
@@ -30,6 +31,43 @@ pub struct ParquetReadBudget {
     row_groups: Arc<Semaphore>,
     bytes: Arc<Semaphore>,
     byte_permits: u32,
+    oversized_warning_logged: AtomicBool,
+    diagnostics: Arc<ParquetReadDiagnostics>,
+}
+
+#[derive(Debug)]
+struct ParquetReadDiagnostics {
+    enabled: AtomicBool,
+    row_group_count: AtomicU64,
+    projected_bytes_min: AtomicU64,
+    projected_bytes_max: AtomicU64,
+    projected_bytes_total: AtomicU64,
+    current_inflight: AtomicUsize,
+    peak_inflight: AtomicUsize,
+}
+
+impl Default for ParquetReadDiagnostics {
+    fn default() -> Self {
+        Self {
+            enabled: AtomicBool::new(false),
+            row_group_count: AtomicU64::new(0),
+            projected_bytes_min: AtomicU64::new(u64::MAX),
+            projected_bytes_max: AtomicU64::new(0),
+            projected_bytes_total: AtomicU64::new(0),
+            current_inflight: AtomicUsize::new(0),
+            peak_inflight: AtomicUsize::new(0),
+        }
+    }
+}
+
+#[derive(Debug, Default, PartialEq, Eq)]
+pub(crate) struct ParquetReadDiagnosticsSnapshot {
+    pub(crate) row_group_count: u64,
+    pub(crate) projected_bytes_min: u64,
+    pub(crate) projected_bytes_max: u64,
+    pub(crate) projected_bytes_total: u64,
+    pub(crate) current_inflight: usize,
+    pub(crate) peak_inflight: usize,
 }
 
 impl ParquetReadBudget {
@@ -59,13 +97,66 @@ impl ParquetReadBudget {
             row_groups: Arc::new(Semaphore::new(parallelism)),
             bytes: Arc::new(Semaphore::new(byte_permits as usize)),
             byte_permits,
+            oversized_warning_logged: AtomicBool::new(false),
+            diagnostics: Arc::new(ParquetReadDiagnostics::default()),
         })
     }
 
     pub fn parallelism(&self) -> usize {
         self.parallelism
     }
 
+    pub(crate) fn enable_diagnostics(&self) {
+        self.diagnostics.enabled.store(true, Ordering::Relaxed);
+    }
+
+    pub(crate) fn diagnostics_enabled(&self) -> bool {
+        self.diagnostics.enabled.load(Ordering::Relaxed)
+    }
+
+    pub(crate) fn record_projected_row_groups(&self, projected_bytes: &[u64]) {
+        if !self.diagnostics_enabled() || projected_bytes.is_empty() {
+            return;
+        }
+        self.diagnostics
+            .row_group_count
+            .fetch_add(projected_bytes.len() as u64, Ordering::Relaxed);
+        self.diagnostics.projected_bytes_min.fetch_min(
+            *projected_bytes.iter().min().expect("checked non-empty"),
+            Ordering::Relaxed,
+        );
+        self.diagnostics.projected_bytes_max.fetch_max(
+            *projected_bytes.iter().max().expect("checked non-empty"),
+            Ordering::Relaxed,
+        );
+        self.diagnostics.projected_bytes_total.fetch_add(
+            projected_bytes
+                .iter()
+                .copied()
+                .fold(0u64, u64::saturating_add),
+            Ordering::Relaxed,
+        );
+    }
+
+    pub(crate) fn diagnostics(&self) -> ParquetReadDiagnosticsSnapshot {
+        let row_group_count = self.diagnostics.row_group_count.load(Ordering::Relaxed);
+        ParquetReadDiagnosticsSnapshot {
+            row_group_count,
+            projected_bytes_min: if row_group_count == 0 {
+                0
+            } else {
+                self.diagnostics.projected_bytes_min.load(Ordering::Relaxed)
+            },
+            projected_bytes_max: self.diagnostics.projected_bytes_max.load(Ordering::Relaxed),
+            projected_bytes_total: self
+                .diagnostics
+                .projected_bytes_total
+                .load(Ordering::Relaxed),
+            current_inflight: self.diagnostics.current_inflight.load(Ordering::Relaxed),
+            peak_inflight: self.diagnostics.peak_inflight.load(Ordering::Relaxed),
+        }
+    }
+
     pub(crate) async fn acquire(
         &self,
         projected_uncompressed_bytes: u64,
@@ -81,16 +172,39 @@ impl ParquetReadBudget {
             .max(1)
             .div_ceil(BYTE_PERMIT_UNIT)
             .min(u64::from(self.byte_permits)) as u32;
+        if projected_uncompressed_bytes > u64::from(self.byte_permits) * BYTE_PERMIT_UNIT
+            && !self.oversized_warning_logged.swap(true, Ordering::Relaxed)
+        {
+            log::warn!(
+                "Parquet row group projected size ({projected_uncompressed_bytes} bytes) exceeds \
+                 read.parquet.row-group.max-inflight-bytes ({} bytes); it will consume the entire \
+                 byte budget and may reduce row-group read parallelism; increase the option if \
+                 memory allows",
+                u64::from(self.byte_permits) * BYTE_PERMIT_UNIT
+            );
+        }
         let bytes = Arc::clone(&self.bytes)
             .acquire_many_owned(requested)
             .await
             .map_err(|error| crate::Error::UnexpectedError {
                 message: "Parquet byte read budget was closed".to_string(),
                 source: Some(Box::new(error)),
             })?;
+        let diagnostics = self.diagnostics_enabled().then(|| {
+            let current = self
+                .diagnostics
+                .current_inflight
+                .fetch_add(1, Ordering::Relaxed)
+                + 1;
+            self.diagnostics
+                .peak_inflight
+                .fetch_max(current, Ordering::Relaxed);
+            Arc::clone(&self.diagnostics)
+        });
         Ok(ParquetReadPermit {
             _row_group: row_group,
             _bytes: bytes,
+            diagnostics,
         })
     }
 }
@@ -106,6 +220,15 @@ impl Default for ParquetReadBudget {
 pub(crate) struct ParquetReadPermit {
     _row_group: OwnedSemaphorePermit,
     _bytes: OwnedSemaphorePermit,
+    diagnostics: Option<Arc<ParquetReadDiagnostics>>,
+}
+
+impl Drop for ParquetReadPermit {
+    fn drop(&mut self) {
+        if let Some(diagnostics) = &self.diagnostics {
+            diagnostics.current_inflight.fetch_sub(1, Ordering::Relaxed);
+        }
+    }
 }
 
 #[cfg(test)]
@@ -132,6 +255,32 @@ mod tests {
             .unwrap();
     }
 
+    #[tokio::test]
+    async fn diagnostics_aggregate_shared_row_group_reads() {
+        let budget = Arc::new(ParquetReadBudget::new(2, 2 * BYTE_PERMIT_UNIT).unwrap());
+        budget.enable_diagnostics();
+        budget.record_projected_row_groups(&[300, 100, 200]);
+
+        let first = budget.acquire(1).await.unwrap();
+        let second = budget.acquire(1).await.unwrap();
+        assert_eq!(
+            budget.diagnostics(),
+            ParquetReadDiagnosticsSnapshot {
+                row_group_count: 3,
+                projected_bytes_min: 100,
+                projected_bytes_max: 300,
+                projected_bytes_total: 600,
+                current_inflight: 2,
+                peak_inflight: 2,
+            }
+        );
+
+        drop(first);
+        drop(second);
+        assert_eq!(budget.diagnostics().current_inflight, 0);
+        assert_eq!(budget.diagnostics().peak_inflight, 2);
+    }
+
     #[test]
     fn rejects_invalid_limits() {
         assert!(ParquetReadBudget::new(0, BYTE_PERMIT_UNIT).is_err());
@@ -141,4 +290,47 @@ mod tests {
                 .is_err()
         );
     }
+
+    #[tokio::test]
+    async fn oversized_row_group_consumes_the_budget() {
+        let budget = Arc::new(ParquetReadBudget::new(8, 8 * BYTE_PERMIT_UNIT).unwrap());
+        let first = budget.acquire(100 * BYTE_PERMIT_UNIT).await.unwrap();
+        assert!(budget.oversized_warning_logged.load(Ordering::Relaxed));
+        assert!(
+            tokio::time::timeout(Duration::from_millis(20), budget.acquire(1))
+                .await
+                .is_err(),
+            "an oversized row group must consume the whole byte budget"
+        );
+        drop(first);
+        budget.acquire(1).await.unwrap();
+    }
+
+    #[tokio::test]
+    async fn small_row_groups_keep_exact_accounting() {
+        let budget = Arc::new(ParquetReadBudget::new(4, 4 * BYTE_PERMIT_UNIT).unwrap());
+        let mut permits = Vec::new();
+        for _ in 0..4 {
+            permits.push(budget.acquire(BYTE_PERMIT_UNIT).await.unwrap());
+        }
+        assert!(
+            tokio::time::timeout(Duration::from_millis(20), budget.acquire(BYTE_PERMIT_UNIT))
+                .await
+                .is_err()
+        );
+    }
+
+    #[tokio::test]
+    async fn tiny_budget_still_admits_one_at_a_time() {
+        let budget = Arc::new(ParquetReadBudget::new(8, BYTE_PERMIT_UNIT).unwrap());
+        let first = budget.acquire(100 * BYTE_PERMIT_UNIT).await.unwrap();
+        assert!(
+            tokio::time::timeout(Duration::from_millis(20), budget.acquire(1))
+                .await
+                .is_err(),
+            "a single-permit budget admits exactly one read"
+        );
+        drop(first);
+        budget.acquire(1).await.unwrap();
+    }
 }