diff --git a/crates/paimon/src/spec/core_options.rs b/crates/paimon/src/spec/core_options.rs index 79354eea7..0b6c8b297 100644 --- a/crates/paimon/src/spec/core_options.rs +++ b/crates/paimon/src/spec/core_options.rs @@ -31,6 +31,7 @@ const GLOBAL_INDEX_SEARCH_MODE_OPTION: &str = "global-index.search-mode"; const SCALAR_INDEX_SEARCH_MODE_OPTION: &str = "scalar-index.search-mode"; const VECTOR_INDEX_SEARCH_MODE_OPTION: &str = "vector-index.search-mode"; const FULL_TEXT_INDEX_SEARCH_MODE_OPTION: &str = "full-text-index.search-mode"; +const GLOBAL_INDEX_BUILD_PARALLELISM_OPTION: &str = "global-index.build.parallelism"; const GLOBAL_INDEX_ROW_COUNT_PER_SHARD_OPTION: &str = "global-index.row-count-per-shard"; const GLOBAL_INDEX_THREAD_NUM_OPTION: &str = "global-index.thread-num"; const GLOBAL_INDEX_VINDEX_READ_THREAD_NUM_OPTION: &str = "global-index.vindex.read-thread-num"; @@ -1057,6 +1058,19 @@ impl<'a> CoreOptions<'a> { Ok(value) } + /// Number of generic global-index shards prepared concurrently by a local + /// build, matching PyPaimon's option and default. Each shard may also use + /// native worker threads; the default remains one. + pub fn global_index_build_parallelism(&self) -> crate::Result { + match self.options.get(GLOBAL_INDEX_BUILD_PARALLELISM_OPTION) { + None => Ok(1), + Some(raw) => raw.trim().parse::().ok().filter(|value| *value > 0) + .ok_or_else(|| crate::Error::ConfigInvalid { + message: format!("Option '{GLOBAL_INDEX_BUILD_PARALLELISM_OPTION}' must be greater than 0, got: {raw}"), + }), + } + } + /// Maximum number of concurrent global-index search tasks, mirroring Java /// `CoreOptions.GLOBAL_INDEX_THREAD_NUM` (key `global-index.thread-num`, /// default 32). Used as the per-operation fan-out limit for sorted BTree and @@ -2534,6 +2548,35 @@ mod tests { } } + #[test] + fn test_global_index_build_parallelism() { + assert_eq!( + CoreOptions::new(&HashMap::new()) + .global_index_build_parallelism() + .unwrap(), + 1 + ); + for (raw, expected) in [("1", 1), (" 3 ", 3), ("+2", 2), ("1000000", 1000000)] { + let options = + HashMap::from([(GLOBAL_INDEX_BUILD_PARALLELISM_OPTION.into(), raw.into())]); + assert_eq!( + CoreOptions::new(&options) + .global_index_build_parallelism() + .unwrap(), + expected + ); + } + for raw in ["0", "-1", "invalid", "18446744073709551616"] { + let options = + HashMap::from([(GLOBAL_INDEX_BUILD_PARALLELISM_OPTION.into(), raw.into())]); + assert!(CoreOptions::new(&options) + .global_index_build_parallelism() + .unwrap_err() + .to_string() + .contains("parallelism")); + } + } + #[test] fn test_global_index_thread_num_default_and_custom() { // Default mirrors Java (32) when unset. diff --git a/crates/paimon/src/table/full_text_index_build_builder.rs b/crates/paimon/src/table/full_text_index_build_builder.rs index 0ee0924fb..49802a1b5 100644 --- a/crates/paimon/src/table/full_text_index_build_builder.rs +++ b/crates/paimon/src/table/full_text_index_build_builder.rs @@ -34,7 +34,7 @@ use crate::table::global_index_build_common::{ copy_local_file_to_output, indexed_row_ranges, validate_existing_index_overlap, }; use crate::table::global_index_types::FULL_TEXT_GLOBAL_INDEX_TYPE; -use crate::table::{CommitMessage, DataSplitBuilder, RowRange, Table, TableCommit}; +use crate::table::{CommitMessage, DataSplitBuilder, RowRange, Table}; use crate::{Error, Result}; use arrow_array::cast::AsArray; use arrow_array::{Array, Int64Array, LargeStringArray, RecordBatch, StringArray, StringViewArray}; @@ -98,8 +98,9 @@ impl<'a> FullTextIndexBuildBuilder<'a> { // procedure options. let mut merged_options = self.table.schema().options().clone(); merged_options.extend(self.options.clone()); - let rows_per_shard = - CoreOptions::new(&merged_options).global_index_row_count_per_shard()?; + let core_options = CoreOptions::new(&merged_options); + let rows_per_shard = core_options.global_index_row_count_per_shard()?; + let parallelism = core_options.global_index_build_parallelism()?; let index_field = find_index_field(self.table, index_column)?; validate_text_field(index_field)?; @@ -161,43 +162,34 @@ impl<'a> FullTextIndexBuildBuilder<'a> { ) .await?; - let commit = TableCommit::new( - self.table.clone(), - format!( - "global-index-{}-create-{}", - FULL_TEXT_GLOBAL_INDEX_TYPE, - uuid::Uuid::new_v4() - ), - ); - let mut messages: Vec = Vec::with_capacity(shards.len()); - for shard in shards { - let build_result = self - .build_shard( - &shard, - index_column, - &native_options, - index_field.id(), - &index_meta, - ) - .await; - let index_file = match build_result { - Ok(index_file) => index_file, - Err(error) => { - let _ = commit.abort(&messages).await; - return Err(error); + let prepared = super::global_index_build_common::preparation::prepare_shards( + self.table, + shards, + parallelism, + |shard| { + let native_options = &native_options; + let index_meta = &index_meta; + async move { + let file = self + .build_shard( + &shard, + index_column, + native_options, + index_field.id(), + index_meta, + ) + .await?; + // Java skips a shard whose writer saw no rows. + Ok(file.map(|file| { + let mut message = CommitMessage::new(shard.partition_bytes, 0, vec![]); + message.new_index_files = vec![file]; + (message, ()) + })) } - }; - // Java skips a shard whose writer saw no rows. - let Some(index_file) = index_file else { - continue; - }; - let mut message = CommitMessage::new(shard.partition_bytes.clone(), 0, vec![]); - message.new_index_files = vec![index_file]; - messages.push(message); - } - if messages.is_empty() { - return Ok((Some(snapshot.id()), vec![])); - } + }, + ) + .await?; + let messages = prepared.into_iter().map(|(message, ())| message).collect(); Ok((Some(snapshot.id()), messages)) } diff --git a/crates/paimon/src/table/global_index_build_builder.rs b/crates/paimon/src/table/global_index_build_builder.rs index 428de8225..a16e6a14a 100644 --- a/crates/paimon/src/table/global_index_build_builder.rs +++ b/crates/paimon/src/table/global_index_build_builder.rs @@ -98,6 +98,8 @@ impl<'a> GlobalIndexBuildBuilder<'a> { self } + /// Overrides table options for this build. Full-text and vindex builds use + /// `global-index.build.parallelism` to bound concurrent shards (default: 1). pub fn with_options(&mut self, options: HashMap) -> &mut Self { self.options = options; self @@ -346,3 +348,7 @@ mod tests { ); } } + +#[cfg(test)] +#[path = "global_index_build_builder/parallel_tests.rs"] +mod parallel_tests; diff --git a/crates/paimon/src/table/global_index_build_builder/parallel_tests.rs b/crates/paimon/src/table/global_index_build_builder/parallel_tests.rs new file mode 100644 index 000000000..e87376e56 --- /dev/null +++ b/crates/paimon/src/table/global_index_build_builder/parallel_tests.rs @@ -0,0 +1,339 @@ +// 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. + +//! Integration coverage for real generic writers using the shared shard runner. + +use super::*; +use crate::catalog::Identifier; +use crate::io::FileIOBuilder; +use crate::spec::{ArrayType, DataType, FloatType, IntType, Schema, TableSchema, VarCharType}; +use crate::table::TableWrite; +use arrow_array::builder::{Float32Builder, ListBuilder}; +use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch, StringArray}; +use arrow_schema::{DataType as ArrowType, Field, Schema as ArrowSchema}; +use std::sync::Arc; + +fn table(kind: &str, parallelism: usize, external: bool) -> Table { + let location = format!("memory:/generic-parallel-{kind}-{}", uuid::Uuid::new_v4()); + let mut options = HashMap::from([ + ("row-tracking.enabled".into(), "true".into()), + ("data-evolution.enabled".into(), "true".into()), + ("global-index.enabled".into(), "true".into()), + ("bucket".into(), "-1".into()), + ("global-index.row-count-per-shard".into(), "4".into()), + ( + "global-index.build.parallelism".into(), + parallelism.to_string(), + ), + ("vector-index.search-mode".into(), "fast".into()), + ("full-text-index.search-mode".into(), "fast".into()), + ]); + if external { + options.insert( + "global-index.external-path".into(), + format!("{location}-external"), + ); + } + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("text", DataType::VarChar(VarCharType::string_type())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ) + .options(options) + .build() + .unwrap(); + Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "parallel"), + location, + TableSchema::new(0, &schema), + None, + ) +} + +async fn append(table: &Table, start: usize, vectors: Vec>>>) { + let rows = vectors.len(); + let mut values = ListBuilder::new(Float32Builder::new()); + for vector in vectors { + match vector { + Some(vector) => { + for value in vector { + values.values().append_option(value); + } + values.append(true); + } + None => values.append(false), + } + } + let values = values.finish(); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", ArrowType::Int32, false), + Field::new("text", ArrowType::Utf8, true), + Field::new("embedding", values.data_type().clone(), true), + ])); + let columns: Vec = vec![ + Arc::new(Int32Array::from( + (start..start + rows).map(|i| i as i32).collect::>(), + )), + Arc::new(StringArray::from( + (start..start + rows) + .map(|i| (i % 2 == 1).then_some("paimon index")) + .collect::>(), + )), + Arc::new(values), + ]; + let mut write = TableWrite::new(table, "data".into()).unwrap(); + write + .write_arrow_batch(&RecordBatch::try_new(schema, columns).unwrap()) + .await + .unwrap(); + TableCommit::new(table.clone(), "data".into()) + .commit(write.prepare_commit().await.unwrap()) + .await + .unwrap(); +} + +fn builder<'a>(table: &'a Table, kind: &str) -> GlobalIndexBuildBuilder<'a> { + let mut builder = table.new_global_index_build_builder(); + builder + .with_index_column(if kind == "full-text" { + "text" + } else { + "embedding" + }) + .with_index_type(kind); + if kind != "full-text" { + let mut options = HashMap::from([(format!("{kind}.dimension"), "2".into())]); + if kind != "diskann" { + options.insert(format!("{kind}.nlist"), "1".into()); + } + if kind == "ivf-pq" { + options.insert("ivf-pq.pq.m".into(), "1".into()); + } + if kind == "diskann" { + options.insert("diskann.pq.bits".into(), "4".into()); + } + builder.with_options(options); + } + builder +} + +fn path(table: &Table, file: &crate::spec::IndexFileMeta) -> String { + file.external_path + .clone() + .unwrap_or_else(|| format!("{}/index/{}", table.location(), file.file_name)) +} + +#[tokio::test] +async fn parallel_native_writers_match_serial_shards_and_incremental_coverage() { + let mut kinds = vec!["ivf-flat", "ivf-pq", "ivf-sq", "ivf-rq", "diskann"]; + if cfg!(feature = "fulltext") { + kinds.push("full-text"); + } + for kind in kinds { + for external in [false, true] { + let serial = table(kind, 1, external); + let parallel = table(kind, 3, external); + // The first vector shard is all NULL; subsequent shards have sparse IDs. + let rows = (0..20) + .map(|i| (i >= 4 && i % 2 == 1).then_some(vec![Some(1.), Some(0.)])) + .collect::>(); + append(&serial, 0, rows.clone()).await; + append(¶llel, 0, rows).await; + let control = builder(&serial, kind).build().await.unwrap(); + let messages = builder(¶llel, kind).build().await.unwrap(); + assert_eq!(messages.len(), control.len(), "{kind}"); + assert_eq!( + parallel + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap() + .id(), + 1 + ); + for (actual, expected) in messages.iter().zip(&control) { + let file = &actual.new_index_files[0]; + let reference = &expected.new_index_files[0]; + assert_eq!(file.row_count, reference.row_count); + let mut actual_meta = file.global_index_meta.clone().unwrap(); + let mut expected_meta = reference.global_index_meta.clone().unwrap(); + let actual_options: serde_json::Value = + serde_json::from_slice(&actual_meta.index_meta.take().unwrap()).unwrap(); + let expected_options: serde_json::Value = + serde_json::from_slice(&expected_meta.index_meta.take().unwrap()).unwrap(); + assert_eq!(actual_meta, expected_meta); + assert_eq!(actual_options, expected_options); + let actual = parallel + .file_io() + .new_input(&path(¶llel, file)) + .unwrap() + .read() + .await + .unwrap(); + let expected = serial + .file_io() + .new_input(&path(&serial, reference)) + .unwrap() + .read() + .await + .unwrap(); + // Full-text archives contain generated segment identities; vector files are deterministic. + if kind != "full-text" { + assert_eq!(actual, expected, "{kind}"); + } + } + TableCommit::new(parallel.clone(), "indexes".into()) + .commit(messages) + .await + .unwrap(); + assert!(builder(¶llel, kind).build().await.unwrap().is_empty()); + append( + ¶llel, + 20, + vec![ + None, + Some(vec![Some(0.), Some(1.)]), + None, + Some(vec![Some(1.), Some(1.)]), + ], + ) + .await; + let new = builder(¶llel, kind).build().await.unwrap(); + assert_eq!(new.len(), 1); + assert_eq!( + new[0].new_index_files[0] + .global_index_meta + .as_ref() + .unwrap() + .row_range_start, + 20 + ); + TableCommit::new(parallel.clone(), "new-indexes".into()) + .commit(new) + .await + .unwrap(); + if kind != "full-text" { + let result = parallel + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1., 1.]) + .with_limit(64) + .execute() + .await + .unwrap(); + let mut ids = result.row_ids().unwrap().row_ids.clone(); + ids.sort_unstable(); + assert_eq!(ids, vec![5, 7, 9, 11, 13, 15, 17, 19, 21, 23], "{kind}"); + } + #[cfg(feature = "fulltext")] + if kind == "full-text" { + let result = parallel + .new_full_text_search_builder() + .with_text_column("text") + .with_query_text("paimon") + .with_limit(64) + .execute() + .await + .unwrap(); + let mut ids = result + .into_iter() + .map(|range| range.from()) + .collect::>(); + ids.sort_unstable(); + assert_eq!(ids, (1..24).step_by(2).collect::>()); + } + } + } +} + +#[tokio::test] +async fn failed_parallel_shards_clean_private_indexes_and_retain_handed_off_files() { + for external in [false, true] { + let table = table("ivf-flat", 3, external); + append(&table, 0, vec![Some(vec![Some(1.), Some(0.)]); 12]).await; + // These files are already owned by a caller and must survive a later failure. + let handed_off = builder(&table, "ivf-sq").build().await.unwrap(); + append(&table, 12, vec![Some(vec![Some(1.), None]); 4]).await; + let before = table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap() + .id(); + let error = builder(&table, "ivf-flat").execute().await.unwrap_err(); + assert!(error.to_string().contains("null vector element")); + assert_eq!( + table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap() + .id(), + before + ); + for message in &handed_off { + assert!(table + .file_io() + .exists(&path(&table, &message.new_index_files[0])) + .await + .unwrap()); + } + let directory = if external { + format!("{}-external", table.location()) + } else { + format!("{}/index", table.location()) + }; + let remaining = table.file_io().list_status(&directory).await.unwrap(); + assert_eq!( + remaining.len(), + handed_off.len(), + "private files from in-flight workers leaked" + ); + } +} + +#[tokio::test] +async fn invalid_parallelism_is_rejected_for_every_generic_family_before_snapshot_reads() { + let mut kinds = vec!["ivf-flat", "ivf-pq", "ivf-sq", "ivf-rq", "diskann"]; + if cfg!(feature = "fulltext") { + kinds.push("full-text"); + } + for kind in kinds { + let table = table(kind, 1, false); + for value in ["0", "-1", "invalid"] { + let mut builder = builder(&table, kind); + let mut options = builder.options.clone(); + options.insert("global-index.build.parallelism".into(), value.into()); + builder.with_options(options); + assert!( + builder + .build() + .await + .unwrap_err() + .to_string() + .contains("parallelism"), + "{kind}" + ); + } + } +} diff --git a/crates/paimon/src/table/global_index_build_common.rs b/crates/paimon/src/table/global_index_build_common.rs index 77775b18d..8466624e3 100644 --- a/crates/paimon/src/table/global_index_build_common.rs +++ b/crates/paimon/src/table/global_index_build_common.rs @@ -33,6 +33,7 @@ use bytes::Bytes; use std::path::Path; use tokio::io::AsyncReadExt; +pub(crate) mod preparation; pub(crate) mod vector; const COPY_BUFFER_SIZE: usize = 1024 * 1024; diff --git a/crates/paimon/src/table/global_index_build_common/preparation.rs b/crates/paimon/src/table/global_index_build_common/preparation.rs new file mode 100644 index 000000000..6dafba3ef --- /dev/null +++ b/crates/paimon/src/table/global_index_build_common/preparation.rs @@ -0,0 +1,127 @@ +// 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. + +//! Bounded preparation of independently owned global-index shards. + +use crate::table::{CommitMessage, Table, TableCommit}; +use crate::{Error, Result}; +use futures::stream::FuturesUnordered; +use futures::{FutureExt, StreamExt}; +use std::future::Future; +use std::panic::AssertUnwindSafe; + +/// Prepare shards against a snapshot selected by the caller. Each shard has its +/// own reader and writer, like Java GenericIndexTopoBuilder.BuildIndexOperator. +/// +/// Results are returned in plan order, even if shards finish out of order. On +/// failure, stop admitting work and drain every started shard before cleaning +/// this preparation's private outputs. Dropping in-flight futures is unsafe: +/// their native blocking workers could still be writing the index files. +/// +/// This helper owns every returned message until it returns successfully. Never +/// use it to clean up messages already handed to callers or submitted for commit. +pub(crate) async fn prepare_shards( + table: &Table, + shards: Vec, + parallelism: usize, + mut build: F, +) -> Result> +where + F: FnMut(S) -> Fut, + Fut: Future>>, +{ + if parallelism == 0 { + return Err(Error::ConfigInvalid { + message: "Option 'global-index.build.parallelism' must be greater than 0.".into(), + }); + } + let mut pending = shards.into_iter().enumerate(); + let mut active = FuturesUnordered::new(); + for _ in 0..parallelism { + let Some((ordinal, shard)) = pending.next() else { + break; + }; + active.push(run_shard(ordinal, build(shard))); + } + let mut completed = Vec::new(); + let mut failure = None; + while let Some((ordinal, result)) = active.next().await { + match result { + Ok(Some(output)) => completed.push((ordinal, output)), + Ok(None) => {} + Err(error) => { + if failure.is_none() { + failure = Some(error); + } + } + } + if failure.is_none() { + if let Some((ordinal, shard)) = pending.next() { + active.push(run_shard(ordinal, build(shard))); + } + } + } + if let Some(error) = failure { + // Every started worker has finished; no file can appear after cleanup. + let messages = completed + .into_iter() + .map(|(_, (message, _))| message) + .collect::>(); + abort_private_outputs(table, &messages).await; + return Err(error); + } + completed.sort_unstable_by_key(|(ordinal, _)| *ordinal); + Ok(completed.into_iter().map(|(_, output)| output).collect()) +} + +async fn run_shard( + ordinal: usize, + build: Fut, +) -> (usize, Result>) +where + Fut: Future>>, +{ + let result = AssertUnwindSafe(build) + .catch_unwind() + .await + .unwrap_or_else(|panic| { + let reason = panic + .downcast_ref::() + .map(String::as_str) + .or_else(|| panic.downcast_ref::<&str>().copied()) + .unwrap_or("unknown panic"); + Err(Error::UnexpectedError { + message: format!("Global index shard {ordinal} panicked: {reason}"), + source: None, + }) + }); + (ordinal, result) +} + +async fn abort_private_outputs(table: &Table, messages: &[CommitMessage]) { + if !messages.is_empty() { + let commit = TableCommit::new( + table.clone(), + format!("global-index-private-build-{}", uuid::Uuid::new_v4()), + ); + // Preserve the build error if best-effort cleanup itself fails. + let _ = commit.abort(messages).await; + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/paimon/src/table/global_index_build_common/preparation/tests.rs b/crates/paimon/src/table/global_index_build_common/preparation/tests.rs new file mode 100644 index 000000000..8da72656a --- /dev/null +++ b/crates/paimon/src/table/global_index_build_common/preparation/tests.rs @@ -0,0 +1,348 @@ +// 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::*; +use crate::catalog::Identifier; +use crate::io::FileIOBuilder; +use crate::spec::{DataType, GlobalIndexMeta, IndexFileMeta, IntType, Schema, TableSchema}; +use bytes::Bytes; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::{mpsc, oneshot}; +use tokio::task::JoinHandle; + +fn table() -> Table { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .build() + .unwrap(); + Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "concurrency"), + format!("memory:/index-concurrency-{}", uuid::Uuid::new_v4()), + TableSchema::new(0, &schema), + None, + ) +} + +#[derive(Clone, Copy)] +enum Outcome { + File, + ExternalFile, + Empty, + Fail, + Panic, +} + +#[derive(Default)] +struct Counters { + admitted: AtomicUsize, + active: AtomicUsize, + peak: AtomicUsize, +} + +struct Active(Arc); +impl Active { + fn new(counters: Arc) -> Self { + let count = counters.active.fetch_add(1, Ordering::SeqCst) + 1; + counters.peak.fetch_max(count, Ordering::SeqCst); + Self(counters) + } +} +impl Drop for Active { + fn drop(&mut self) { + self.0.active.fetch_sub(1, Ordering::SeqCst); + } +} + +async fn private_file(table: &Table, ordinal: usize, external: bool) -> CommitMessage { + let name = format!("private-{ordinal}.index"); + let external_path = external.then(|| format!("{}-external/{name}", table.location())); + let path = external_path + .clone() + .unwrap_or_else(|| format!("{}/index/{name}", table.location())); + table + .file_io() + .new_output(&path) + .unwrap() + .write(Bytes::from_static(b"private")) + .await + .unwrap(); + let mut message = CommitMessage::new(vec![], 0, vec![]); + message.new_index_files.push(IndexFileMeta { + index_type: "ivf-flat".into(), + file_name: name, + file_size: 7, + row_count: 1, + deletion_vectors_ranges: None, + external_path, + global_index_meta: Some(GlobalIndexMeta { + row_range_start: ordinal as i64, + row_range_end: ordinal as i64, + index_field_id: 0, + extra_field_ids: None, + source_meta: None, + index_meta: None, + }), + }); + message +} + +type Build = JoinHandle>>; +struct Running { + task: Build, + started: mpsc::UnboundedReceiver, + releases: Vec>>, + counters: Arc, + finished: mpsc::UnboundedReceiver, +} + +fn start(table: &Table, parallelism: usize, outcomes: Vec) -> Running { + let (started, receiver) = mpsc::unbounded_channel(); + let (finished, finished_receiver) = mpsc::unbounded_channel(); + let mut releases = Vec::new(); + let jobs = outcomes + .into_iter() + .enumerate() + .map(|(ordinal, outcome)| { + let (sender, receiver) = oneshot::channel(); + releases.push(Some(sender)); + (ordinal, outcome, receiver) + }) + .collect(); + let owner = table.clone(); + let worker_table = table.clone(); + let counters = Arc::new(Counters::default()); + let observed = counters.clone(); + let task = tokio::spawn(async move { + prepare_shards( + &owner, + jobs, + parallelism, + move |(ordinal, outcome, release): (_, _, oneshot::Receiver<()>)| { + let table = worker_table.clone(); + let started = started.clone(); + let observed = observed.clone(); + let finished = finished.clone(); + observed.admitted.fetch_add(1, Ordering::SeqCst); + async move { + let _active = Active::new(observed); + let message = match outcome { + Outcome::File | Outcome::ExternalFile => Some( + private_file(&table, ordinal, matches!(outcome, Outcome::ExternalFile)) + .await, + ), + _ => None, + }; + started.send(ordinal).unwrap(); + release.await.unwrap(); + finished.send(ordinal).unwrap(); + match outcome { + Outcome::Empty => Ok(None), + Outcome::Fail => Err(Error::UnexpectedError { + message: "original shard failure".into(), + source: Some(Box::new(std::io::Error::other("original native cause"))), + }), + Outcome::Panic => panic!("original shard panic"), + _ => Ok(message.map(|message| (message, ordinal))), + } + } + }, + ) + .await + }); + Running { + task, + started: receiver, + releases, + counters, + finished: finished_receiver, + } +} + +impl Running { + async fn next_started(&mut self) -> usize { + tokio::time::timeout(Duration::from_secs(10), self.started.recv()) + .await + .expect("admitted shard must start") + .expect("build still has work") + } + fn release(&mut self, ordinal: usize) { + self.releases[ordinal].take().unwrap().send(()).unwrap(); + } + async fn finish(self) -> Result> { + tokio::time::timeout(Duration::from_secs(10), self.task) + .await + .expect("build must finish") + .unwrap() + } +} + +#[tokio::test] +async fn bounded_parallel_jobs_return_messages_in_plan_order() { + let table = table(); + let mut running = start(&table, 2, vec![Outcome::File; 3]); + assert_eq!(running.next_started().await, 0); + assert_eq!(running.next_started().await, 1); + assert_eq!(running.counters.admitted.load(Ordering::SeqCst), 2); + running.release(1); + assert_eq!(running.next_started().await, 2); + running.release(2); + running.release(0); + let counters = running.counters.clone(); + let prepared = running.finish().await.unwrap(); + assert_eq!(counters.peak.load(Ordering::SeqCst), 2); + assert_eq!(counters.active.load(Ordering::SeqCst), 0); + assert_eq!( + prepared.iter().map(|(_, value)| *value).collect::>(), + vec![0, 1, 2] + ); + for (message, ordinal) in prepared { + assert_eq!( + message.new_index_files[0] + .global_index_meta + .as_ref() + .unwrap() + .row_range_start, + ordinal as i64 + ); + assert!(table + .file_io() + .exists(&format!( + "{}/index/private-{ordinal}.index", + table.location() + )) + .await + .unwrap()); + } +} + +#[tokio::test] +async fn one_worker_remains_sequential_and_empty_outputs_do_not_change_order() { + let table = table(); + let mut running = start( + &table, + 1, + vec![Outcome::File, Outcome::Empty, Outcome::File], + ); + for ordinal in 0..3 { + assert_eq!(running.next_started().await, ordinal); + assert_eq!(running.counters.active.load(Ordering::SeqCst), 1); + assert_eq!( + running.counters.admitted.load(Ordering::SeqCst), + ordinal + 1 + ); + running.release(ordinal); + } + let prepared = running.finish().await.unwrap(); + assert_eq!( + prepared + .into_iter() + .map(|(_, value)| value) + .collect::>(), + vec![0, 2] + ); +} + +#[tokio::test] +async fn a_large_limit_only_admits_available_shards() { + let table = table(); + let mut running = start(&table, usize::MAX, vec![Outcome::Empty; 2]); + assert_eq!(running.next_started().await, 0); + assert_eq!(running.next_started().await, 1); + assert_eq!(running.counters.admitted.load(Ordering::SeqCst), 2); + running.release(0); + running.release(1); + assert!(running.finish().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn failure_drains_inflight_jobs_before_deleting_only_private_outputs() { + for outcome in [Outcome::Fail, Outcome::Panic] { + for external in [false, true] { + let table = table(); + let protected = private_file(&table, 99, external).await; + let protected_path = protected.new_index_files[0] + .external_path + .clone() + .unwrap_or_else(|| format!("{}/index/private-99.index", table.location())); + let mode = if external { + Outcome::ExternalFile + } else { + Outcome::File + }; + let mut running = start(&table, 2, vec![mode, outcome, mode]); + assert_eq!(running.next_started().await, 0); + assert_eq!(running.next_started().await, 1); + let path = if external { + format!("{}-external/private-0.index", table.location()) + } else { + format!("{}/index/private-0.index", table.location()) + }; + running.release(1); + // The failed worker returns immediately after this signal; the + // coordinator then waits on job 0 before it can perform cleanup. + assert_eq!( + tokio::time::timeout(Duration::from_secs(10), running.finished.recv()) + .await + .unwrap(), + Some(1) + ); + assert_eq!(running.counters.active.load(Ordering::SeqCst), 1); + assert!(!running.task.is_finished()); + assert!(table.file_io().exists(&path).await.unwrap()); + running.release(0); + let counters = running.counters.clone(); + let error = running.finish().await.unwrap_err(); + match outcome { + Outcome::Fail => match error { + Error::UnexpectedError { message, source } => { + assert_eq!(message, "original shard failure"); + assert_eq!(source.unwrap().to_string(), "original native cause"); + } + _ => panic!("original failure was replaced"), + }, + Outcome::Panic => assert!(error.to_string().contains("original shard panic")), + _ => unreachable!(), + } + assert_eq!(counters.active.load(Ordering::SeqCst), 0); + assert_eq!(counters.admitted.load(Ordering::SeqCst), 2); + assert!(!table.file_io().exists(&path).await.unwrap()); + assert!(table.file_io().exists(&protected_path).await.unwrap()); + } + } +} + +#[tokio::test] +async fn empty_input_and_zero_limit_never_construct_a_worker() { + let table = table(); + let make = |_| async { + panic!("no worker may be constructed"); + #[allow(unreachable_code)] + Ok::<_, Error>(None::<(CommitMessage, ())>) + }; + assert!(prepare_shards(&table, Vec::<()>::new(), 1, make) + .await + .unwrap() + .is_empty()); + assert!(prepare_shards(&table, vec![()], 0, make) + .await + .unwrap_err() + .to_string() + .contains("parallelism")); +} diff --git a/crates/paimon/src/table/vindex_index_build_builder.rs b/crates/paimon/src/table/vindex_index_build_builder.rs index 10a4637ea..ddb4470e5 100644 --- a/crates/paimon/src/table/vindex_index_build_builder.rs +++ b/crates/paimon/src/table/vindex_index_build_builder.rs @@ -26,7 +26,7 @@ use planning::plan_vindex_shards; use validation::{checked_i32, find_index_field, validate_table_options, validate_vector_field}; use crate::spec::{CoreOptions, Predicate}; -use crate::table::{CommitMessage, RowRange, Table, TableCommit}; +use crate::table::{CommitMessage, RowRange, Table}; use crate::vindex::{is_vindex_index_type, VindexVectorIndexOptions}; use crate::{Error, Result}; use std::collections::HashMap; @@ -92,21 +92,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { let core_options = CoreOptions::new(&merged_options); validate_table_options(self.table, &core_options)?; let rows_per_shard = core_options.global_index_row_count_per_shard()?; - if let Some(value) = merged_options.get("global-index.build.parallelism") { - let parallelism = value - .trim() - .parse::() - .map_err(|_| Error::ConfigInvalid { - message: "Option 'global-index.build.parallelism' must be greater than 0." - .into(), - })?; - if parallelism == 0 { - return Err(Error::ConfigInvalid { - message: "Option 'global-index.build.parallelism' must be greater than 0." - .into(), - }); - } - } + let parallelism = core_options.global_index_build_parallelism()?; let index_field = find_index_field(self.table, index_column)?; validate_vector_field(index_field)?; @@ -175,42 +161,38 @@ impl<'a> VindexIndexBuildBuilder<'a> { ) .await?; - let commit = TableCommit::new( - self.table.clone(), - format!( - "global-index-{}-create-{}", - self.index_type, - uuid::Uuid::new_v4() - ), - ); - let shard_count = shards.len(); - let mut messages = Vec::with_capacity(shard_count); - let mut timings = Vec::with_capacity(shard_count); - for shard in shards { - let built = match self - .build_index_file( - &shard, - index_column, - dimension, - index_field.id(), - &vindex_options, - index_meta.clone(), - ) - .await - { - Ok(index_file) => index_file, - Err(error) => { - // Earlier outputs are private to this preparation; nothing - // has been returned to a caller or submitted for commit. - let _ = commit.abort(&messages).await; - return Err(error); + let prepared = super::global_index_build_common::preparation::prepare_shards( + self.table, + shards, + parallelism, + |shard| { + let vindex_options = &vindex_options; + let index_meta = index_meta.clone(); + async move { + let built = self + .build_index_file( + &shard, + index_column, + dimension, + index_field.id(), + vindex_options, + index_meta, + ) + .await?; + Ok(built.map(|built| { + let mut message = CommitMessage::new(shard.partition_bytes, 0, vec![]); + message.new_index_files = vec![built.meta]; + (message, built.timing) + })) } - }; - let Some(built) = built else { continue }; - let mut message = CommitMessage::new(shard.partition_bytes.clone(), 0, vec![]); - message.new_index_files = vec![built.meta]; + }, + ) + .await?; + let mut messages = Vec::with_capacity(prepared.len()); + let mut timings = Vec::with_capacity(prepared.len()); + for (message, timing) in prepared { messages.push(message); - if let Some(timing) = built.timing { + if let Some(timing) = timing { timings.push(timing); } } diff --git a/docs/src/sql.md b/docs/src/sql.md index bd8d4a8fc..e017bedb6 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -1256,6 +1256,13 @@ Invalid options are rejected before any data is read. NULL values are counted in a shard's row count but are not indexed. A later call indexes only rows that no `full-text` index on the column covers yet. +Full-text and vindex builds honor `global-index.build.parallelism` from table +or build options. The default is `1`; higher values bound the number of shards +built concurrently. Each shard owns its writer and may also use native worker +threads. Prepared commit messages retain shard plan order. On a build failure, +all started shards finish before private outputs are cleaned. Files in messages +already returned to a caller are retained, including after a failed commit. + The current global-index builders require a row-tracking data-evolution table with global indexes enabled. They do not support primary-key tables or tables with deletion vectors enabled: