leaves12138 commented on code in PR #800: URL: https://github.com/apache/paimon-rust/pull/800#discussion_r4043858471
########## crates/paimon/src/table/vindex_index_build_builder/pipeline.rs: ########## @@ -0,0 +1,1113 @@ +// 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. + +use super::extraction::{ + data_split_for_shard_ranges, extract_vector_batch, validate_vector_batch_ranges, +}; +use super::planning::VindexIndexShard; +use super::timing::{vector_index_build_timing_enabled, VectorIndexBuildTiming}; +use super::validation::{ + checked_row_count, checked_training_sample_index, checked_training_vector_count, + checked_vector_bytes, +}; +use super::writer::BuiltIndexFile; +use super::VindexIndexBuildBuilder; +use crate::arrow::format::parquet::{ + coalesced_parquet_range_bytes, parquet_granules, ParquetGranule, +}; +use crate::spec::ROW_ID_FIELD_NAME; +use crate::table::data_file_reader::DataFileReadTiming; +use crate::table::table_read::configured_parquet_read_budget; +use crate::table::{merge_row_ranges, ArrowRecordBatchStream, RowRange}; +use crate::vindex::VindexVectorIndexOptions; +use crate::{Error, Result}; +use arrow_array::RecordBatch; +use arrow_buffer::MutableBuffer; +use futures::{StreamExt, TryStreamExt}; +use paimon_vindex_core::autotune::default_training_vector_count; +use paimon_vindex_core::index::{VectorIndexTrainer, VectorIndexTraining, VectorIndexWriter}; +use std::collections::{HashMap, HashSet}; +use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write}; +use std::ops::Range; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; + +const MIN_STRATA: usize = 256; +const ROWS_PER_STRATUM: usize = 128; +const FIRST_BYTES_NUMERATOR: u64 = 3; +const FIRST_BYTES_DENOMINATOR: u64 = 10; +const QUEUE_CAPACITY: usize = 2; +const BUFFER_BYTES: usize = 8 * 1024 * 1024; +const REPLAY_TARGET_BYTES: usize = 32 * 1024 * 1024; + +#[derive(Clone, Debug)] +struct Granule { + range: RowRange, + file_index: usize, + byte_ranges: Vec<Range<u64>>, +} + +#[derive(Debug)] +pub(super) struct GranulePlan { + pub(super) first: Vec<RowRange>, + pub(super) rest: Vec<RowRange>, + first_rows: usize, +} + +struct SpillRecord { + ids: Vec<i64>, + bytes: Vec<u8>, +} + +enum AddItem { + Batch(RecordBatch, Vec<i64>), + Spilled(Vec<i64>, MutableBuffer), +} + +type SpillTask = JoinHandle<std::io::Result<(std::fs::File, u64, Duration)>>; +type ConsumerTask = JoinHandle<Result<(VectorIndexWriter, usize, usize, Duration)>>; +type TrainingTask = JoinHandle<std::io::Result<(VectorIndexTraining, Duration)>>; +type ReplayTask = JoinHandle<std::io::Result<(usize, Duration)>>; + +struct SpillWriter { + sender: mpsc::Sender<SpillRecord>, + task: SpillTask, +} + +impl SpillWriter { + async fn finish(self) -> Result<(std::fs::File, u64, Duration)> { + drop(self.sender); + join_spill(self.task).await + } +} + +struct LivePipeline { + sender: mpsc::Sender<AddItem>, + consumer: ConsumerTask, + replay: ReplayTask, + spill_bytes: u64, + spill_write: Duration, +} + +fn spawn_spill_writer(timing_enabled: bool) -> Result<SpillWriter> { + let file = tempfile::tempfile().map_err(|e| Error::UnexpectedError { + message: format!("Failed to create temporary vindex vector file: {e}"), + source: Some(Box::new(e)), + })?; + let (sender, mut receiver) = mpsc::channel::<SpillRecord>(QUEUE_CAPACITY); + let task = tokio::task::spawn_blocking(move || -> std::io::Result<_> { + let mut writer = BufWriter::with_capacity(BUFFER_BYTES, file); + let mut spill_bytes = 0u64; + let mut spill_write = Duration::ZERO; + while let Some(record) = receiver.blocking_recv() { + let write_start = timing_enabled.then(Instant::now); + let count = record.ids.len() as u64; + writer.write_all(&count.to_le_bytes())?; + for id in record.ids { + writer.write_all(&id.to_le_bytes())?; + } + writer.write_all(&record.bytes)?; + spill_bytes = spill_bytes + .saturating_add(8) + .saturating_add(count.saturating_mul(8)) + .saturating_add(record.bytes.len() as u64); + if let Some(start) = write_start { + spill_write = spill_write.saturating_add(start.elapsed()); + } + } + let write_start = timing_enabled.then(Instant::now); + writer.flush()?; + let file = writer.into_inner().map_err(|e| e.into_error())?; + if let Some(start) = write_start { + spill_write = spill_write.saturating_add(start.elapsed()); + } + Ok((file, spill_bytes, spill_write)) + }); + Ok(SpillWriter { sender, task }) +} + +async fn join_spill(task: SpillTask) -> Result<(std::fs::File, u64, Duration)> { + task.await + .map_err(|e| Error::UnexpectedError { + message: format!("vindex spill task failed: {e}"), + source: None, + })? + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to spill vindex vectors: {e}"), + source: Some(Box::new(e)), + }) +} + +fn spawn_add_consumer( + writer: VectorIndexWriter, + mut receiver: mpsc::Receiver<AddItem>, + index_column: String, + dimension: usize, + timing_enabled: bool, +) -> ConsumerTask { + tokio::task::spawn_blocking(move || -> Result<_> { + let mut writer = writer; + let mut rows_added = 0usize; + let mut replay_rows = 0usize; + let mut index_add = Duration::ZERO; + while let Some(item) = receiver.blocking_recv() { + let add_start = timing_enabled.then(Instant::now); + match item { + AddItem::Batch(batch, ids) => { + let vectors = extract_vector_batch(&batch, &index_column, dimension)?; + if ids.len() != vectors.row_count { + return Err(Error::DataInvalid { + message: "vindex add batch id count mismatch".to_string(), + source: None, + }); + } + writer + .add_vectors(&ids, vectors.values, vectors.row_count) + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to add vectors to vindex index: {e}"), + source: Some(Box::new(e)), + })?; + rows_added += vectors.row_count; + } + AddItem::Spilled(ids, buffer) => { + let values = buffer.typed_data::<f32>(); + if values.len() != ids.len() * dimension { + return Err(Error::DataInvalid { + message: "vindex spilled vector length mismatch".to_string(), + source: None, + }); + } + writer.add_vectors(&ids, values, ids.len()).map_err(|e| { + Error::UnexpectedError { + message: format!("Failed to add spilled vectors to vindex index: {e}"), + source: Some(Box::new(e)), + } + })?; + rows_added += ids.len(); + replay_rows += ids.len(); + } + } + if let Some(start) = add_start { + index_add = index_add.saturating_add(start.elapsed()); + } + } + Ok((writer, rows_added, replay_rows, index_add)) + }) +} + +async fn join_consumer(task: ConsumerTask) -> Result<(VectorIndexWriter, usize, usize, Duration)> { + task.await.map_err(|e| Error::UnexpectedError { + message: format!("vindex add task failed: {e}"), + source: None, + })? +} + +async fn join_training(task: TrainingTask) -> Result<(VectorIndexTraining, Duration)> { + task.await + .map_err(|e| Error::UnexpectedError { + message: format!("vindex training task failed: {e}"), + source: None, + })? + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to train vindex index: {e}"), + source: Some(Box::new(e)), + }) +} + +async fn start_live_pipeline( + training: TrainingTask, + spill: SpillWriter, + index_column: String, + dimension: usize, + timing_enabled: bool, +) -> Result<(LivePipeline, Duration)> { + let trained = join_training(training).await; + let spilled = spill.finish().await; + let (trained, train_finish) = trained?; + let (file, spill_bytes, spill_write) = spilled?; + let (sender, receiver) = mpsc::channel(QUEUE_CAPACITY); + let consumer = spawn_add_consumer( + VectorIndexWriter::new(trained), + receiver, + index_column, + dimension, + timing_enabled, + ); + let replay = spawn_replay(file, sender.clone(), dimension, timing_enabled); + Ok(( + LivePipeline { + sender, + consumer, + replay, + spill_bytes, + spill_write, + }, + train_finish, + )) +} + +async fn finish_live_pipeline( + pipeline: LivePipeline, +) -> ( + Result<(VectorIndexWriter, usize, usize, Duration)>, + Result<(usize, Duration)>, + u64, + Duration, +) { + let replay = pipeline + .replay + .await + .map_err(|e| Error::UnexpectedError { + message: format!("vindex replay task failed: {e}"), + source: None, + }) + .and_then(|result| { + result.map_err(|e| Error::UnexpectedError { + message: format!("Failed to replay spilled vindex vectors: {e}"), + source: Some(Box::new(e)), + }) + }); + drop(pipeline.sender); + let consumer = join_consumer(pipeline.consumer).await; + (consumer, replay, pipeline.spill_bytes, pipeline.spill_write) +} + +fn spawn_replay( + mut file: std::fs::File, + sender: mpsc::Sender<AddItem>, + dimension: usize, + timing_enabled: bool, +) -> ReplayTask { + tokio::task::spawn_blocking(move || -> std::io::Result<(usize, Duration)> { + let mut spill_read = Duration::ZERO; + let read_start = timing_enabled.then(Instant::now); + file.seek(SeekFrom::Start(0))?; + if let Some(start) = read_start { + spill_read = spill_read.saturating_add(start.elapsed()); + } + let mut reader = BufReader::with_capacity(BUFFER_BYTES, file); + let mut rows = 0usize; + let mut ids = Vec::new(); + let mut vectors = MutableBuffer::new(REPLAY_TARGET_BYTES); + loop { + let read_start = timing_enabled.then(Instant::now); + let mut header = [0u8; 8]; + let read = reader.read(&mut header)?; + if read == 0 { + if let Some(start) = read_start { + spill_read = spill_read.saturating_add(start.elapsed()); + } + break; + } + reader.read_exact(&mut header[read..])?; + let count = usize::try_from(u64::from_le_bytes(header)).map_err(|_| { + std::io::Error::new(std::io::ErrorKind::InvalidData, "invalid spill row count") + })?; + let id_bytes_len = count.checked_mul(8).ok_or_else(|| { + std::io::Error::new(std::io::ErrorKind::InvalidData, "spill id length overflow") + })?; + let mut id_bytes = vec![0u8; id_bytes_len]; + reader.read_exact(&mut id_bytes)?; + let (id_chunks, remainder) = id_bytes.as_chunks::<8>(); + debug_assert!(remainder.is_empty()); + ids.extend(id_chunks.iter().map(|bytes| i64::from_le_bytes(*bytes))); + let vector_bytes = count + .checked_mul(dimension) + .and_then(|value| value.checked_mul(4)) + .ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + "spill vector length overflow", + ) + })?; + let offset = vectors.len(); + vectors.resize(offset + vector_bytes, 0); + reader.read_exact(&mut vectors.as_slice_mut()[offset..])?; + if let Some(start) = read_start { + spill_read = spill_read.saturating_add(start.elapsed()); + } + rows += count; + if vectors.len() >= REPLAY_TARGET_BYTES { + let item = AddItem::Spilled( + std::mem::take(&mut ids), + std::mem::replace(&mut vectors, MutableBuffer::new(REPLAY_TARGET_BYTES)), + ); + if sender.blocking_send(item).is_err() { + return Ok((rows, spill_read)); + } + } + } + if !ids.is_empty() { + let _ = sender.blocking_send(AddItem::Spilled(ids, vectors)); + } + Ok((rows, spill_read)) + }) +} + +fn pick_indices(total: usize, count: usize) -> impl Iterator<Item = usize> { + (0..count).map(move |index| ((2 * index + 1) * total) / (2 * count)) Review Comment: [P2] Avoid fixed midpoint granule sampling that can exclude an entire distribution Verified on `f934e8a0fdae78d382adaffe2b05696e30070398`, using the PR's locked `paimon-vindex-core 0.4.0` dependency. The midpoint selection here is deterministic systematic sampling, rather than randomized stratified sampling. For 4,096 granules and a target of 512, it always selects zero-based granule indices `4, 12, 20, ...`. Data correlated with file/page order can therefore be completely absent from training. The trainer's subsequent reservoir sampling cannot recover a distribution that never reaches it. I reproduced this through the actual granule planner, Parquet reads, index build/serialization, and search, not just a simulation of the selection formula: - 4,096 real Parquet files in the memory filesystem, 256 rows per file, totaling 1,048,576 rows in one shard. - Every eighth file (zero-based file index `% 8 == 7`) contains `[100.0]`; the other files alternate `[0.0]` and `[1.0]`. The excluded cluster represents 12.5% of all rows. - IVF-SQ, L2, `dimension=1`, `nlist=nprobe=1`, and `train.sample-ratio=1.0`. - The planner selects 512 first ranges / 131,072 candidate rows, containing no `[100.0]` vectors. - With `vindex.build.granule.enabled=false`, querying `[100.0]` returns 10/10 results from that cluster. With it enabled, the same query returns 0/10. These are cluster-hit counts, not exact-ID recall, since distances tie within a cluster. All rows are still added to the index; the regression is in training quality. SQ learns only the `[0, 1]` range and clips the unseen `[100]` vectors to the same upper code. I also reproduced a tail-only distribution failure in a separate actual-planner-plus-trainer comparison. Could we avoid fixed granule positions and add a regression test for periodic/file-correlated distributions? One conservative option is seeded randomized stratification over logical row positions, expanding selected rows to pages for I/O without treating every fetched row as a training sample. If the required page reads eliminate the savings, retain the full-shard sampling fallback. Randomizing whole granules is less intrusive, but still needs explicit evaluation of rare-cluster coverage and unequal granule sizes. Until that trade-off is validated, I would keep the conservative path as the default. The existing 24 vindex-build tests and 58 Parquet tests pass; they do not cover this case. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
