From b3b263d35453f6407fc69ac5244cc99d405537df Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Tue, 1 Sep 2026 10:22:02 +0800 Subject: [PATCH 1/2] refactor: split global index implementations --- .../src/table/global_index_build_common.rs | 2 + .../table/global_index_build_common/vector.rs | 258 ++ .../paimon/src/table/global_index_scanner.rs | 3864 +---------------- .../global_index_scanner/deletion_vectors.rs | 122 + .../src/table/global_index_scanner/entry.rs | 212 + .../table/global_index_scanner/evaluator.rs | 609 +++ .../table/global_index_scanner/predicates.rs | 115 + .../table/global_index_scanner/query_plan.rs | 88 + .../src/table/global_index_scanner/reader.rs | 380 ++ .../table/global_index_scanner/row_ranges.rs | 296 ++ .../src/table/global_index_scanner/tests.rs | 2136 +++++++++ .../src/table/lumina_index_build_builder.rs | 1938 +-------- .../lumina_index_build_builder/extraction.rs | 199 + .../lumina_index_build_builder/planning.rs | 49 + .../table/lumina_index_build_builder/tests.rs | 1257 ++++++ .../lumina_index_build_builder/validation.rs | 135 + .../lumina_index_build_builder/writer.rs | 221 + .../sorted_global_index_build_builder.rs | 3348 +------------- .../extraction.rs | 343 ++ .../planning.rs | 245 ++ .../tests.rs | 2332 ++++++++++ .../validation.rs | 179 + .../writer.rs | 365 ++ .../src/table/vindex_index_build_builder.rs | 1876 +------- .../vindex_index_build_builder/extraction.rs | 190 + .../vindex_index_build_builder/planning.rs | 47 + .../table/vindex_index_build_builder/tests.rs | 863 ++++ .../vindex_index_build_builder/timing.rs | 102 + .../vindex_index_build_builder/validation.rs | 155 + .../vindex_index_build_builder/writer.rs | 442 ++ 30 files changed, 11462 insertions(+), 10906 deletions(-) create mode 100644 crates/paimon/src/table/global_index_build_common/vector.rs create mode 100644 crates/paimon/src/table/global_index_scanner/deletion_vectors.rs create mode 100644 crates/paimon/src/table/global_index_scanner/entry.rs create mode 100644 crates/paimon/src/table/global_index_scanner/evaluator.rs create mode 100644 crates/paimon/src/table/global_index_scanner/predicates.rs create mode 100644 crates/paimon/src/table/global_index_scanner/query_plan.rs create mode 100644 crates/paimon/src/table/global_index_scanner/reader.rs create mode 100644 crates/paimon/src/table/global_index_scanner/row_ranges.rs create mode 100644 crates/paimon/src/table/global_index_scanner/tests.rs create mode 100644 crates/paimon/src/table/lumina_index_build_builder/extraction.rs create mode 100644 crates/paimon/src/table/lumina_index_build_builder/planning.rs create mode 100644 crates/paimon/src/table/lumina_index_build_builder/tests.rs create mode 100644 crates/paimon/src/table/lumina_index_build_builder/validation.rs create mode 100644 crates/paimon/src/table/lumina_index_build_builder/writer.rs create mode 100644 crates/paimon/src/table/sorted_global_index_build_builder/extraction.rs create mode 100644 crates/paimon/src/table/sorted_global_index_build_builder/planning.rs create mode 100644 crates/paimon/src/table/sorted_global_index_build_builder/tests.rs create mode 100644 crates/paimon/src/table/sorted_global_index_build_builder/validation.rs create mode 100644 crates/paimon/src/table/sorted_global_index_build_builder/writer.rs create mode 100644 crates/paimon/src/table/vindex_index_build_builder/extraction.rs create mode 100644 crates/paimon/src/table/vindex_index_build_builder/planning.rs create mode 100644 crates/paimon/src/table/vindex_index_build_builder/tests.rs create mode 100644 crates/paimon/src/table/vindex_index_build_builder/timing.rs create mode 100644 crates/paimon/src/table/vindex_index_build_builder/validation.rs create mode 100644 crates/paimon/src/table/vindex_index_build_builder/writer.rs diff --git a/crates/paimon/src/table/global_index_build_common.rs b/crates/paimon/src/table/global_index_build_common.rs index 355a56e47..3b255a259 100644 --- a/crates/paimon/src/table/global_index_build_common.rs +++ b/crates/paimon/src/table/global_index_build_common.rs @@ -28,6 +28,8 @@ use crate::spec::{FileKind, IndexManifest}; use crate::table::{merge_row_ranges, RowRange, Table}; use crate::{Error, Result}; +pub(crate) mod vector; + /// Java `sameExtraFieldIds`: null/empty are equal; otherwise exact ordered equality. pub(crate) fn same_extra_field_ids(a: Option<&[i32]>, b: Option<&[i32]>) -> bool { let a = a.unwrap_or(&[]); diff --git a/crates/paimon/src/table/global_index_build_common/vector.rs b/crates/paimon/src/table/global_index_build_common/vector.rs new file mode 100644 index 000000000..72a37d68a --- /dev/null +++ b/crates/paimon/src/table/global_index_build_common/vector.rs @@ -0,0 +1,258 @@ +// 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. + +//! Shared shard planning for contiguous vector global indexes. + +use crate::spec::{ + bucket_dir_name, BinaryRow, CoreOptions, DataField, DataFileMeta, DataType, FileKind, + ManifestEntry, PartitionComputer, +}; +use crate::table::source::exclude_row_ranges; +use crate::table::{RowRange, Table}; +use crate::{Error, Result}; +use std::collections::HashMap; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct VectorIndexShard { + pub(crate) partition: BinaryRow, + pub(crate) partition_bytes: Vec, + pub(crate) files: Vec, + pub(crate) row_range_start: i64, + pub(crate) row_range_end: i64, + pub(crate) snapshot_id: i64, + pub(crate) source_bucket: i32, + pub(crate) total_buckets: i32, + pub(crate) bucket_path: String, +} + +pub(crate) fn find_index_field<'a>(table: &'a Table, column: &str) -> Result<&'a DataField> { + table + .schema() + .fields() + .iter() + .find(|field| field.name() == column) + .ok_or_else(|| Error::ColumnNotExist { + full_name: table.identifier().full_name(), + column: column.to_string(), + }) +} + +pub(crate) fn validate_vector_field(field: &DataField, index_name: &str) -> Result<()> { + let is_array_float = matches!( + field.data_type(), + DataType::Array(array) if matches!(array.element_type(), DataType::Float(_)) + ); + let is_vector_float = matches!( + field.data_type(), + DataType::Vector(vector) if matches!(vector.element_type(), DataType::Float(_)) + ); + if !is_array_float && !is_vector_float { + return Err(Error::DataInvalid { + message: format!( + "{index_name} index requires ARRAY or VECTOR column, got {:?} for column '{}'", + field.data_type(), + field.name() + ), + source: None, + }); + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn plan_vector_index_shards( + table_location: &str, + partition_keys: &[String], + schema_fields: &[DataField], + core_options: &CoreOptions, + snapshot_id: i64, + entries: Vec, + rows_per_shard: i64, + indexed: &[RowRange], + index_name: &str, +) -> Result> { + if rows_per_shard <= 0 { + return Err(Error::DataInvalid { + message: format!( + "Option 'global-index.row-count-per-shard' must be greater than 0, got: {rows_per_shard}" + ), + source: None, + }); + } + + let mut by_partition_bucket: HashMap<(Vec, i32, i32), Vec> = HashMap::new(); + for entry in entries { + if *entry.kind() != FileKind::Add { + continue; + } + if entry.file().first_row_id.is_none() { + return Err(missing_row_id_error(entry.file(), index_name)); + } + let (partition, bucket, total_buckets, file) = entry.into_parts(); + by_partition_bucket + .entry((partition, bucket, total_buckets)) + .or_default() + .push(file); + } + + let mut result = Vec::new(); + for ((partition_bytes, source_bucket, total_buckets), files) in by_partition_bucket { + let partition = if partition_keys.is_empty() { + BinaryRow::new(0) + } else { + BinaryRow::from_serialized_bytes(&partition_bytes)? + }; + let bucket_path = bucket_path( + table_location, + partition_keys, + schema_fields, + core_options, + &partition, + source_bucket, + )?; + let mut files_by_shard: HashMap> = HashMap::new(); + for file in files { + let (file_start, file_end) = file + .row_id_range() + .ok_or_else(|| missing_row_id_error(&file, index_name))?; + let start_shard = file_start / rows_per_shard; + let end_shard = file_end / rows_per_shard; + for shard_id in start_shard..=end_shard { + files_by_shard + .entry(shard_id * rows_per_shard) + .or_default() + .push(file.clone()); + } + } + + let mut shard_starts = files_by_shard.keys().copied().collect::>(); + shard_starts.sort_unstable(); + for shard_start in shard_starts { + let shard_end = shard_start + rows_per_shard - 1; + let mut shard_files = files_by_shard.remove(&shard_start).unwrap_or_default(); + shard_files.sort_by_key(|file| file.first_row_id); + for group in group_contiguous_files(shard_files, index_name)? { + let group_start = group + .first() + .and_then(|file| file.first_row_id) + .expect("planned groups are non-empty and row-id assigned"); + let group_end = group + .iter() + .map(|file| file.row_id_range().unwrap().1) + .max() + .unwrap(); + let coverage_start = group_start.max(shard_start); + let coverage_end = group_end.min(shard_end); + let build_segments = + exclude_row_ranges(&[RowRange::new(coverage_start, coverage_end)], indexed); + for segment in build_segments { + result.push(VectorIndexShard { + partition: partition.clone(), + partition_bytes: partition_bytes.clone(), + files: group.clone(), + row_range_start: segment.from(), + row_range_end: segment.to(), + snapshot_id, + source_bucket, + total_buckets, + bucket_path: bucket_path.clone(), + }); + } + } + } + } + result.sort_by(|left, right| { + left.partition + .to_serialized_bytes() + .cmp(&right.partition.to_serialized_bytes()) + .then(left.source_bucket.cmp(&right.source_bucket)) + .then(left.row_range_start.cmp(&right.row_range_start)) + }); + Ok(result) +} + +fn missing_row_id_error(file: &DataFileMeta, index_name: &str) -> Error { + Error::DataInvalid { + message: format!( + "Data file '{}' is missing first_row_id; cannot build a complete {index_name} index", + file.file_name + ), + source: None, + } +} + +fn group_contiguous_files( + mut files: Vec, + index_name: &str, +) -> Result>> { + if files.is_empty() { + return Ok(Vec::new()); + } + files.sort_by_key(|file| file.first_row_id); + let mut groups = Vec::new(); + let mut current = Vec::new(); + let mut current_end = None; + for file in files { + let (file_start, file_end) = file + .row_id_range() + .ok_or_else(|| missing_row_id_error(&file, index_name))?; + match current_end { + None => { + current.push(file); + current_end = Some(file_end); + } + Some(end) if file_start <= end + 1 => { + current.push(file); + current_end = Some(end.max(file_end)); + } + Some(_) => { + groups.push(std::mem::take(&mut current)); + current.push(file); + current_end = Some(file_end); + } + } + } + if !current.is_empty() { + groups.push(current); + } + Ok(groups) +} + +fn bucket_path( + table_location: &str, + partition_keys: &[String], + schema_fields: &[DataField], + core_options: &CoreOptions, + partition: &BinaryRow, + bucket: i32, +) -> Result { + let base = table_location.trim_end_matches('/'); + if partition_keys.is_empty() { + return Ok(format!("{base}/{}", bucket_dir_name(bucket))); + } + let computer = PartitionComputer::new( + partition_keys, + schema_fields, + core_options.partition_default_name(), + core_options.legacy_partition_name(), + )?; + Ok(format!( + "{base}/{}{}", + computer.generate_partition_path(partition)?, + bucket_dir_name(bucket) + )) +} diff --git a/crates/paimon/src/table/global_index_scanner.rs b/crates/paimon/src/table/global_index_scanner.rs index 7d137ce9f..746fb8a00 100644 --- a/crates/paimon/src/table/global_index_scanner.rs +++ b/crates/paimon/src/table/global_index_scanner.rs @@ -20,30 +20,33 @@ //! //! Reference: [org.apache.paimon.index.GlobalIndexScanner](https://github.com/apache/paimon/blob/master/paimon-core/src/main/java/org/apache/paimon/index/GlobalIndexScanner.java) -use super::bitmap_global_index_format::{ - is_bitmap_floating_residual_sensitive_op, make_bitmap_key_comparator, serialize_bitmap_datum, +mod deletion_vectors; +mod entry; +mod evaluator; +mod predicates; +mod query_plan; +mod reader; +mod row_ranges; + +pub(crate) use deletion_vectors::deleted_row_ranges_for_data_evolution_dvs; +use entry::{validate_fm_file_sets, GlobalIndexEntry, GlobalIndexEntryMeta, GlobalIndexFileKind}; +use row_ranges::unindexed_ranges_for_indexed_coverage; +pub(crate) use row_ranges::{ + search_limit_with_deleted_rows, unindexed_ranges_for_global_index_entries, RowRangeIndex, }; -use super::bitmap_global_index_reader::BitmapGlobalIndexReader; + use super::global_index_types::{ normalize_queryable_global_index_type, BITMAP_GLOBAL_INDEX_TYPE, BTREE_GLOBAL_INDEX_TYPE, FM_GLOBAL_INDEX_TYPE, MULTIVALUE_GLOBAL_INDEX_TYPE, }; -use crate::btree::query::{extract_between, BetweenInfo, IndexQuery}; -use crate::btree::{make_key_comparator, serialize_datum, BTreeIndexMeta, BTreeIndexReader}; -use crate::deletion_vector::DeletionVectorFactory; -use crate::fm_index::{manifest_row_range, FMGlobalIndexReader, FMReadContext, FMReadOptions}; +use crate::btree::{BTreeIndexMeta, BTreeIndexReader}; +use crate::fm_index::{manifest_row_range, FMReadContext, FMReadOptions}; use crate::io::FileIO; -use crate::spec::{ - DataField, DataType, Datum, FileKind, GlobalIndexSearchMode, IndexFileMeta, IndexManifestEntry, - Predicate, PredicateOperator, -}; -use crate::table::{DeletionFile, RowRange, Table}; +use crate::spec::{DataField, FileKind, GlobalIndexSearchMode, IndexManifestEntry, Predicate}; +use crate::table::RowRange; use crate::{Error, Result}; -use futures::{StreamExt, TryStreamExt}; -use roaring::RoaringTreemap; use std::cmp::Ordering; -use std::collections::{HashMap, HashSet}; -use std::future::Future; +use std::collections::HashMap; use std::sync::{Arc, Mutex}; use tokio::sync::Semaphore; @@ -52,34 +55,9 @@ use std::sync::atomic::{AtomicUsize as TestAtomicUsize, Ordering as TestOrdering type BoxedCmp = Box Ordering + Send + Sync>; -type EvaluateFuture<'a> = std::pin::Pin< - Box>> + Send + 'a>, ->; - -type PredicateTuple<'a> = (PredicateOperator, &'a [Datum], &'a DataType); - const DELETION_VECTORS_INDEX_TYPE: &str = "DELETION_VECTORS"; const INDEX_DIR: &str = "index"; -async fn try_fold_bounded( - futures: impl IntoIterator, - max_concurrency: usize, - mut accumulator: Acc, - mut fold: Fold, -) -> Result -where - Fut: Future>, - Fold: FnMut(&mut Acc, T), -{ - debug_assert!(max_concurrency > 0); - let stream = futures::stream::iter(futures).buffer_unordered(max_concurrency); - futures::pin_mut!(stream); - while let Some(value) = stream.try_next().await? { - fold(&mut accumulator, value); - } - Ok(accumulator) -} - #[cfg(test)] #[derive(Default)] struct QueryIoProbe { @@ -114,12 +92,6 @@ impl Drop for QueryIoProbeGuard<'_> { } } -struct GlobalIndexScanResult { - row_ranges: Vec, - evaluated_field_ids: HashSet, - indexed_coverage: Vec, -} - /// Evaluates global index predicates and returns matching row ranges. /// /// The scanner filters index manifest entries for global index files, @@ -149,283 +121,6 @@ pub(crate) struct GlobalIndexScanner { query_io_probe: Option>, } -/// A resolved global index entry with parsed metadata. -struct GlobalIndexEntry { - file_name: String, - index_type: GlobalIndexFileKind, - file_size: i64, - row_range_start: i64, - row_range_end: i64, - meta: GlobalIndexEntryMeta, -} - -fn sorted_entry_meta(entry: &GlobalIndexEntry) -> &BTreeIndexMeta { - match &entry.meta { - GlobalIndexEntryMeta::Sorted(meta) => meta, - GlobalIndexEntryMeta::FM { .. } => unreachable!("FM entries do not have sorted metadata"), - } -} - -enum GlobalIndexEntryMeta { - Sorted(BTreeIndexMeta), - FM { - bytes: Vec, - first_row_id: u64, - row_count: u64, - }, -} - -struct FMFileRowRange<'a> { - file_name: &'a str, - first_row_id: u64, - row_count: u64, -} - -fn validate_fm_file_sets(entries_by_field: &HashMap>) -> Result<()> { - for (field_id, entries) in entries_by_field { - let mut groups: HashMap<(i64, i64), Vec>> = HashMap::new(); - for entry in entries { - let GlobalIndexEntryMeta::FM { - first_row_id, - row_count, - .. - } = &entry.meta - else { - continue; - }; - groups - .entry((entry.row_range_start, entry.row_range_end)) - .or_default() - .push(FMFileRowRange { - file_name: &entry.file_name, - first_row_id: *first_row_id, - row_count: *row_count, - }); - } - - for ((range_start, range_end), mut files) in groups { - let expected_row_count = range_end - .checked_sub(range_start) - .and_then(|count| count.checked_add(1)) - .and_then(|count| u64::try_from(count).ok()) - .ok_or_else(|| Error::DataInvalid { - message: format!( - "Invalid FM global index source row range [{range_start}, {range_end}] for field {field_id}" - ), - source: None, - })?; - files.sort_unstable_by_key(|file| file.first_row_id); - let mut expected_first_row_id = 0u64; - for file in files { - let FMFileRowRange { - file_name, - first_row_id, - row_count, - } = file; - if row_count == 0 || first_row_id != expected_first_row_id { - return Err(Error::DataInvalid { - message: format!( - "FM global index files do not exactly cover source row range [{range_start}, {range_end}] for field {field_id}: expected relative row {expected_first_row_id}, file '{file_name}' starts at {first_row_id}" - ), - source: None, - }); - } - expected_first_row_id = - first_row_id - .checked_add(row_count) - .ok_or_else(|| Error::DataInvalid { - message: format!( - "FM global index row range overflows for file '{file_name}'" - ), - source: None, - })?; - if expected_first_row_id > expected_row_count { - return Err(Error::DataInvalid { - message: format!( - "FM global index file '{file_name}' extends beyond source row range [{range_start}, {range_end}]" - ), - source: None, - }); - } - } - if expected_first_row_id != expected_row_count { - return Err(Error::DataInvalid { - message: format!( - "FM global index files cover {expected_first_row_id} rows, expected {expected_row_count} for source row range [{range_start}, {range_end}] and field {field_id}" - ), - source: None, - }); - } - } - } - Ok(()) -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum GlobalIndexFileKind { - BTree, - Bitmap, - Multivalue, - FM, -} - -fn is_floating_point(data_type: &DataType) -> bool { - matches!(data_type, DataType::Float(_) | DataType::Double(_)) -} - -fn bitmap_meta_may_match( - meta: &BTreeIndexMeta, - op: PredicateOperator, - data_type: &DataType, - serialized_literals: &[Vec], - cmp: &dyn Fn(&[u8], &[u8]) -> Ordering, -) -> bool { - if is_floating_point(data_type) && is_bitmap_floating_residual_sensitive_op(op) { - !meta.only_nulls() - } else { - meta.may_match(op, serialized_literals, cmp) - } -} - -fn bitmap_meta_may_match_between( - meta: &BTreeIndexMeta, - data_type: &DataType, - from_key: &[u8], - to_key: &[u8], - cmp: &dyn Fn(&[u8], &[u8]) -> Ordering, -) -> bool { - if is_floating_point(data_type) - && is_bitmap_floating_residual_sensitive_op(PredicateOperator::Between) - { - !meta.only_nulls() - } else { - meta.may_match_between(from_key, to_key, cmp) - } -} - -fn multivalue_meta_may_match( - meta: &BTreeIndexMeta, - op: PredicateOperator, - serialized_literals: &[Vec], - cmp: &dyn Fn(&[u8], &[u8]) -> Ordering, -) -> bool { - match op { - PredicateOperator::ArrayContains => { - meta.may_match(PredicateOperator::Eq, serialized_literals, cmp) - } - PredicateOperator::ArraysOverlap => { - meta.may_match(PredicateOperator::In, serialized_literals, cmp) - } - PredicateOperator::ArrayContainsAll => serialized_literals.iter().all(|literal| { - meta.may_match(PredicateOperator::Eq, std::slice::from_ref(literal), cmp) - }), - _ => false, - } -} - -impl GlobalIndexFileKind { - fn name(self) -> &'static str { - match self { - Self::BTree => "BTree", - Self::Bitmap => "bitmap", - Self::Multivalue => "multivalue", - Self::FM => "FM", - } - } -} - -enum OpenedGlobalIndexReader { - BTree(BTreeIndexReader), - Bitmap(BitmapGlobalIndexReader), - FM(FMGlobalIndexReader), -} - -#[derive(Clone, Copy, Default)] -struct FallbackScanPlan { - selected_btree: usize, - selected_bitmap: usize, - allow_btree: bool, - allow_bitmap: bool, -} - -struct EntryQueryPlan { - entry_idx: usize, - between_matches: bool, - between_evaluated: bool, - matching_predicates: Vec, -} - -struct EntryQueryResult { - bitmap: Option, - declined: bool, -} - -impl FallbackScanPlan { - fn allowed(self, kind: GlobalIndexFileKind) -> bool { - match kind { - GlobalIndexFileKind::BTree => self.allow_btree, - GlobalIndexFileKind::Bitmap | GlobalIndexFileKind::Multivalue => self.allow_bitmap, - GlobalIndexFileKind::FM => true, - } - } -} - -impl OpenedGlobalIndexReader { - async fn query( - &self, - op: PredicateOperator, - literals: &[Datum], - data_type: &DataType, - ) -> std::io::Result> { - match self { - Self::BTree(reader) => reader.query(op, literals, data_type).await.map(Some), - Self::Bitmap(reader) => reader.query(op, literals, data_type).await.map(Some), - Self::FM(reader) => match op { - PredicateOperator::Contains => { - let literal = literals.first().ok_or_else(|| { - std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "FM contains requires one literal", - ) - })?; - reader - .contains(&serialize_bitmap_datum(literal, data_type)) - .await - } - PredicateOperator::IsNull => reader.is_null().await.map(Some), - PredicateOperator::IsNotNull => reader.is_not_null().await.map(Some), - _ => Ok(None), - }, - } - } - - async fn range_query( - &self, - from: &[u8], - to: &[u8], - data_type: &DataType, - from_inclusive: bool, - to_inclusive: bool, - ) -> std::io::Result { - match self { - Self::BTree(reader) => { - reader - .range_query(from, to, from_inclusive, to_inclusive) - .await - } - Self::Bitmap(reader) => { - reader - .range_query(from, to, data_type, from_inclusive, to_inclusive) - .await - } - Self::FM(_) => Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "FM index does not support ordered range queries", - )), - } - } -} - impl GlobalIndexScanner { /// Create a scanner from index manifest entries. /// Returns `Ok(None)` if there are no global index entries. @@ -620,1359 +315,45 @@ impl GlobalIndexScanner { query_io_probe: None, })) } - - /// Evaluate a predicate against the global indexes and return matching row ranges. - /// Returns `None` if the predicate cannot be evaluated by the global index. - fn evaluate<'a>(&'a self, predicate: &'a Predicate) -> EvaluateFuture<'a> { - Box::pin(async move { - match predicate { - Predicate::Leaf { - column, - op, - literals, - data_type, - .. - } => { - if !is_sorted_global_index_supported_op(*op) { - return Ok(None); - } - let field_id = self.find_field_id_by_name(column)?; - let field_id = match field_id { - Some(id) => id, - None => return Ok(None), - }; - let entries = match self.entries_for_field(field_id) { - Some(e) => e, - None => return Ok(None), - }; - if !entries_support_predicate(entries, *op, literals) { - return Ok(None); - } - let predicates = [(*op, literals.as_slice(), data_type)]; - let selected_entries = select_entries_for_predicates(entries, &predicates); - self.evaluate_leaf(&selected_entries, &predicates) - .await - .map(|result| { - result.map(|(row_ranges, indexed_coverage)| GlobalIndexScanResult { - row_ranges, - evaluated_field_ids: HashSet::from([field_id]), - indexed_coverage, - }) - }) - } - Predicate::And(children) => { - // Group leaf predicates by field_id to reuse readers - let mut leaf_groups: std::collections::HashMap>> = - std::collections::HashMap::new(); - let mut non_leaf_children = Vec::new(); - - for child in children { - if let Predicate::Leaf { - column, - op, - literals, - data_type, - .. - } = child - { - if is_sorted_global_index_supported_op(*op) { - if let Some(field_id) = self.find_field_id_by_name(column)? { - if self.entries_for_field(field_id).is_some_and(|entries| { - entries_support_predicate(entries, *op, literals) - }) { - leaf_groups.entry(field_id).or_default().push(( - *op, - literals.as_slice(), - data_type, - )); - continue; - } - } - } - } - non_leaf_children.push(child); - } - - // Evaluate independent fields concurrently while keeping predicates for the - // same field together so each index file is opened only once. - let mut leaf_futures = Vec::with_capacity(leaf_groups.len()); - for (field_id, predicates) in &leaf_groups { - if let Some(entries) = self.entries_for_field(*field_id) { - let field_id = *field_id; - let mut selected_predicates = predicates.clone(); - let mut selected_entries = - select_entries_for_predicates(entries, &selected_predicates); - if selected_entries.is_empty() { - for predicate in predicates { - let candidate_predicates = vec![*predicate]; - let candidate_entries = select_entries_for_predicates( - entries, - &candidate_predicates, - ); - if !candidate_entries.is_empty() { - selected_predicates = candidate_predicates; - selected_entries = candidate_entries; - break; - } - } - } - leaf_futures.push(async move { - let result = self - .evaluate_leaf(&selected_entries, &selected_predicates) - .await?; - Ok((field_id, result)) - }); - } - } - let leaf_group_count = leaf_futures.len(); - let (mut row_ranges, mut indexed_coverage, mut evaluated_field_ids) = - try_fold_bounded( - leaf_futures, - leaf_group_count.max(1), - (None::>, None::>, HashSet::new()), - |(row_ranges, indexed_coverage, evaluated_field_ids), - (field_id, result)| { - if let Some((ranges, coverage)) = result { - *row_ranges = Some(match row_ranges.take() { - None => ranges, - Some(existing) => { - intersect_sorted_ranges(&existing, &ranges) - } - }); - *indexed_coverage = Some(match indexed_coverage.take() { - None => coverage, - Some(existing) => { - intersect_sorted_ranges(&existing, &coverage) - } - }); - evaluated_field_ids.insert(field_id); - } - }, - ) - .await?; - - // Evaluate non-leaf children recursively - for child in non_leaf_children { - if let Some(child_result) = self.evaluate(child).await? { - row_ranges = Some(match row_ranges { - None => child_result.row_ranges, - Some(existing) => { - intersect_sorted_ranges(&existing, &child_result.row_ranges) - } - }); - evaluated_field_ids.extend(child_result.evaluated_field_ids); - indexed_coverage = Some(match indexed_coverage { - None => child_result.indexed_coverage, - Some(existing) => intersect_sorted_ranges( - &existing, - &child_result.indexed_coverage, - ), - }); - } - } - - Ok(row_ranges.map(|row_ranges| GlobalIndexScanResult { - row_ranges, - evaluated_field_ids, - indexed_coverage: indexed_coverage.unwrap_or_default(), - })) - } - Predicate::Or(children) => { - let mut all_ranges: Vec = Vec::new(); - let mut evaluated_field_ids = HashSet::new(); - let mut indexed_coverage: Option> = None; - for child in children { - match self.evaluate(child).await? { - Some(child_result) => { - all_ranges.extend(child_result.row_ranges); - evaluated_field_ids.extend(child_result.evaluated_field_ids); - indexed_coverage = Some(match indexed_coverage { - None => child_result.indexed_coverage, - Some(existing) => intersect_sorted_ranges( - &existing, - &child_result.indexed_coverage, - ), - }); - } - None => return Ok(None), - } - } - let row_ranges = if all_ranges.is_empty() { - Vec::new() - } else { - super::merge_row_ranges(all_ranges) - }; - Ok(Some(GlobalIndexScanResult { - row_ranges, - evaluated_field_ids, - indexed_coverage: indexed_coverage.unwrap_or_default(), - })) - } - _ => Ok(None), - } - }) - } - - /// Evaluate multiple predicates against the same set of index entries. - /// Opens each file once and evaluates all predicates, intersecting results. - /// Detects between patterns (GtEq/Gt + LtEq/Lt) and merges them into a single range query. - async fn evaluate_leaf( - &self, - entries: &[&GlobalIndexEntry], - predicates: &[(PredicateOperator, &[Datum], &DataType)], - ) -> Result, Vec)>> { - let normalized_predicates = predicates - .iter() - .map(|(op, literals, data_type)| { - let key_type = if is_multivalue_predicate(*op) { - let DataType::Array(array) = data_type else { - return Err(Error::DataInvalid { - message: format!( - "Array global-index predicate {op} requires an ARRAY field type" - ), - source: None, - }); - }; - array.element_type() - } else { - *data_type - }; - Ok((*op, *literals, key_type)) - }) - .collect::>>()?; - let predicates = normalized_predicates.as_slice(); - - // Try to detect between pattern and split into (between, remaining) - let (between, remaining) = extract_between(predicates); - - let effective_predicates = if between.is_some() { - &remaining - } else { - predicates - }; - - // Pre-compute comparators and serialized keys for file-level pruning per predicate - let pruning_info: Vec<_> = effective_predicates - .iter() - .map(|(op, literals, data_type)| { - let btree_cmp = make_key_comparator(data_type); - let btree_serialized = literals - .iter() - .map(|l| serialize_datum(l, data_type)) - .collect::>(); - let bitmap_cmp = make_bitmap_key_comparator(data_type); - let bitmap_serialized = literals - .iter() - .map(|l| serialize_bitmap_datum(l, data_type)) - .collect::>(); - ( - *op, - *data_type, - btree_cmp, - btree_serialized, - bitmap_cmp, - bitmap_serialized, - ) - }) - .collect(); - - let predicate_matches: Vec> = pruning_info - .iter() - .map( - |(op, data_type, btree_cmp, btree_serialized, bitmap_cmp, bitmap_serialized)| { - entries - .iter() - .map(|entry| match entry.index_type { - GlobalIndexFileKind::BTree => { - sorted_entry_meta(entry).may_match(*op, btree_serialized, btree_cmp) - } - GlobalIndexFileKind::Bitmap => bitmap_meta_may_match( - sorted_entry_meta(entry), - *op, - data_type, - bitmap_serialized, - bitmap_cmp.as_ref(), - ), - GlobalIndexFileKind::Multivalue => multivalue_meta_may_match( - sorted_entry_meta(entry), - *op, - bitmap_serialized, - bitmap_cmp.as_ref(), - ), - GlobalIndexFileKind::FM => true, - }) - .collect() - }, - ) - .collect(); - let predicate_fallback_plans: Vec> = effective_predicates - .iter() - .enumerate() - .map(|(i, (op, _, _))| { - requires_fallback_scan(*op) - .then(|| self.fallback_scan_plan(entries, &predicate_matches[i])) - }) - .collect(); - - let between_matches_by_entry: Vec = - match between.as_ref() { - Some(b) => { - let btree_cmp = make_key_comparator(b.data_type); - let btree_from = serialize_datum(b.from, b.data_type); - let btree_to = serialize_datum(b.to, b.data_type); - let bitmap_cmp = make_bitmap_key_comparator(b.data_type); - let bitmap_from = serialize_bitmap_datum(b.from, b.data_type); - let bitmap_to = serialize_bitmap_datum(b.to, b.data_type); - entries - .iter() - .map(|entry| match entry.index_type { - GlobalIndexFileKind::BTree => sorted_entry_meta(entry) - .may_match_between(&btree_from, &btree_to, &btree_cmp), - GlobalIndexFileKind::Bitmap => bitmap_meta_may_match_between( - sorted_entry_meta(entry), - b.data_type, - &bitmap_from, - &bitmap_to, - bitmap_cmp.as_ref(), - ), - GlobalIndexFileKind::Multivalue | GlobalIndexFileKind::FM => false, - }) - .collect() - } - None => Vec::new(), - }; - let between_fallback_plan = between - .as_ref() - .map(|_| self.fallback_scan_plan(entries, &between_matches_by_entry)); - - let mut query_plans = Vec::with_capacity(entries.len()); - for (entry_idx, entry) in entries.iter().enumerate() { - // Also check if between range may match - let between_matches = between - .as_ref() - .is_some_and(|_| between_matches_by_entry[entry_idx]); - let between_evaluated_for_entry = between_fallback_plan.is_some_and(|plan| { - fallback_plan_evaluates_entry(plan, entry.index_type, between_matches) - }); - - // When a Between conjunct exists but the file does not overlap its - // range, the whole AND cannot match — drop the file regardless of - // how the remaining predicates evaluate. Without this guard, a file - // outside the Between range but matched by some remaining predicate - // (e.g. `BETWEEN 10 AND 20 AND id >= 0` on a file [30, 40]) would - // be retained because `file_result` is initialized from the - // remaining bitmap, silently dropping the Between conjunct. - if between_evaluated_for_entry && !between_matches { - continue; - } - - let mut file_evaluated = between_evaluated_for_entry; - let mut file_cannot_match = false; - let mut file_has_unsupported_match = - between_matches && !between_evaluated_for_entry && between_fallback_plan.is_some(); - let matching_predicates: Vec = (0..effective_predicates.len()) - .filter(|&i| { - let predicate_matches_entry = predicate_matches[i][entry_idx]; - let predicate_evaluated_for_entry = - predicate_fallback_plans[i].is_none_or(|plan| { - fallback_plan_evaluates_entry( - plan, - entry.index_type, - predicate_matches_entry, - ) - }); - if !predicate_evaluated_for_entry { - file_has_unsupported_match |= predicate_matches_entry; - return false; - } - file_evaluated = true; - if !predicate_matches_entry { - file_cannot_match = true; - return false; - } - true - }) - .collect(); - if file_cannot_match { - continue; - } - if !file_evaluated { - if file_has_unsupported_match { - return Ok(None); - } - continue; - } - - query_plans.push(EntryQueryPlan { - entry_idx, - between_matches, - between_evaluated: between_evaluated_for_entry, - matching_predicates, - }); - } - - // Complete all pruning and fallback decisions before starting shard I/O. - // A later unsupported shard must fall back to the normal scan without an - // earlier shard racing it with an I/O or query error. - let data_type = between - .as_ref() - .map(|b| b.data_type) - .or_else(|| effective_predicates.first().map(|p| p.2)) - .unwrap_or(predicates[0].2); - let between = between.as_ref(); - let futures = - query_plans.into_iter().map(|plan| async move { - let entry = &entries[plan.entry_idx]; - let _permit = self.query_semaphore.acquire().await.map_err(|error| { - Error::UnexpectedError { - message: "global-index query concurrency budget was closed".to_string(), - source: Some(Box::new(error)), - } - })?; - #[cfg(test)] - let _query_io_probe_guard = match &self.query_io_probe { - Some(probe) => Some(probe.enter().await), - None => None, - }; - let result = self - .query_entry(entry, data_type, between, &plan, effective_predicates) - .await?; - Ok((entry.row_range_start, result)) - }); - let (all_row_ids, declined) = try_fold_bounded( - futures, - self.global_index_thread_num, - (RoaringTreemap::new(), false), - |(all_row_ids, declined), (row_range_start, file_result)| { - if file_result.declined { - *declined = true; - return; - } - if let Some(bitmap) = file_result.bitmap { - for row_id in bitmap.iter() { - all_row_ids.insert(row_id + row_range_start as u64); - } - } - }, - ) - .await?; - - if declined { - return Ok(None); - } - - let coverage = super::merge_row_ranges( - entries - .iter() - .map(|entry| RowRange::new(entry.row_range_start, entry.row_range_end)) - .collect(), - ); - Ok(Some((bitmap_to_ranges(&all_row_ids), coverage))) - } - - async fn query_entry( - &self, - entry: &GlobalIndexEntry, - data_type: &DataType, - between: Option<&BetweenInfo<'_>>, - plan: &EntryQueryPlan, - effective_predicates: &[(PredicateOperator, &[Datum], &DataType)], - ) -> Result { - let mut reader = if (plan.between_matches && plan.between_evaluated) - || !plan.matching_predicates.is_empty() - { - Some(self.open_reader_for_entry(entry, data_type).await?) - } else { - None - }; - let mut file_result = None; - - if plan.between_matches && plan.between_evaluated { - let between = between.expect("evaluated between query is present"); - let serialize_key = match entry.index_type { - GlobalIndexFileKind::BTree => serialize_datum, - GlobalIndexFileKind::Bitmap | GlobalIndexFileKind::Multivalue => { - serialize_bitmap_datum - } - GlobalIndexFileKind::FM => unreachable!("FM range query was rejected in planning"), - }; - let from_key = serialize_key(between.from, between.data_type); - let to_key = serialize_key(between.to, between.data_type); - let bitmap = reader - .as_ref() - .expect("reader is opened when between matches") - .range_query( - &from_key, - &to_key, - between.data_type, - between.from_inclusive, - between.to_inclusive, - ) - .await - .map_err(|error| Self::query_error(entry, error))?; - file_result = Some(bitmap); - } - - for &idx in &plan.matching_predicates { - let (op, literals, data_type) = &effective_predicates[idx]; - let Some(bitmap) = reader - .as_ref() - .expect("reader is opened when predicates match") - .query(*op, literals, data_type) - .await - .map_err(|error| Self::query_error(entry, error))? - else { - return Ok(EntryQueryResult { - bitmap: None, - declined: true, - }); - }; - file_result = Some(match file_result { - None => bitmap, - Some(mut existing) => { - existing &= bitmap; - existing - } - }); - } - - // Each concurrent task owns its reader. Only return it to the shared - // cache after all predicates for this shard have completed. - if let Some(OpenedGlobalIndexReader::BTree(reader)) = reader.take() { - self.return_reader(entry.file_name.clone(), reader); - } - Ok(EntryQueryResult { - bitmap: file_result, - declined: false, - }) - } - - fn query_error(entry: &GlobalIndexEntry, error: std::io::Error) -> Error { - Error::DataInvalid { - message: format!( - "Global index query failed for {} file '{}'", - entry.index_type.name(), - entry.file_name - ), - source: Some(Box::new(error)), - } - } - - /// Get a cached reader or open a new one for the given file. - async fn get_or_open_reader( - &self, - entry: &GlobalIndexEntry, - meta: &BTreeIndexMeta, - data_type: &DataType, - ) -> Result { - // Try to take from cache - { - let mut cache = self.reader_cache.lock().unwrap(); - if let Some(reader) = cache.remove(&entry.file_name) { - return Ok(OpenedGlobalIndexReader::BTree(reader)); - } - } - - // Open new reader - let path = format!("{}/{INDEX_DIR}/{}", self.table_path, entry.file_name); - let input = self.file_io.new_input(&path)?; - let file_size = if entry.file_size > 0 { - entry.file_size as u64 - } else { - input.metadata().await?.size - }; - let file_reader = input.reader().await?; - - let cmp = make_key_comparator(data_type); - BTreeIndexReader::open(Box::new(file_reader), file_size, meta, cmp) - .await - .map(OpenedGlobalIndexReader::BTree) - .map_err(|e| crate::Error::DataInvalid { - message: format!("Failed to open BTree index file: {}", entry.file_name), - source: Some(Box::new(e)), - }) - } - - async fn open_reader_for_entry( - &self, - entry: &GlobalIndexEntry, - data_type: &DataType, - ) -> Result { - match entry.index_type { - GlobalIndexFileKind::BTree => { - self.get_or_open_reader(entry, sorted_entry_meta(entry), data_type) - .await - } - GlobalIndexFileKind::Bitmap => self - .open_bitmap_reader(entry) - .await - .map(OpenedGlobalIndexReader::Bitmap) - .map_err(|e| crate::Error::DataInvalid { - message: format!( - "Failed to open bitmap global index file: {}", - entry.file_name - ), - source: Some(Box::new(e)), - }), - GlobalIndexFileKind::Multivalue => self - .open_bitmap_reader(entry) - .await - .map(OpenedGlobalIndexReader::Bitmap) - .map_err(|e| crate::Error::DataInvalid { - message: format!( - "Failed to open multivalue global index file: {}", - entry.file_name - ), - source: Some(Box::new(e)), - }), - GlobalIndexFileKind::FM => self - .open_fm_reader(entry) - .await - .map(OpenedGlobalIndexReader::FM) - .map_err(|e| crate::Error::DataInvalid { - message: format!("Failed to open FM global index file: {}", entry.file_name), - source: Some(Box::new(e)), - }), - } - } - - async fn open_fm_reader( - &self, - entry: &GlobalIndexEntry, - ) -> std::io::Result { - let GlobalIndexEntryMeta::FM { - bytes: manifest_meta, - .. - } = &entry.meta - else { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "FM entry has non-FM manifest metadata", - )); - }; - let path = format!("{}/{INDEX_DIR}/{}", self.table_path, entry.file_name); - let input = self - .file_io - .new_input(&path) - .map_err(|e| std::io::Error::other(e.to_string()))?; - let file_size = if entry.file_size > 0 { - entry.file_size as u64 - } else { - input - .metadata() - .await - .map_err(|e| std::io::Error::other(e.to_string()))? - .size - }; - let file_reader = input - .reader() - .await - .map_err(|e| std::io::Error::other(e.to_string()))?; - FMGlobalIndexReader::open_with_context( - Box::new(file_reader), - file_size, - manifest_meta, - self.fm_read_options, - Arc::clone(&self.fm_read_context), - entry.file_name.clone(), - ) - .await - } - - async fn open_bitmap_reader( - &self, - entry: &GlobalIndexEntry, - ) -> std::io::Result { - let path = format!("{}/{INDEX_DIR}/{}", self.table_path, entry.file_name); - let input = self - .file_io - .new_input(&path) - .map_err(|e| std::io::Error::other(e.to_string()))?; - let file_size = if entry.file_size > 0 { - entry.file_size as u64 - } else { - input - .metadata() - .await - .map_err(|e| std::io::Error::other(e.to_string()))? - .size - }; - let file_reader = input - .reader() - .await - .map_err(|e| std::io::Error::other(e.to_string()))?; - BitmapGlobalIndexReader::open(Box::new(file_reader), file_size).await - } - - fn fallback_scan_plan( - &self, - entries: &[&GlobalIndexEntry], - selected: &[bool], - ) -> FallbackScanPlan { - let mut plan = FallbackScanPlan::default(); - let mut btree_total = 0i64; - let mut bitmap_total = 0i64; - let mut btree_valid = true; - let mut bitmap_valid = true; - - for (entry, selected) in entries.iter().zip(selected) { - if !selected { - continue; - } - match entry.index_type { - GlobalIndexFileKind::BTree => { - plan.selected_btree += 1; - btree_valid &= add_file_size(&mut btree_total, entry.file_size); - } - GlobalIndexFileKind::Bitmap => { - plan.selected_bitmap += 1; - bitmap_valid &= add_file_size(&mut bitmap_total, entry.file_size); - } - GlobalIndexFileKind::Multivalue => { - plan.selected_bitmap += 1; - bitmap_valid &= add_file_size(&mut bitmap_total, entry.file_size); - } - GlobalIndexFileKind::FM => {} - } - } - - plan.allow_btree = plan.selected_btree > 0 - && btree_valid - && self.btree_fallback_scan_max_size > 0 - && btree_total <= self.btree_fallback_scan_max_size; - plan.allow_bitmap = plan.selected_bitmap > 0 - && bitmap_valid - && self.bitmap_fallback_scan_max_size > 0 - && bitmap_total <= self.bitmap_fallback_scan_max_size; - plan - } - - /// Return a reader to the cache for future reuse. - fn return_reader(&self, file_name: String, reader: BTreeIndexReader) { - let mut cache = self.reader_cache.lock().unwrap(); - cache.insert(file_name, reader); - } - - fn find_field_id_by_name(&self, column: &str) -> Result> { - Ok(crate::table::find_field_id_by_name( - &self.schema_fields, - column, - )) - } - - fn entries_for_field(&self, field_id: i32) -> Option<&[GlobalIndexEntry]> { - self.entries_by_field - .iter() - .find(|(id, _)| *id == field_id) - .map(|(_, entries)| entries.as_slice()) - } - - /// Return row ranges not covered by global indexes for this predicate. - /// - /// `full` uses `[0, snapshot.next_row_id - 1]`; `detail` uses actual - /// data-file row ranges collected by the scan. The caller unions these - /// ranges with indexed matches, and the normal read filter evaluates the - /// predicate on the raw rows. - #[cfg(test)] - fn unindexed_ranges( - &self, - predicate: &Predicate, - search_mode: GlobalIndexSearchMode, - next_row_id: Option, - data_ranges: &[RowRange], - ) -> Result> { - let field_ids = self.collect_field_ids(predicate)?; - Ok(self.unindexed_ranges_for_field_ids(&field_ids, search_mode, next_row_id, data_ranges)) - } - - #[cfg(test)] - fn unindexed_ranges_for_field_ids( - &self, - field_ids: &HashSet, - search_mode: GlobalIndexSearchMode, - next_row_id: Option, - data_ranges: &[RowRange], - ) -> Vec { - unindexed_ranges_for_coverage( - &self.coverage_by_field, - field_ids, - search_mode, - next_row_id, - data_ranges, - ) - } - - #[cfg(test)] - fn collect_field_ids(&self, predicate: &Predicate) -> Result> { - let mut field_ids = HashSet::new(); - self.collect_field_ids_inner(predicate, &mut field_ids)?; - Ok(field_ids) - } - - #[cfg(test)] - fn collect_field_ids_inner( - &self, - predicate: &Predicate, - field_ids: &mut HashSet, - ) -> Result<()> { - match predicate { - Predicate::Leaf { column, .. } => { - if let Some(field_id) = self.find_field_id_by_name(column)? { - field_ids.insert(field_id); - } - } - Predicate::And(children) | Predicate::Or(children) => { - for child in children { - self.collect_field_ids_inner(child, field_ids)?; - } - } - Predicate::Not(inner) => self.collect_field_ids_inner(inner, field_ids)?, - Predicate::AlwaysTrue | Predicate::AlwaysFalse => {} - } - Ok(()) - } -} - -/// Whether the sorted global index can evaluate this operator directly. -/// Operators that fall outside this set bypass the index and are evaluated -/// later in the read pipeline (stats prune + parquet row filter). -fn is_sorted_global_index_supported_op(op: PredicateOperator) -> bool { - matches!( - op, - PredicateOperator::Eq - | PredicateOperator::NotEq - | PredicateOperator::Lt - | PredicateOperator::LtEq - | PredicateOperator::Gt - | PredicateOperator::GtEq - | PredicateOperator::In - | PredicateOperator::NotIn - | PredicateOperator::IsNull - | PredicateOperator::IsNotNull - | PredicateOperator::Between - | PredicateOperator::NotBetween - | PredicateOperator::StartsWith - | PredicateOperator::EndsWith - | PredicateOperator::Contains - | PredicateOperator::Like - | PredicateOperator::ArrayContains - | PredicateOperator::ArraysOverlap - | PredicateOperator::ArrayContainsAll - ) -} - -fn is_multivalue_predicate(op: PredicateOperator) -> bool { - matches!( - op, - PredicateOperator::ArrayContains - | PredicateOperator::ArraysOverlap - | PredicateOperator::ArrayContainsAll - ) } -fn entry_supports_predicate( - entry: &GlobalIndexEntry, - op: PredicateOperator, - literals: &[Datum], -) -> bool { - match entry.index_type { - GlobalIndexFileKind::Multivalue => { - is_multivalue_predicate(op) - && !(matches!(op, PredicateOperator::ArrayContainsAll) && literals.is_empty()) - } - GlobalIndexFileKind::FM => { - matches!(op, PredicateOperator::IsNull | PredicateOperator::IsNotNull) - || (op == PredicateOperator::Contains && literals.len() == 1) - } - GlobalIndexFileKind::BTree | GlobalIndexFileKind::Bitmap => !is_multivalue_predicate(op), - } +/// Create a GlobalIndexScanner and evaluate predicates, returning row ranges. +/// This is the main entry point for the table scan integration. +/// +/// Returns `None` if global index is not available or predicates can't be evaluated. +pub(crate) struct GlobalIndexEvaluation<'a> { + pub(crate) file_io: &'a FileIO, + pub(crate) table_path: &'a str, + pub(crate) index_entries: &'a [IndexManifestEntry], + pub(crate) predicates: &'a [Predicate], + pub(crate) schema_fields: &'a [DataField], + pub(crate) search_mode: GlobalIndexSearchMode, + pub(crate) global_index_thread_num: usize, + pub(crate) btree_fallback_scan_max_size: i64, + pub(crate) bitmap_fallback_scan_max_size: i64, + pub(crate) fm_read_options: FMReadOptions, + pub(crate) next_row_id: Option, + pub(crate) data_ranges: &'a [RowRange], } -fn entries_support_predicate( - entries: &[GlobalIndexEntry], - op: PredicateOperator, - literals: &[Datum], -) -> bool { - entries - .iter() - .any(|entry| entry_supports_predicate(entry, op, literals)) -} +pub(crate) async fn evaluate_global_index( + evaluation: GlobalIndexEvaluation<'_>, +) -> Result>> { + let scanner = match GlobalIndexScanner::create_with_fm_options( + evaluation.file_io, + evaluation.table_path, + evaluation.global_index_thread_num, + evaluation.btree_fallback_scan_max_size, + evaluation.bitmap_fallback_scan_max_size, + evaluation.index_entries, + evaluation.schema_fields, + evaluation.fm_read_options, + )? { + Some(s) => s, + None => return Ok(None), + }; -fn select_entries_for_predicates<'a>( - entries: &'a [GlobalIndexEntry], - predicates: &[(PredicateOperator, &[Datum], &DataType)], -) -> Vec<&'a GlobalIndexEntry> { - let compatible = entries - .iter() - .filter(|entry| { - predicates - .iter() - .all(|(op, literals, _)| entry_supports_predicate(entry, *op, literals)) - }) - .collect::>(); - - if predicates - .iter() - .any(|(op, _, _)| *op == PredicateOperator::Contains) - && compatible - .iter() - .any(|entry| entry.index_type == GlobalIndexFileKind::FM) - { - compatible - .into_iter() - .filter(|entry| entry.index_type == GlobalIndexFileKind::FM) - .collect() - } else { - compatible - } -} - -fn requires_fallback_scan(op: PredicateOperator) -> bool { - matches!( - op, - PredicateOperator::Lt - | PredicateOperator::LtEq - | PredicateOperator::Gt - | PredicateOperator::GtEq - | PredicateOperator::Between - | PredicateOperator::NotBetween - | PredicateOperator::EndsWith - | PredicateOperator::Contains - | PredicateOperator::Like - ) -} - -fn fallback_plan_evaluates_entry( - plan: FallbackScanPlan, - kind: GlobalIndexFileKind, - selected: bool, -) -> bool { - !selected || plan.allowed(kind) -} - -fn add_file_size(total: &mut i64, file_size: i64) -> bool { - if file_size < 0 { - return false; - } - match total.checked_add(file_size) { - Some(next) => { - *total = next; - true - } - None => false, - } -} - -/// Convert a RoaringTreemap to merged RowRanges (already sorted and deduplicated). -fn bitmap_to_ranges(bitmap: &RoaringTreemap) -> Vec { - if bitmap.is_empty() { - return Vec::new(); - } - let mut ranges = Vec::new(); - let mut iter = bitmap.iter(); - let first = iter.next().unwrap(); - let mut start = first as i64; - let mut end = start; - - for id in iter { - let id = id as i64; - if id == end + 1 { - end = id; - } else { - ranges.push(RowRange::new(start, end)); - start = id; - end = id; - } - } - ranges.push(RowRange::new(start, end)); - ranges -} - -/// Intersect two sorted range lists using RowRangeIndex for efficient binary search. -fn intersect_sorted_ranges(a: &[RowRange], b: &[RowRange]) -> Vec { - let idx = RowRangeIndex::create(a.to_vec()); - let mut result = Vec::new(); - for r in b { - result.extend(idx.intersected_ranges(r.from(), r.to())); - } - result -} - -fn data_ranges_for_search_mode( - search_mode: GlobalIndexSearchMode, - next_row_id: Option, - data_ranges: &[RowRange], -) -> Option> { - match search_mode { - GlobalIndexSearchMode::Fast => None, - GlobalIndexSearchMode::Full => match next_row_id { - Some(next_row_id) if next_row_id > 0 => Some(vec![RowRange::new(0, next_row_id - 1)]), - _ => None, - }, - GlobalIndexSearchMode::Detail => { - if data_ranges.is_empty() { - None - } else { - Some(data_ranges.to_vec()) - } - } - } -} - -fn indexed_ranges_from_coverage( - coverage_by_field: &HashMap>, - field_ids: &HashSet, -) -> Vec { - let mut ranges: Option> = None; - for field_id in field_ids { - let Some(field_ranges) = coverage_by_field.get(field_id) else { - return Vec::new(); - }; - if field_ranges.is_empty() { - return Vec::new(); - } - let field_ranges = super::merge_row_ranges(field_ranges.clone()); - ranges = Some(match ranges { - None => field_ranges, - Some(existing) => intersect_sorted_ranges(&existing, &field_ranges), - }); - } - ranges.map(super::merge_row_ranges).unwrap_or_default() -} - -fn unindexed_ranges_for_coverage( - coverage_by_field: &HashMap>, - field_ids: &HashSet, - search_mode: GlobalIndexSearchMode, - next_row_id: Option, - data_ranges: &[RowRange], -) -> Vec { - let Some(data_ranges) = data_ranges_for_search_mode(search_mode, next_row_id, data_ranges) - else { - return Vec::new(); - }; - let indexed_ranges = indexed_ranges_from_coverage(coverage_by_field, field_ids); - super::source::exclude_row_ranges(&data_ranges, &indexed_ranges) -} - -fn unindexed_ranges_for_indexed_coverage( - indexed_ranges: &[RowRange], - search_mode: GlobalIndexSearchMode, - next_row_id: Option, - data_ranges: &[RowRange], -) -> Vec { - let Some(data_ranges) = data_ranges_for_search_mode(search_mode, next_row_id, data_ranges) - else { - return Vec::new(); - }; - super::source::exclude_row_ranges( - &data_ranges, - &super::merge_row_ranges(indexed_ranges.to_vec()), - ) -} - -/// Compute row ranges not covered by a family of global index files. -/// -/// This mirrors Java `GlobalIndexCoverage`: `full` compares index coverage -/// against `[0, snapshot.next_row_id - 1]`, while `detail` compares against -/// exact data-file row ranges supplied by the caller. -pub(crate) fn unindexed_ranges_for_global_index_entries( - index_entries: &[IndexManifestEntry], - field_ids: &HashSet, - search_mode: GlobalIndexSearchMode, - next_row_id: Option, - data_ranges: &[RowRange], - index_file_filter: impl Fn(&IndexFileMeta) -> bool, -) -> Vec { - let mut coverage_by_field: HashMap> = HashMap::new(); - for entry in index_entries { - if entry.kind != FileKind::Add || !index_file_filter(&entry.index_file) { - continue; - } - let Some(global_meta) = entry.index_file.global_index_meta.as_ref() else { - continue; - }; - let row_range = RowRange::new(global_meta.row_range_start, global_meta.row_range_end); - coverage_by_field - .entry(global_meta.index_field_id) - .or_default() - .push(row_range.clone()); - if let Some(extra_field_ids) = global_meta.extra_field_ids.as_ref() { - for extra_field_id in extra_field_ids { - coverage_by_field - .entry(*extra_field_id) - .or_default() - .push(row_range.clone()); - } - } - } - unindexed_ranges_for_coverage( - &coverage_by_field, - field_ids, - search_mode, - next_row_id, - data_ranges, - ) -} - -/// Resolve live deletion-vector index entries into global row-id ranges. -/// -/// Data-evolution DV entries are keyed by the normal anchor data file. The DV -/// bitmap positions are local to that anchor file's `first_row_id`, so this -/// helper joins index metadata with live data-file metadata before converting -/// deleted positions to global row IDs. -pub(crate) async fn deleted_row_ranges_for_data_evolution_dvs( - table: &Table, - index_entries: &[IndexManifestEntry], -) -> Result> { - if !index_entries.iter().any(|entry| { - entry.kind == FileKind::Add && entry.index_file.index_type == DELETION_VECTORS_INDEX_TYPE - }) { - return Ok(Vec::new()); - } - - let plan = table - .new_read_builder() - .new_scan() - .with_scan_all_files() - .plan() - .await?; - - let mut first_row_ids: HashMap<(Vec, i32, String), i64> = HashMap::new(); - for split in plan.splits() { - let partition = split.partition().to_serialized_bytes(); - let bucket = split.bucket(); - for file in split.data_files() { - if let Some(first_row_id) = file.first_row_id { - first_row_ids.insert( - (partition.clone(), bucket, file.file_name.clone()), - first_row_id, - ); - } - } - } - - let mut ranges = Vec::new(); - let table_path = table.location().trim_end_matches('/'); - for entry in index_entries { - if entry.kind != FileKind::Add || entry.index_file.index_type != DELETION_VECTORS_INDEX_TYPE - { - continue; - } - let Some(dv_ranges) = entry.index_file.deletion_vectors_ranges.as_ref() else { - continue; - }; - let index_path = format!("{table_path}/{INDEX_DIR}/{}", entry.index_file.file_name); - for (data_file_name, meta) in dv_ranges { - let key = ( - entry.partition.clone(), - entry.bucket, - data_file_name.clone(), - ); - let first_row_id = first_row_ids.get(&key).copied().ok_or_else(|| { - crate::Error::DataInvalid { - message: format!( - "Deletion vector references data file '{}' but no live row-tracked file was found", - data_file_name - ), - source: None, - } - })?; - let deletion_file = DeletionFile::new( - index_path.clone(), - meta.offset as i64, - meta.length as i64, - meta.cardinality, - ); - let deletion_vector = - DeletionVectorFactory::read(table.file_io(), &deletion_file).await?; - for deleted in deletion_vector.iter() { - let deleted = i64::try_from(deleted).map_err(|_| crate::Error::DataInvalid { - message: format!( - "Deleted position {deleted} for data file '{}' exceeds i64::MAX", - data_file_name - ), - source: None, - })?; - let row_id = - first_row_id - .checked_add(deleted) - .ok_or_else(|| crate::Error::DataInvalid { - message: format!( - "Deleted row id overflows i64 for data file '{}'", - data_file_name - ), - source: None, - })?; - ranges.push(RowRange::new(row_id, row_id)); - } - } - } - - Ok(super::merge_row_ranges(ranges)) -} - -/// Index for row ranges. Stores sorted, non-overlapping ranges and supports -/// efficient intersection queries via binary search. -/// -/// Reference: [org.apache.paimon.utils.RowRangeIndex](https://github.com/apache/paimon/blob/master/paimon-common/src/main/java/org/apache/paimon/utils/RowRangeIndex.java) -#[derive(Debug, Clone)] -pub(crate) struct RowRangeIndex { - ranges: Vec, - starts: Vec, - ends: Vec, -} - -impl RowRangeIndex { - /// Create a new `RowRangeIndex` from a list of ranges. - /// Ranges are sorted and merged to eliminate overlaps. - pub fn create(ranges: Vec) -> Self { - let ranges = super::merge_row_ranges(ranges); - let starts: Vec = ranges.iter().map(|r| r.from()).collect(); - let ends: Vec = ranges.iter().map(|r| r.to()).collect(); - Self { - ranges, - starts, - ends, - } - } - - /// Returns the underlying ranges. - #[cfg(test)] - pub fn ranges(&self) -> &[RowRange] { - &self.ranges - } - - /// Returns true if the index has any range that intersects `[start, end]`. - pub fn intersects(&self, start: i64, end: i64) -> bool { - let candidate = lower_bound(&self.ends, start); - candidate < self.starts.len() && self.starts[candidate] <= end - } - - /// Counts rows in this index that intersect `[start, end]`. - pub fn intersection_row_count(&self, start: i64, end: i64) -> usize { - if start > end { - return 0; - } - self.intersected_ranges(start, end) - .into_iter() - .fold(0usize, |total, range| { - let len = range.to().saturating_sub(range.from()).saturating_add(1); - total.saturating_add(usize::try_from(len).unwrap_or(usize::MAX)) - }) - } - - /// Returns the sub-ranges of this index that intersect `[start, end]`, - /// clipped to the intersection boundaries. - pub fn intersected_ranges(&self, start: i64, end: i64) -> Vec { - let left = lower_bound(&self.ends, start); - if left >= self.ranges.len() || self.starts[left] > end { - return Vec::new(); - } - - let mut right = lower_bound(&self.ends, end); - if right >= self.ranges.len() { - right = self.ranges.len() - 1; - } - - let mut result = Vec::new(); - - // First range: clip from the left - let first = &self.ranges[left]; - result.push(RowRange::new(start.max(first.from()), end.min(first.to()))); - - // Middle ranges: fully included - if right > left + 1 { - for r in &self.ranges[left + 1..right] { - result.push(r.clone()); - } - } - - // Last range (if different from first): clip from the right - if right != left { - let last = &self.ranges[right]; - if last.from() <= end { - result.push(RowRange::new(start.max(last.from()), end.min(last.to()))); - } - } - - result - } -} - -pub(crate) fn search_limit_with_deleted_rows( - limit: usize, - row_range_start: i64, - row_range_end: i64, - deleted_rows: Option<&RowRangeIndex>, -) -> usize { - let Some(range_len) = row_range_end - .checked_sub(row_range_start) - .and_then(|len| len.checked_add(1)) - .and_then(|len| usize::try_from(len).ok()) - else { - return limit; - }; - - let deleted_count = deleted_rows - .map(|index| index.intersection_row_count(row_range_start, row_range_end)) - .unwrap_or(0) - .min(range_len); - limit.saturating_add(deleted_count).min(range_len) -} - -/// Binary search: find the first index where `sorted[index] >= target`. -fn lower_bound(sorted: &[i64], target: i64) -> usize { - let mut left = 0; - let mut right = sorted.len(); - while left < right { - let mid = left + (right - left) / 2; - if sorted[mid] < target { - left = mid + 1; - } else { - right = mid; - } - } - left -} - -/// Create a GlobalIndexScanner and evaluate predicates, returning row ranges. -/// This is the main entry point for the table scan integration. -/// -/// Returns `None` if global index is not available or predicates can't be evaluated. -pub(crate) struct GlobalIndexEvaluation<'a> { - pub(crate) file_io: &'a FileIO, - pub(crate) table_path: &'a str, - pub(crate) index_entries: &'a [IndexManifestEntry], - pub(crate) predicates: &'a [Predicate], - pub(crate) schema_fields: &'a [DataField], - pub(crate) search_mode: GlobalIndexSearchMode, - pub(crate) global_index_thread_num: usize, - pub(crate) btree_fallback_scan_max_size: i64, - pub(crate) bitmap_fallback_scan_max_size: i64, - pub(crate) fm_read_options: FMReadOptions, - pub(crate) next_row_id: Option, - pub(crate) data_ranges: &'a [RowRange], -} - -pub(crate) async fn evaluate_global_index( - evaluation: GlobalIndexEvaluation<'_>, -) -> Result>> { - let scanner = match GlobalIndexScanner::create_with_fm_options( - evaluation.file_io, - evaluation.table_path, - evaluation.global_index_thread_num, - evaluation.btree_fallback_scan_max_size, - evaluation.bitmap_fallback_scan_max_size, - evaluation.index_entries, - evaluation.schema_fields, - evaluation.fm_read_options, - )? { - Some(s) => s, - None => return Ok(None), - }; - - let combined = Predicate::and(evaluation.predicates.to_vec()); + let combined = Predicate::and(evaluation.predicates.to_vec()); let scan_result = match scanner.evaluate(&combined).await? { Some(scan_result) => scan_result, @@ -1989,2141 +370,4 @@ pub(crate) async fn evaluate_global_index( } #[cfg(test)] -mod tests { - use super::*; - use crate::btree::test_util::VecFileWrite; - use crate::btree::{BTreeIndexWriter, BlockCompressionType}; - use crate::fm_index::{FMGlobalIndexWriter, FMWriteOptions}; - use crate::table::bitmap_global_index_writer::BitmapGlobalIndexWriter; - use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; - use std::sync::Arc; - - #[tokio::test] - async fn test_try_fold_bounded_respects_concurrency_limit() { - for limit in [1, 3] { - let active = Arc::new(AtomicUsize::new(0)); - let peak = Arc::new(AtomicUsize::new(0)); - let futures = (0..9usize).map(|value| { - let active = Arc::clone(&active); - let peak = Arc::clone(&peak); - async move { - let current = active.fetch_add(1, AtomicOrdering::SeqCst) + 1; - peak.fetch_max(current, AtomicOrdering::SeqCst); - tokio::task::yield_now().await; - active.fetch_sub(1, AtomicOrdering::SeqCst); - Ok::<_, crate::Error>(value) - } - }); - - let mut values = try_fold_bounded(futures, limit, Vec::new(), |values, value| { - values.push(value) - }) - .await - .unwrap(); - values.sort_unstable(); - - assert_eq!(values, (0..9).collect::>()); - assert_eq!(peak.load(AtomicOrdering::SeqCst), limit); - } - } - - #[test] - fn test_bitmap_to_ranges() { - assert_eq!( - bitmap_to_ranges(&RoaringTreemap::new()), - Vec::::new() - ); - - let mut bm = RoaringTreemap::new(); - bm.insert(5); - assert_eq!(bitmap_to_ranges(&bm), vec![RowRange::new(5, 5)]); - - let mut bm = RoaringTreemap::new(); - for id in [1, 2, 3, 5, 6, 10] { - bm.insert(id); - } - assert_eq!( - bitmap_to_ranges(&bm), - vec![ - RowRange::new(1, 3), - RowRange::new(5, 6), - RowRange::new(10, 10), - ] - ); - } - - #[test] - fn test_intersect_sorted_ranges() { - let a = vec![RowRange::new(0, 10), RowRange::new(20, 30)]; - let b = vec![RowRange::new(5, 25)]; - let result = intersect_sorted_ranges(&a, &b); - assert_eq!(result, vec![RowRange::new(5, 10), RowRange::new(20, 25)]); - } - - #[test] - fn test_intersect_no_overlap() { - let a = vec![RowRange::new(0, 5)]; - let b = vec![RowRange::new(10, 20)]; - assert!(intersect_sorted_ranges(&a, &b).is_empty()); - } - - #[test] - fn test_serialize_datum_int() { - let key = serialize_datum(&Datum::Int(42), &DataType::Int(crate::spec::IntType::new())); - assert_eq!(key, 42i32.to_le_bytes().to_vec()); - } - - #[test] - fn test_serialize_datum_string() { - let key = serialize_datum( - &Datum::String("hello".to_string()), - &DataType::VarChar(crate::spec::VarCharType::new(100).unwrap()), - ); - assert_eq!(key, b"hello".to_vec()); - } - - fn assert_bitmap_floating_meta_policy( - data_type: DataType, - min: Datum, - max: Datum, - outside: Datum, - nan: Datum, - ) { - let cmp = make_bitmap_key_comparator(&data_type); - let min_key = serialize_bitmap_datum(&min, &data_type); - let max_key = serialize_bitmap_datum(&max, &data_type); - let outside_key = serialize_bitmap_datum(&outside, &data_type); - let nan_key = serialize_bitmap_datum(&nan, &data_type); - let meta = BTreeIndexMeta::new(Some(min_key.clone()), Some(max_key), false); - - assert!(!bitmap_meta_may_match( - &meta, - PredicateOperator::Eq, - &data_type, - std::slice::from_ref(&outside_key), - cmp.as_ref(), - )); - assert!(!bitmap_meta_may_match( - &meta, - PredicateOperator::In, - &data_type, - std::slice::from_ref(&outside_key), - cmp.as_ref(), - )); - assert!(!bitmap_meta_may_match( - &meta, - PredicateOperator::IsNull, - &data_type, - &[], - cmp.as_ref(), - )); - assert!(bitmap_meta_may_match( - &meta, - PredicateOperator::IsNotNull, - &data_type, - &[], - cmp.as_ref(), - )); - - let nan_meta = BTreeIndexMeta::new(Some(min_key), Some(nan_key.clone()), false); - assert!(bitmap_meta_may_match( - &nan_meta, - PredicateOperator::Eq, - &data_type, - std::slice::from_ref(&nan_key), - cmp.as_ref(), - )); - assert!(bitmap_meta_may_match( - &nan_meta, - PredicateOperator::In, - &data_type, - std::slice::from_ref(&nan_key), - cmp.as_ref(), - )); - - assert!(bitmap_meta_may_match( - &meta, - PredicateOperator::Gt, - &data_type, - std::slice::from_ref(&outside_key), - cmp.as_ref(), - )); - assert!(bitmap_meta_may_match_between( - &meta, - &data_type, - &outside_key, - &outside_key, - cmp.as_ref(), - )); - - let only_nulls = BTreeIndexMeta::new(None, None, true); - assert!(bitmap_meta_may_match( - &only_nulls, - PredicateOperator::IsNull, - &data_type, - &[], - cmp.as_ref(), - )); - assert!(!bitmap_meta_may_match( - &only_nulls, - PredicateOperator::IsNotNull, - &data_type, - &[], - cmp.as_ref(), - )); - assert!(!bitmap_meta_may_match( - &only_nulls, - PredicateOperator::NotEq, - &data_type, - std::slice::from_ref(&outside_key), - cmp.as_ref(), - )); - assert!(!bitmap_meta_may_match_between( - &only_nulls, - &data_type, - &outside_key, - &outside_key, - cmp.as_ref(), - )); - } - - #[test] - fn test_bitmap_floating_meta_prunes_equality_and_fails_open_for_ranges() { - assert_bitmap_floating_meta_policy( - DataType::Float(crate::spec::FloatType::new()), - Datum::Float(-1.0), - Datum::Float(1.0), - Datum::Float(2.0), - Datum::Float(f32::NAN), - ); - assert_bitmap_floating_meta_policy( - DataType::Double(crate::spec::DoubleType::new()), - Datum::Double(-1.0), - Datum::Double(1.0), - Datum::Double(2.0), - Datum::Double(f64::NAN), - ); - } - - #[test] - fn test_row_range_index_merges_overlapping() { - let idx = RowRangeIndex::create(vec![ - RowRange::new(0, 5), - RowRange::new(3, 10), - RowRange::new(20, 30), - ]); - assert_eq!(idx.ranges().len(), 2); - assert_eq!(idx.ranges()[0], RowRange::new(0, 10)); - assert_eq!(idx.ranges()[1], RowRange::new(20, 30)); - } - - #[test] - fn test_row_range_index_merges_adjacent() { - let idx = RowRangeIndex::create(vec![RowRange::new(0, 5), RowRange::new(6, 10)]); - assert_eq!(idx.ranges().len(), 1); - assert_eq!(idx.ranges()[0], RowRange::new(0, 10)); - } - - #[test] - fn test_row_range_index_intersects() { - let idx = RowRangeIndex::create(vec![RowRange::new(10, 20), RowRange::new(30, 40)]); - assert!(idx.intersects(15, 25)); - assert!(idx.intersects(5, 10)); - assert!(idx.intersects(20, 30)); - assert!(!idx.intersects(0, 9)); - assert!(!idx.intersects(21, 29)); - assert!(!idx.intersects(41, 50)); - } - - #[test] - fn test_row_range_index_intersected_ranges() { - let idx = RowRangeIndex::create(vec![ - RowRange::new(10, 20), - RowRange::new(30, 40), - RowRange::new(50, 60), - ]); - let result = idx.intersected_ranges(15, 55); - assert_eq!( - result, - vec![ - RowRange::new(15, 20), - RowRange::new(30, 40), - RowRange::new(50, 55), - ] - ); - } - - #[test] - fn test_row_range_index_intersection_row_count() { - let idx = RowRangeIndex::create(vec![ - RowRange::new(10, 20), - RowRange::new(30, 40), - RowRange::new(50, 60), - ]); - - assert_eq!(idx.intersection_row_count(15, 55), 23); - assert_eq!(idx.intersection_row_count(21, 29), 0); - assert_eq!(idx.intersection_row_count(55, 15), 0); - } - - #[test] - fn test_search_limit_with_deleted_rows_expands_and_caps() { - let idx = RowRangeIndex::create(vec![RowRange::new(2, 4), RowRange::new(8, 10)]); - - assert_eq!(search_limit_with_deleted_rows(5, 0, 19, Some(&idx)), 11); - assert_eq!(search_limit_with_deleted_rows(18, 0, 19, Some(&idx)), 20); - assert_eq!(search_limit_with_deleted_rows(5, 0, 19, None), 5); - } - - #[test] - fn test_row_range_index_empty() { - let idx = RowRangeIndex::create(Vec::new()); - assert!(!idx.intersects(0, 100)); - assert!(idx.intersected_ranges(0, 100).is_empty()); - } - - fn le_int_key(v: i32) -> Vec { - v.to_le_bytes().to_vec() - } - - /// Set up a temp dir with `index/{file_name}` containing the btree testdata file, - /// and return (FileIO, table_path, file_name, _tmp_dir). - fn setup_testdata_table(testdata_name: &str) -> (FileIO, String, String, tempfile::TempDir) { - let src = format!( - "{}/testdata/btree/{testdata_name}", - env!("CARGO_MANIFEST_DIR") - ); - let tmp = tempfile::tempdir().unwrap(); - let index_dir = tmp.path().join("index"); - std::fs::create_dir_all(&index_dir).unwrap(); - std::fs::copy(&src, index_dir.join(testdata_name)).unwrap(); - - let table_path = format!("file://{}", tmp.path().display()); - let file_io = crate::io::FileIOBuilder::new("file").build().unwrap(); - (file_io, table_path, testdata_name.to_string(), tmp) - } - - type BitmapTestdataTable = (FileIO, String, String, BTreeIndexMeta, tempfile::TempDir); - - fn setup_bitmap_testdata_table(file_name: &str) -> BitmapTestdataTable { - let src = format!("{}/testdata/bitmap/{file_name}", env!("CARGO_MANIFEST_DIR")); - let meta_src = format!( - "{}/testdata/bitmap/{file_name}.meta", - env!("CARGO_MANIFEST_DIR") - ); - let tmp = tempfile::tempdir().unwrap(); - let index_dir = tmp.path().join("index"); - std::fs::create_dir_all(&index_dir).unwrap(); - std::fs::copy(&src, index_dir.join(file_name)).unwrap(); - let meta = BTreeIndexMeta::deserialize(&std::fs::read(meta_src).unwrap()).unwrap(); - - let table_path = format!("file://{}", tmp.path().display()); - let file_io = crate::io::FileIOBuilder::new("file").build().unwrap(); - (file_io, table_path, file_name.to_string(), meta, tmp) - } - - fn setup_java_bitmap_testdata_table() -> BitmapTestdataTable { - setup_bitmap_testdata_table("bitmap_varchar_java.index") - } - - fn make_global_index_entry( - file_name: &str, - field_id: i32, - row_range_start: i64, - row_range_end: i64, - meta: &BTreeIndexMeta, - ) -> crate::spec::IndexManifestEntry { - make_global_index_entry_with_type( - BTREE_GLOBAL_INDEX_TYPE, - file_name, - field_id, - row_range_start, - row_range_end, - meta, - ) - } - - fn make_global_index_entry_with_type( - index_type: &str, - file_name: &str, - field_id: i32, - row_range_start: i64, - row_range_end: i64, - meta: &BTreeIndexMeta, - ) -> crate::spec::IndexManifestEntry { - use crate::spec::{GlobalIndexMeta, IndexFileMeta}; - IndexManifestEntry { - version: 1, - kind: FileKind::Add, - partition: vec![], - bucket: 0, - index_file: IndexFileMeta { - index_type: index_type.to_string(), - file_name: file_name.to_string(), - file_size: 0, - row_count: 0, - deletion_vectors_ranges: None, - global_index_meta: Some(GlobalIndexMeta { - row_range_start, - row_range_end, - index_field_id: field_id, - extra_field_ids: None, - source_meta: None, - index_meta: Some(meta.serialize()), - }), - }, - } - } - - async fn make_fm_index_entry( - file_name: &str, - field_id: i32, - row_range_start: i64, - row_range_end: i64, - first_row_id: u64, - row_count: u64, - ) -> IndexManifestEntry { - let output = VecFileWrite::new(); - let mut writer = FMGlobalIndexWriter::new( - Box::new(output.clone()), - FMWriteOptions { - compression: BlockCompressionType::None, - ..FMWriteOptions::default() - }, - ) - .unwrap(); - for row_id in first_row_id..first_row_id + row_count { - writer.write(Some(b"value"), row_id).await.unwrap(); - } - let result = writer.finish().await.unwrap(); - IndexManifestEntry { - version: 1, - kind: FileKind::Add, - partition: vec![], - bucket: 0, - index_file: IndexFileMeta { - index_type: FM_GLOBAL_INDEX_TYPE.to_string(), - file_name: file_name.to_string(), - file_size: output.to_vec().len() as i64, - row_count: result.row_count as i64, - deletion_vectors_ranges: None, - global_index_meta: Some(crate::spec::GlobalIndexMeta { - row_range_start, - row_range_end, - index_field_id: field_id, - extra_field_ids: None, - source_meta: None, - index_meta: Some(result.index_meta), - }), - }, - } - } - - #[tokio::test] - async fn test_fm_file_set_must_exactly_cover_source_range() { - let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); - let first = make_fm_index_entry("first.fm", 1, 10, 13, 0, 2).await; - let error = match GlobalIndexScanner::create( - &file_io, - "memory:/table", - 1, - i64::MAX, - i64::MAX, - std::slice::from_ref(&first), - &string_schema_fields(), - ) { - Err(error) => error, - Ok(_) => panic!("a partial Java FM file set must fail closed"), - }; - assert!( - matches!(error, Error::DataInvalid { message, .. } if message.contains("cover 2 rows, expected 4")) - ); - - let second = make_fm_index_entry("second.fm", 1, 10, 13, 2, 2).await; - assert!(GlobalIndexScanner::create( - &file_io, - "memory:/table", - 1, - i64::MAX, - i64::MAX, - &[first, second], - &string_schema_fields(), - ) - .unwrap() - .is_some()); - } - - #[test] - fn test_mixed_fm_and_btree_select_compatible_index_family() { - let btree = GlobalIndexEntry { - file_name: "name.btree".to_string(), - index_type: GlobalIndexFileKind::BTree, - file_size: 1, - row_range_start: 0, - row_range_end: 9, - meta: GlobalIndexEntryMeta::Sorted(BTreeIndexMeta::new(None, None, false)), - }; - let fm = GlobalIndexEntry { - file_name: "name.fm".to_string(), - index_type: GlobalIndexFileKind::FM, - file_size: 1, - row_range_start: 0, - row_range_end: 9, - meta: GlobalIndexEntryMeta::FM { - bytes: Vec::new(), - first_row_id: 0, - row_count: 10, - }, - }; - let entries = [btree, fm]; - let data_type = DataType::VarChar(crate::spec::VarCharType::string_type()); - let literal = [Datum::String("needle".to_string())]; - - let contains = [(PredicateOperator::Contains, literal.as_slice(), &data_type)]; - let selected = select_entries_for_predicates(&entries, &contains); - assert_eq!(selected.len(), 1); - assert_eq!(selected[0].index_type, GlobalIndexFileKind::FM); - - let equals = [(PredicateOperator::Eq, literal.as_slice(), &data_type)]; - let selected = select_entries_for_predicates(&entries, &equals); - assert_eq!(selected.len(), 1); - assert_eq!(selected[0].index_type, GlobalIndexFileKind::BTree); - } - - fn int_schema_fields() -> Vec { - vec![DataField::new( - 1, - "id".to_string(), - DataType::Int(crate::spec::IntType::new()), - )] - } - - fn string_schema_fields() -> Vec { - vec![DataField::new( - 1, - "name".to_string(), - DataType::VarChar(crate::spec::VarCharType::string_type()), - )] - } - - async fn evaluate_global_index_fast( - file_io: &FileIO, - table_path: &str, - entries: &[IndexManifestEntry], - predicates: &[Predicate], - fields: &[DataField], - ) -> Result>> { - evaluate_global_index_fast_with_fallback_size( - file_io, - table_path, - entries, - predicates, - fields, - i64::MAX, - i64::MAX, - ) - .await - } - - async fn evaluate_global_index_fast_with_fallback_size( - file_io: &FileIO, - table_path: &str, - entries: &[IndexManifestEntry], - predicates: &[Predicate], - fields: &[DataField], - btree_fallback_scan_max_size: i64, - bitmap_fallback_scan_max_size: i64, - ) -> Result>> { - super::evaluate_global_index(super::GlobalIndexEvaluation { - file_io, - table_path, - index_entries: entries, - predicates, - schema_fields: fields, - search_mode: GlobalIndexSearchMode::Fast, - global_index_thread_num: 32, - btree_fallback_scan_max_size, - bitmap_fallback_scan_max_size, - fm_read_options: FMReadOptions::default(), - next_row_id: None, - data_ranges: &[], - }) - .await - } - - fn two_field_schema_fields() -> Vec { - vec![ - DataField::new( - 1, - "id".to_string(), - DataType::Int(crate::spec::IntType::new()), - ), - DataField::new( - 2, - "value".to_string(), - DataType::Int(crate::spec::IntType::new()), - ), - ] - } - - fn int_eq(column: &str, index: usize, value: i32) -> Predicate { - Predicate::Leaf { - column: column.to_string(), - index, - data_type: DataType::Int(crate::spec::IntType::new()), - op: PredicateOperator::Eq, - literals: vec![Datum::Int(value)], - } - } - - #[test] - fn test_unindexed_ranges_fast_mode_empty() { - let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); - let meta = BTreeIndexMeta::new(None, None, false); - let entries = vec![make_global_index_entry("idx", 1, 0, 49, &meta)]; - let fields = int_schema_fields(); - let scanner = GlobalIndexScanner::create( - &file_io, - "memory:/t", - 32, - i64::MAX, - i64::MAX, - &entries, - &fields, - ) - .expect("create scanner") - .expect("scanner"); - - let ranges = scanner - .unindexed_ranges( - &int_eq("id", 0, 7), - GlobalIndexSearchMode::Fast, - Some(100), - &[RowRange::new(50, 99)], - ) - .unwrap(); - assert!(ranges.is_empty()); - } - - #[test] - fn test_unindexed_ranges_full_uses_snapshot_next_row_id() { - let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); - let meta = BTreeIndexMeta::new(None, None, false); - let entries = vec![make_global_index_entry("idx", 1, 0, 49, &meta)]; - let fields = int_schema_fields(); - let scanner = GlobalIndexScanner::create( - &file_io, - "memory:/t", - 32, - i64::MAX, - i64::MAX, - &entries, - &fields, - ) - .expect("create scanner") - .expect("scanner"); - - let ranges = scanner - .unindexed_ranges( - &int_eq("id", 0, 7), - GlobalIndexSearchMode::Full, - Some(100), - &[], - ) - .unwrap(); - assert_eq!(ranges, vec![RowRange::new(50, 99)]); - } - - #[test] - fn test_unindexed_ranges_detail_uses_data_file_ranges() { - let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); - let meta = BTreeIndexMeta::new(None, None, false); - let entries = vec![make_global_index_entry("idx", 1, 0, 49, &meta)]; - let fields = int_schema_fields(); - let scanner = GlobalIndexScanner::create( - &file_io, - "memory:/t", - 32, - i64::MAX, - i64::MAX, - &entries, - &fields, - ) - .expect("create scanner") - .expect("scanner"); - - let ranges = scanner - .unindexed_ranges( - &int_eq("id", 0, 7), - GlobalIndexSearchMode::Detail, - Some(100), - &[ - RowRange::new(0, 10), - RowRange::new(40, 60), - RowRange::new(80, 90), - ], - ) - .unwrap(); - assert_eq!(ranges, vec![RowRange::new(50, 60), RowRange::new(80, 90)]); - } - - #[test] - fn test_unindexed_ranges_uses_all_predicate_field_coverage() { - let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); - let meta = BTreeIndexMeta::new(None, None, false); - let entries = vec![ - make_global_index_entry("idx_id", 1, 0, 49, &meta), - make_global_index_entry("idx_value", 2, 0, 99, &meta), - ]; - let fields = two_field_schema_fields(); - let scanner = GlobalIndexScanner::create( - &file_io, - "memory:/t", - 32, - i64::MAX, - i64::MAX, - &entries, - &fields, - ) - .expect("create scanner") - .expect("scanner"); - let predicate = Predicate::and(vec![int_eq("id", 0, 7), int_eq("value", 1, 8)]); - - let ranges = scanner - .unindexed_ranges(&predicate, GlobalIndexSearchMode::Full, Some(100), &[]) - .unwrap(); - assert_eq!(ranges, vec![RowRange::new(50, 99)]); - } - - #[test] - fn test_unindexed_ranges_missing_field_coverage_reads_all_data_ranges() { - let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); - let meta = BTreeIndexMeta::new(None, None, false); - let entries = vec![make_global_index_entry("idx_id", 1, 0, 49, &meta)]; - let fields = two_field_schema_fields(); - let scanner = GlobalIndexScanner::create( - &file_io, - "memory:/t", - 32, - i64::MAX, - i64::MAX, - &entries, - &fields, - ) - .expect("create scanner") - .expect("scanner"); - let predicate = Predicate::and(vec![int_eq("id", 0, 7), int_eq("value", 1, 8)]); - - let ranges = scanner - .unindexed_ranges(&predicate, GlobalIndexSearchMode::Full, Some(100), &[]) - .unwrap(); - assert_eq!(ranges, vec![RowRange::new(0, 99)]); - } - - #[test] - fn test_unindexed_ranges_counts_extra_field_coverage() { - let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); - let meta = BTreeIndexMeta::new(None, None, false); - let mut entry = make_global_index_entry("idx_id_value", 1, 0, 99, &meta); - entry - .index_file - .global_index_meta - .as_mut() - .unwrap() - .extra_field_ids = Some(vec![2]); - let fields = two_field_schema_fields(); - let scanner = GlobalIndexScanner::create( - &file_io, - "memory:/t", - 32, - i64::MAX, - i64::MAX, - &[entry], - &fields, - ) - .expect("create scanner") - .expect("scanner"); - - let ranges = scanner - .unindexed_ranges( - &int_eq("value", 1, 8), - GlobalIndexSearchMode::Full, - Some(100), - &[], - ) - .unwrap(); - assert!(ranges.is_empty()); - } - - #[tokio::test] - async fn test_evaluate_extra_field_only_without_composite_reader_falls_back() { - let (file_io, table_path, file_name, _tmp) = - setup_testdata_table("btree_int_100_no_compress.bin"); - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - let mut entry = make_global_index_entry(&file_name, 1, 0, 99, &meta); - entry - .index_file - .global_index_meta - .as_mut() - .unwrap() - .extra_field_ids = Some(vec![2]); - let fields = two_field_schema_fields(); - let predicates = vec![int_eq("value", 1, 50)]; - - let result = - evaluate_global_index_fast(&file_io, &table_path, &[entry], &predicates, &fields) - .await - .unwrap(); - assert!( - result.is_none(), - "extra-field-only predicates must fall back until composite-key btree reads are supported" - ); - } - - #[tokio::test] - async fn test_evaluate_global_index_eq() { - let (file_io, table_path, file_name, tmp) = - setup_testdata_table("btree_int_100_no_compress.bin"); - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - let mut entry = make_global_index_entry(&file_name, 1, 0, 99, &meta); - entry.index_file.file_size = std::fs::metadata(tmp.path().join("index").join(&file_name)) - .unwrap() - .len() as i64; - let entries = vec![entry]; - let fields = int_schema_fields(); - - // key=50 -> row_id=25, offset by row_range_start=0 -> global row_id=25 - let predicates = vec![Predicate::Leaf { - column: "id".to_string(), - index: 0, - data_type: DataType::Int(crate::spec::IntType::new()), - op: PredicateOperator::Eq, - literals: vec![Datum::Int(50)], - }]; - - let result = - evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) - .await - .unwrap(); - let ranges = result.unwrap(); - assert_eq!(ranges, vec![RowRange::new(25, 25)]); - } - - #[tokio::test] - async fn test_evaluate_global_index_uses_known_file_size() { - let (file_io, table_path, file_name, _tmp) = - setup_testdata_table("btree_int_100_no_compress.bin"); - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - let mut entry = make_global_index_entry(&file_name, 1, 0, 99, &meta); - entry.index_file.file_size = 1; - - let error = evaluate_global_index_fast( - &file_io, - &table_path, - &[entry], - &[int_eq("id", 0, 50)], - &int_schema_fields(), - ) - .await - .expect_err("the known file size should be used without a metadata lookup"); - - assert!(matches!( - error, - crate::Error::DataInvalid { message, .. } - if message.contains("Failed to open BTree index file") - )); - } - - #[tokio::test] - async fn test_missing_index_meta_returns_error() { - let (file_io, table_path, file_name, tmp) = - setup_testdata_table("btree_int_100_no_compress.bin"); - let second_file_name = "btree_int_100_no_compress_2.bin"; - std::fs::copy( - tmp.path().join("index").join(&file_name), - tmp.path().join("index").join(second_file_name), - ) - .unwrap(); - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - let valid_entry = make_global_index_entry(&file_name, 1, 0, 99, &meta); - let mut invalid_entry = make_global_index_entry(second_file_name, 1, 100, 199, &meta); - invalid_entry - .index_file - .global_index_meta - .as_mut() - .unwrap() - .index_meta = None; - - let error = evaluate_global_index_fast( - &file_io, - &table_path, - &[valid_entry, invalid_entry], - &[int_eq("id", 0, 50)], - &int_schema_fields(), - ) - .await - .expect_err("missing sorted index metadata must fail the scan"); - - assert!(matches!( - error, - crate::Error::DataInvalid { message, .. } - if message.contains(second_file_name) - )); - } - - #[tokio::test] - async fn test_invalid_index_meta_returns_error() { - let (file_io, table_path, file_name, tmp) = - setup_testdata_table("btree_int_100_no_compress.bin"); - let second_file_name = "btree_int_100_no_compress_2.bin"; - std::fs::copy( - tmp.path().join("index").join(&file_name), - tmp.path().join("index").join(second_file_name), - ) - .unwrap(); - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - let valid_entry = make_global_index_entry(&file_name, 1, 0, 99, &meta); - let mut invalid_entry = make_global_index_entry(second_file_name, 1, 100, 199, &meta); - let mut invalid_meta = vec![0; 9]; - invalid_meta[..4].copy_from_slice(&10i32.to_le_bytes()); - invalid_entry - .index_file - .global_index_meta - .as_mut() - .unwrap() - .index_meta = Some(invalid_meta); - - let error = evaluate_global_index_fast( - &file_io, - &table_path, - &[valid_entry, invalid_entry], - &[int_eq("id", 0, 50)], - &int_schema_fields(), - ) - .await - .expect_err("invalid sorted index metadata must fail the scan"); - - assert!(matches!( - error, - crate::Error::DataInvalid { - message, - source: Some(_), - } if message.contains(second_file_name) - )); - } - - #[tokio::test] - async fn test_missing_global_index_meta_returns_error() { - let (file_io, table_path, file_name, tmp) = - setup_testdata_table("btree_int_100_no_compress.bin"); - let second_file_name = "btree_int_100_no_compress_2.bin"; - std::fs::copy( - tmp.path().join("index").join(&file_name), - tmp.path().join("index").join(second_file_name), - ) - .unwrap(); - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - let valid_entry = make_global_index_entry(&file_name, 1, 0, 99, &meta); - let mut invalid_entry = make_global_index_entry(second_file_name, 1, 100, 199, &meta); - invalid_entry.index_file.global_index_meta = None; - - let error = evaluate_global_index_fast( - &file_io, - &table_path, - &[valid_entry, invalid_entry], - &[int_eq("id", 0, 50)], - &int_schema_fields(), - ) - .await - .expect_err("missing global index metadata must fail the scan"); - - assert!(matches!( - error, - crate::Error::DataInvalid { message, .. } - if message.contains(second_file_name) - )); - } - - #[tokio::test] - async fn test_evaluate_java_bitmap_golden_index_eq_and_null() { - let data_type = DataType::VarChar(crate::spec::VarCharType::string_type()); - let (file_io, table_path, file_name, meta, tmp) = setup_java_bitmap_testdata_table(); - let mut entry = make_global_index_entry_with_type( - BITMAP_GLOBAL_INDEX_TYPE, - &file_name, - 1, - 100, - 109, - &meta, - ); - entry.index_file.file_size = std::fs::metadata(tmp.path().join("index").join(&file_name)) - .unwrap() - .len() as i64; - let entries = vec![entry]; - let fields = string_schema_fields(); - assert_eq!(meta.first_key, Some(b"alpha".to_vec())); - assert_eq!(meta.last_key, Some(b"office".to_vec())); - assert!(meta.has_nulls); - - let eq_predicates = vec![Predicate::Leaf { - column: "name".to_string(), - index: 0, - data_type: data_type.clone(), - op: PredicateOperator::Eq, - literals: vec![Datum::String("k2".to_string())], - }]; - let eq_result = - evaluate_global_index_fast(&file_io, &table_path, &entries, &eq_predicates, &fields) - .await - .unwrap(); - assert_eq!(eq_result.unwrap(), vec![RowRange::new(105, 106)]); - - let null_predicates = vec![Predicate::Leaf { - column: "name".to_string(), - index: 0, - data_type, - op: PredicateOperator::IsNull, - literals: vec![], - }]; - let null_result = - evaluate_global_index_fast(&file_io, &table_path, &entries, &null_predicates, &fields) - .await - .unwrap(); - assert_eq!(null_result.unwrap(), vec![RowRange::new(104, 104)]); - } - - async fn assert_bitmap_int_fixture(file_name: &str) { - let data_type = DataType::Int(crate::spec::IntType::new()); - let (file_io, table_path, file_name, meta, _tmp) = setup_bitmap_testdata_table(file_name); - let entries = vec![make_global_index_entry_with_type( - BITMAP_GLOBAL_INDEX_TYPE, - &file_name, - 1, - 100, - 105, - &meta, - )]; - let fields = int_schema_fields(); - assert_eq!(meta.first_key, Some(le_int_key(-1))); - assert_eq!(meta.last_key, Some(le_int_key(256))); - assert!(meta.has_nulls); - - let cases = [ - ( - PredicateOperator::Eq, - vec![Datum::Int(0)], - vec![RowRange::new(101, 102)], - ), - ( - PredicateOperator::Eq, - vec![Datum::Int(256)], - vec![RowRange::new(104, 104)], - ), - ( - PredicateOperator::In, - vec![Datum::Int(-1), Datum::Int(1), Datum::Int(256)], - vec![RowRange::new(100, 100), RowRange::new(103, 104)], - ), - ( - PredicateOperator::NotEq, - vec![Datum::Int(0)], - vec![RowRange::new(100, 100), RowRange::new(103, 104)], - ), - ( - PredicateOperator::NotIn, - vec![Datum::Int(-1), Datum::Int(1), Datum::Int(256)], - vec![RowRange::new(101, 102)], - ), - ( - PredicateOperator::IsNull, - vec![], - vec![RowRange::new(105, 105)], - ), - ]; - - for (op, literals, expected) in cases { - let predicates = vec![Predicate::Leaf { - column: "id".to_string(), - index: 0, - data_type: data_type.clone(), - op, - literals, - }]; - let result = - evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) - .await - .unwrap() - .unwrap(); - assert_eq!(result, expected, "{file_name}: {op}"); - } - } - - #[tokio::test] - async fn test_evaluate_java_logical_order_bitmap_int_fixture() { - assert_bitmap_int_fixture("bitmap_int_logical_java.index").await; - } - - async fn assert_bitmap_nan_equality_uses_direct_lookup( - data_type: DataType, - nan_literals: [Datum; 3], - zero: Datum, - ) { - let output = VecFileWrite::new(); - let captured = output.clone(); - let mut writer = BitmapGlobalIndexWriter::new( - Box::new(output), - 1, - BlockCompressionType::None, - make_bitmap_key_comparator(&data_type), - ); - for (row_id, literal) in nan_literals.iter().enumerate() { - let key = serialize_bitmap_datum(literal, &data_type); - writer.write(Some(&key), row_id as i64).unwrap(); - } - let zero_key = serialize_bitmap_datum(&zero, &data_type); - writer.write(Some(&zero_key), 3).unwrap(); - let write_result = writer.finish().await.unwrap(); - let bytes = captured.to_vec(); - - let tmp = tempfile::tempdir().unwrap(); - let index_dir = tmp.path().join("index"); - std::fs::create_dir_all(&index_dir).unwrap(); - let file_name = "bitmap-current.index"; - std::fs::write(index_dir.join(file_name), &bytes).unwrap(); - let table_path = format!("file://{}", tmp.path().display()); - let file_io = crate::io::FileIOBuilder::new("file").build().unwrap(); - - let mut entry = make_global_index_entry_with_type( - BITMAP_GLOBAL_INDEX_TYPE, - file_name, - 1, - 100, - 103, - &write_result.meta, - ); - entry.index_file.file_size = bytes.len() as i64; - let entries = vec![entry]; - let fields = vec![DataField::new(1, "id".to_string(), data_type.clone())]; - let cases = [ - (PredicateOperator::Eq, vec![nan_literals[0].clone()]), - ( - PredicateOperator::In, - vec![nan_literals[1].clone(), nan_literals[2].clone()], - ), - ]; - - for (op, literals) in cases { - let predicates = vec![Predicate::Leaf { - column: "id".to_string(), - index: 0, - data_type: data_type.clone(), - op, - literals, - }]; - let result = evaluate_global_index_fast_with_fallback_size( - &file_io, - &table_path, - &entries, - &predicates, - &fields, - i64::MAX, - 0, - ) - .await - .unwrap() - .unwrap(); - assert_eq!(result, vec![RowRange::new(100, 102)], "{data_type:?}: {op}"); - } - } - - #[tokio::test] - async fn test_bitmap_nan_equality_uses_direct_lookup_with_fallback_scan_disabled() { - assert_bitmap_nan_equality_uses_direct_lookup( - DataType::Float(crate::spec::FloatType::new()), - [ - Datum::Float(f32::from_bits(0xffc0_0001)), - Datum::Float(f32::from_bits(0x7fc0_0010)), - Datum::Float(f32::NAN), - ], - Datum::Float(0.0), - ) - .await; - assert_bitmap_nan_equality_uses_direct_lookup( - DataType::Double(crate::spec::DoubleType::new()), - [ - Datum::Double(f64::from_bits(0xfff8_0000_0000_0001)), - Datum::Double(f64::from_bits(0x7ff8_0000_0000_0010)), - Datum::Double(f64::NAN), - ], - Datum::Double(0.0), - ) - .await; - } - - fn legacy_floating_comparator(data_type: &DataType) -> BoxedCmp { - match data_type { - DataType::Float(_) => Box::new(|left, right| { - let left = f32::from_le_bytes(left.try_into().unwrap()); - let right = f32::from_le_bytes(right.try_into().unwrap()); - left.total_cmp(&right) - }), - DataType::Double(_) => Box::new(|left, right| { - let left = f64::from_le_bytes(left.try_into().unwrap()); - let right = f64::from_le_bytes(right.try_into().unwrap()); - left.total_cmp(&right) - }), - _ => unreachable!("legacy floating comparator requires Float or Double"), - } - } - - async fn assert_legacy_floating_btree( - file_name: &str, - data_type: DataType, - nan_keys: Vec>, - nan_literals: Vec, - zero_key: Vec, - zero_literal: Datum, - ) { - let mut rows = nan_keys - .into_iter() - .enumerate() - .map(|(row_id, key)| (key, row_id as i64)) - .collect::>(); - rows.push((zero_key, 3)); - let cmp = legacy_floating_comparator(&data_type); - rows.sort_by(|left, right| cmp(&left.0, &right.0)); - let expected_first_key = rows.first().unwrap().0.clone(); - let expected_last_key = rows.last().unwrap().0.clone(); - - let output = VecFileWrite::new(); - let captured = output.clone(); - let mut writer = - BTreeIndexWriter::with_comparator(Box::new(output), 1, BlockCompressionType::None, cmp); - for (key, row_id) in rows { - writer.write(Some(&key), row_id).await.unwrap(); - } - let write_result = writer.finish().await.unwrap(); - assert_eq!(write_result.meta.first_key, Some(expected_first_key)); - assert_eq!(write_result.meta.last_key, Some(expected_last_key)); - - let tmp = tempfile::tempdir().unwrap(); - let index_dir = tmp.path().join("index"); - std::fs::create_dir_all(&index_dir).unwrap(); - std::fs::write(index_dir.join(file_name), captured.to_vec()).unwrap(); - let table_path = format!("file://{}", tmp.path().display()); - let file_io = crate::io::FileIOBuilder::new("file").build().unwrap(); - let entries = vec![make_global_index_entry( - file_name, - 1, - 100, - 103, - &write_result.meta, - )]; - let fields = vec![DataField::new(1, "id".to_string(), data_type.clone())]; - let cases = [ - ( - PredicateOperator::Eq, - vec![zero_literal.clone()], - vec![RowRange::new(103, 103)], - ), - ( - PredicateOperator::Eq, - vec![nan_literals[0].clone()], - vec![RowRange::new(100, 100)], - ), - ( - PredicateOperator::In, - vec![ - nan_literals[0].clone(), - nan_literals[1].clone(), - zero_literal, - ], - vec![RowRange::new(100, 101), RowRange::new(103, 103)], - ), - ]; - - for (op, literals, expected) in cases { - let predicates = vec![Predicate::Leaf { - column: "id".to_string(), - index: 0, - data_type: data_type.clone(), - op, - literals, - }]; - let result = - evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) - .await - .unwrap() - .unwrap(); - assert_eq!(result, expected, "{file_name}: {op}"); - } - } - - #[tokio::test] - async fn test_evaluate_legacy_float_btree() { - let nan_bits = [0xffc0_0001u32, 0xffc0_0010, 0xffff_1234]; - assert_legacy_floating_btree( - "btree_float_legacy_rust.index", - DataType::Float(crate::spec::FloatType::new()), - nan_bits - .iter() - .map(|bits| bits.to_le_bytes().to_vec()) - .collect(), - nan_bits - .iter() - .map(|bits| Datum::Float(f32::from_bits(*bits))) - .collect(), - 0.0f32.to_le_bytes().to_vec(), - Datum::Float(0.0), - ) - .await; - } - - #[tokio::test] - async fn test_evaluate_legacy_double_btree() { - let nan_bits = [ - 0xfff8_0000_0000_0001u64, - 0xfff8_0000_0000_0010, - 0xffff_1234_5678_9abc, - ]; - assert_legacy_floating_btree( - "btree_double_legacy_rust.index", - DataType::Double(crate::spec::DoubleType::new()), - nan_bits - .iter() - .map(|bits| bits.to_le_bytes().to_vec()) - .collect(), - nan_bits - .iter() - .map(|bits| Datum::Double(f64::from_bits(*bits))) - .collect(), - 0.0f64.to_le_bytes().to_vec(), - Datum::Double(0.0), - ) - .await; - } - - #[tokio::test] - async fn test_evaluate_java_bitmap_golden_index_string_fallback_scan() { - let data_type = DataType::VarChar(crate::spec::VarCharType::string_type()); - let (file_io, table_path, file_name, meta, tmp) = setup_java_bitmap_testdata_table(); - let file_size = std::fs::metadata(tmp.path().join("index").join(&file_name)) - .unwrap() - .len() as i64; - let mut entry = make_global_index_entry_with_type( - BITMAP_GLOBAL_INDEX_TYPE, - &file_name, - 1, - 100, - 109, - &meta, - ); - entry.index_file.file_size = file_size; - let entries = vec![entry]; - let fields = string_schema_fields(); - - let ends_with_predicates = vec![Predicate::Leaf { - column: "name".to_string(), - index: 0, - data_type: data_type.clone(), - op: PredicateOperator::EndsWith, - literals: vec![Datum::String("ta".to_string())], - }]; - let ends_with_result = evaluate_global_index_fast( - &file_io, - &table_path, - &entries, - &ends_with_predicates, - &fields, - ) - .await - .unwrap(); - assert_eq!( - ends_with_result.unwrap(), - vec![RowRange::new(101, 101), RowRange::new(103, 103)] - ); - - let contains_predicates = vec![Predicate::Leaf { - column: "name".to_string(), - index: 0, - data_type: data_type.clone(), - op: PredicateOperator::Contains, - literals: vec![Datum::String("ph".to_string())], - }]; - let contains_result = evaluate_global_index_fast( - &file_io, - &table_path, - &entries, - &contains_predicates, - &fields, - ) - .await - .unwrap(); - assert_eq!( - contains_result.unwrap(), - vec![RowRange::new(100, 100), RowRange::new(102, 102)] - ); - - let like_predicates = vec![Predicate::Leaf { - column: "name".to_string(), - index: 0, - data_type: data_type.clone(), - op: PredicateOperator::Like, - literals: vec![Datum::String("%ha%".to_string())], - }]; - let like_result = - evaluate_global_index_fast(&file_io, &table_path, &entries, &like_predicates, &fields) - .await - .unwrap(); - assert_eq!( - like_result.unwrap(), - vec![RowRange::new(100, 100), RowRange::new(102, 102)] - ); - - let less_than_predicates = vec![Predicate::Leaf { - column: "name".to_string(), - index: 0, - data_type: data_type.clone(), - op: PredicateOperator::Lt, - literals: vec![Datum::String("delta".to_string())], - }]; - let less_than_result = evaluate_global_index_fast( - &file_io, - &table_path, - &entries, - &less_than_predicates, - &fields, - ) - .await - .unwrap(); - assert_eq!(less_than_result.unwrap(), vec![RowRange::new(100, 102)]); - - let mut over_limit_entry = make_global_index_entry_with_type( - BITMAP_GLOBAL_INDEX_TYPE, - &file_name, - 1, - 100, - 109, - &meta, - ); - over_limit_entry.index_file.file_size = file_size; - let over_limit_entries = vec![over_limit_entry]; - let over_limit_less_than = evaluate_global_index_fast_with_fallback_size( - &file_io, - &table_path, - &over_limit_entries, - &less_than_predicates, - &fields, - i64::MAX, - file_size - 1, - ) - .await - .unwrap(); - assert!( - over_limit_less_than.is_none(), - "range predicates require fallback dictionary scans and should be unsupported over budget" - ); - - let no_match_contains = vec![Predicate::Leaf { - column: "name".to_string(), - index: 0, - data_type: data_type.clone(), - op: PredicateOperator::Contains, - literals: vec![Datum::String("zz".to_string())], - }]; - let over_limit_result = evaluate_global_index_fast_with_fallback_size( - &file_io, - &table_path, - &over_limit_entries, - &no_match_contains, - &fields, - i64::MAX, - file_size - 1, - ) - .await - .unwrap(); - assert!( - over_limit_result.is_none(), - "fallback scans over budget should be unsupported instead of returning full coverage" - ); - - let direct_with_over_limit_fallback = vec![Predicate::and(vec![ - Predicate::Leaf { - column: "name".to_string(), - index: 0, - data_type: data_type.clone(), - op: PredicateOperator::Eq, - literals: vec![Datum::String("k2".to_string())], - }, - Predicate::Leaf { - column: "name".to_string(), - index: 0, - data_type, - op: PredicateOperator::Contains, - literals: vec![Datum::String("zz".to_string())], - }, - ])]; - let direct_result = evaluate_global_index_fast_with_fallback_size( - &file_io, - &table_path, - &over_limit_entries, - &direct_with_over_limit_fallback, - &fields, - i64::MAX, - file_size - 1, - ) - .await - .unwrap(); - assert_eq!(direct_result.unwrap(), vec![RowRange::new(105, 106)]); - } - - #[tokio::test] - async fn test_btree_fallback_scan_over_limit_is_unsupported() { - let (file_io, table_path, file_name, tmp) = - setup_testdata_table("btree_varchar_100_no_compress.bin"); - let meta = BTreeIndexMeta::new(Some(b"a".to_vec()), Some(b"yyyy".to_vec()), false); - let fields = string_schema_fields(); - let data_type = DataType::VarChar(crate::spec::VarCharType::string_type()); - let predicates = vec![Predicate::Leaf { - column: "name".to_string(), - index: 0, - data_type, - op: PredicateOperator::Contains, - literals: vec![Datum::String("not-present".to_string())], - }]; - - let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; - let exact_result = evaluate_global_index_fast_with_fallback_size( - &file_io, - &table_path, - &entries, - &predicates, - &fields, - i64::MAX, - i64::MAX, - ) - .await - .unwrap(); - assert_eq!(exact_result.unwrap(), Vec::::new()); - - let mut over_limit_entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; - over_limit_entries[0].index_file.file_size = 2; - let over_limit_result = evaluate_global_index_fast_with_fallback_size( - &file_io, - &table_path, - &over_limit_entries, - &predicates, - &fields, - 1, - i64::MAX, - ) - .await - .unwrap(); - assert!( - over_limit_result.is_none(), - "fallback scans over budget should be unsupported instead of returning full coverage" - ); - - let second_file_name = "btree_varchar_100_no_compress_2.bin"; - std::fs::copy( - tmp.path().join("index").join(&file_name), - tmp.path().join("index").join(second_file_name), - ) - .unwrap(); - let mut first = make_global_index_entry(&file_name, 1, 0, 99, &meta); - first.index_file.file_size = 1; - let mut second = make_global_index_entry(second_file_name, 1, 100, 199, &meta); - second.index_file.file_size = 1; - let total_over_limit_result = evaluate_global_index_fast_with_fallback_size( - &file_io, - &table_path, - &[first, second], - &predicates, - &fields, - 1, - i64::MAX, - ) - .await - .unwrap(); - assert!( - total_over_limit_result.is_none(), - "fallback budget should use selected files' total size, not per-file size" - ); - } - - #[tokio::test] - async fn test_fallback_scan_over_limit_with_mixed_index_kinds_is_unsupported() { - let (file_io, table_path, file_name, _tmp) = - setup_testdata_table("btree_varchar_100_no_compress.bin"); - let btree_meta = BTreeIndexMeta::new(Some(b"a".to_vec()), Some(b"yyyy".to_vec()), false); - let bitmap_meta = BTreeIndexMeta::new(Some(b"m".to_vec()), Some(b"z".to_vec()), false); - let fields = string_schema_fields(); - let predicates = vec![Predicate::Leaf { - column: "name".to_string(), - index: 0, - data_type: DataType::VarChar(crate::spec::VarCharType::string_type()), - op: PredicateOperator::Lt, - literals: vec![Datum::String("delta".to_string())], - }]; - - let mut btree = make_global_index_entry_with_type( - BTREE_GLOBAL_INDEX_TYPE, - &file_name, - 1, - 0, - 99, - &btree_meta, - ); - btree.index_file.file_size = 2; - let mut bitmap = make_global_index_entry_with_type( - BITMAP_GLOBAL_INDEX_TYPE, - "bitmap-no-match.index", - 1, - 100, - 199, - &bitmap_meta, - ); - bitmap.index_file.file_size = 1; - - let result = evaluate_global_index_fast_with_fallback_size( - &file_io, - &table_path, - &[btree, bitmap], - &predicates, - &fields, - 1, - i64::MAX, - ) - .await - .unwrap(); - assert!( - result.is_none(), - "an over-budget selected BTree file must stay unsupported even if bitmap files are pruned by metadata" - ); - } - - #[tokio::test] - async fn test_fallback_preflight_happens_before_shard_io() { - let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); - let table_path = "memory:/missing-index-files"; - let meta = BTreeIndexMeta::new(Some(b"a".to_vec()), Some(b"z".to_vec()), false); - let fields = string_schema_fields(); - let predicates = vec![Predicate::Leaf { - column: "name".to_string(), - index: 0, - data_type: DataType::VarChar(crate::spec::VarCharType::string_type()), - op: PredicateOperator::Contains, - literals: vec![Datum::String("middle".to_string())], - }]; - - let mut btree = make_global_index_entry_with_type( - BTREE_GLOBAL_INDEX_TYPE, - "missing-btree.index", - 1, - 0, - 99, - &meta, - ); - btree.index_file.file_size = 1; - let mut bitmap = make_global_index_entry_with_type( - BITMAP_GLOBAL_INDEX_TYPE, - "missing-bitmap.index", - 1, - 100, - 199, - &meta, - ); - bitmap.index_file.file_size = 1; - - let result = evaluate_global_index_fast_with_fallback_size( - &file_io, - table_path, - &[btree, bitmap], - &predicates, - &fields, - 1, - 0, - ) - .await - .expect("fallback must be decided before opening an earlier shard"); - - assert!(result.is_none()); - } - - #[tokio::test] - async fn test_evaluate_global_index_full_mode_includes_unindexed_tail() { - let (file_io, table_path, file_name, _tmp) = - setup_testdata_table("btree_int_100_no_compress.bin"); - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; - let fields = int_schema_fields(); - let predicates = vec![int_eq("id", 0, 50)]; - - let result = super::evaluate_global_index(super::GlobalIndexEvaluation { - file_io: &file_io, - table_path: &table_path, - index_entries: &entries, - predicates: &predicates, - schema_fields: &fields, - search_mode: GlobalIndexSearchMode::Full, - global_index_thread_num: 32, - btree_fallback_scan_max_size: i64::MAX, - bitmap_fallback_scan_max_size: i64::MAX, - fm_read_options: FMReadOptions::default(), - next_row_id: Some(150), - data_ranges: &[], - }) - .await - .unwrap(); - - assert_eq!( - result.unwrap(), - vec![RowRange::new(25, 25), RowRange::new(100, 149)] - ); - } - - #[tokio::test] - async fn test_evaluate_global_index_and_uses_evaluated_field_coverage_for_raw_fallback() { - let src = format!( - "{}/testdata/btree/btree_int_100_no_compress.bin", - env!("CARGO_MANIFEST_DIR") - ); - let tmp = tempfile::tempdir().unwrap(); - let index_dir = tmp.path().join("index"); - std::fs::create_dir_all(&index_dir).unwrap(); - std::fs::copy(&src, index_dir.join("index_part1.bin")).unwrap(); - std::fs::copy(&src, index_dir.join("index_part2.bin")).unwrap(); - - let table_path = format!("file://{}", tmp.path().display()); - let file_io = crate::io::FileIOBuilder::new("file").build().unwrap(); - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - - let mut first = make_global_index_entry("index_part1.bin", 1, 0, 49, &meta); - first - .index_file - .global_index_meta - .as_mut() - .unwrap() - .extra_field_ids = Some(vec![2]); - let second = make_global_index_entry("index_part2.bin", 1, 50, 99, &meta); - let entries = vec![first, second]; - let fields = two_field_schema_fields(); - - let predicates = vec![Predicate::and(vec![ - int_eq("id", 0, 50), - int_eq("value", 1, 8), - ])]; - let result = super::evaluate_global_index(super::GlobalIndexEvaluation { - file_io: &file_io, - table_path: &table_path, - index_entries: &entries, - predicates: &predicates, - schema_fields: &fields, - search_mode: GlobalIndexSearchMode::Full, - global_index_thread_num: 32, - btree_fallback_scan_max_size: i64::MAX, - bitmap_fallback_scan_max_size: i64::MAX, - fm_read_options: FMReadOptions::default(), - next_row_id: Some(100), - data_ranges: &[], - }) - .await - .unwrap(); - - assert_eq!( - result.unwrap(), - vec![RowRange::new(25, 25), RowRange::new(75, 75)], - "raw fallback should use only the id field that was actually evaluated; \ - the unevaluated extra field must not widen or narrow fallback coverage" - ); - } - - #[tokio::test] - async fn test_evaluate_global_index_detail_mode_uses_data_ranges() { - let (file_io, table_path, file_name, _tmp) = - setup_testdata_table("btree_int_100_no_compress.bin"); - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; - let fields = int_schema_fields(); - let predicates = vec![int_eq("id", 0, 50)]; - - let data_ranges = [RowRange::new(90, 120), RowRange::new(140, 145)]; - let result = super::evaluate_global_index(super::GlobalIndexEvaluation { - file_io: &file_io, - table_path: &table_path, - index_entries: &entries, - predicates: &predicates, - schema_fields: &fields, - search_mode: GlobalIndexSearchMode::Detail, - global_index_thread_num: 32, - btree_fallback_scan_max_size: i64::MAX, - bitmap_fallback_scan_max_size: i64::MAX, - fm_read_options: FMReadOptions::default(), - next_row_id: Some(150), - data_ranges: &data_ranges, - }) - .await - .unwrap(); - - assert_eq!( - result.unwrap(), - vec![ - RowRange::new(25, 25), - RowRange::new(100, 120), - RowRange::new(140, 145), - ] - ); - } - - #[tokio::test] - async fn test_evaluate_global_index_range() { - let (file_io, table_path, file_name, _tmp) = - setup_testdata_table("btree_int_100_no_compress.bin"); - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; - let fields = int_schema_fields(); - - // keys 10..=20 -> keys 10,12,14,16,18,20 -> row_ids 5,6,7,8,9,10 - let predicates = vec![ - Predicate::Leaf { - column: "id".to_string(), - index: 0, - data_type: DataType::Int(crate::spec::IntType::new()), - op: PredicateOperator::GtEq, - literals: vec![Datum::Int(10)], - }, - Predicate::Leaf { - column: "id".to_string(), - index: 0, - data_type: DataType::Int(crate::spec::IntType::new()), - op: PredicateOperator::LtEq, - literals: vec![Datum::Int(20)], - }, - ]; - - let result = - evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) - .await - .unwrap(); - let ranges = result.unwrap(); - assert_eq!(ranges, vec![RowRange::new(5, 10)]); - - let mut over_limit_entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; - over_limit_entries[0].index_file.file_size = 2; - let over_limit_result = evaluate_global_index_fast_with_fallback_size( - &file_io, - &table_path, - &over_limit_entries, - &predicates, - &fields, - 1, - i64::MAX, - ) - .await - .unwrap(); - assert!( - over_limit_result.is_none(), - "between/range predicates require fallback scans and should be unsupported over budget" - ); - } - - #[tokio::test] - async fn test_evaluate_global_index_in() { - let (file_io, table_path, file_name, _tmp) = - setup_testdata_table("btree_int_100_no_compress.bin"); - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; - let fields = int_schema_fields(); - - // IN(0, 50, 198) -> row_ids 0, 25, 99 - let predicates = vec![Predicate::Leaf { - column: "id".to_string(), - index: 0, - data_type: DataType::Int(crate::spec::IntType::new()), - op: PredicateOperator::In, - literals: vec![Datum::Int(0), Datum::Int(50), Datum::Int(198)], - }]; - - let result = - evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) - .await - .unwrap(); - let ranges = result.unwrap(); - assert_eq!( - ranges, - vec![ - RowRange::new(0, 0), - RowRange::new(25, 25), - RowRange::new(99, 99) - ] - ); - } - - #[tokio::test] - async fn test_evaluate_global_index_no_match() { - let (file_io, table_path, file_name, _tmp) = - setup_testdata_table("btree_int_100_no_compress.bin"); - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; - let fields = int_schema_fields(); - - // key=999 doesn't exist - let predicates = vec![Predicate::Leaf { - column: "id".to_string(), - index: 0, - data_type: DataType::Int(crate::spec::IntType::new()), - op: PredicateOperator::Eq, - literals: vec![Datum::Int(999)], - }]; - - let result = - evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) - .await - .unwrap(); - let ranges = result.unwrap(); - assert!(ranges.is_empty()); - } - - #[tokio::test] - async fn test_evaluate_global_index_with_row_offset() { - let (file_io, table_path, file_name, _tmp) = - setup_testdata_table("btree_int_100_no_compress.bin"); - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - // row_range_start=1000 simulates an offset - let entries = vec![make_global_index_entry(&file_name, 1, 1000, 1099, &meta)]; - let fields = int_schema_fields(); - - // key=50 -> local row_id=25, offset -> global row_id=1025 - let predicates = vec![Predicate::Leaf { - column: "id".to_string(), - index: 0, - data_type: DataType::Int(crate::spec::IntType::new()), - op: PredicateOperator::Eq, - literals: vec![Datum::Int(50)], - }]; - - let result = - evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) - .await - .unwrap(); - let ranges = result.unwrap(); - assert_eq!(ranges, vec![RowRange::new(1025, 1025)]); - } - - #[tokio::test] - async fn test_evaluate_global_index_unknown_column() { - let (file_io, table_path, file_name, _tmp) = - setup_testdata_table("btree_int_100_no_compress.bin"); - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; - let fields = int_schema_fields(); - - // Column "unknown" not in schema -> None (can't evaluate) - let predicates = vec![Predicate::Leaf { - column: "unknown".to_string(), - index: 0, - data_type: DataType::Int(crate::spec::IntType::new()), - op: PredicateOperator::Eq, - literals: vec![Datum::Int(50)], - }]; - - let result = - evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) - .await - .unwrap(); - assert!(result.is_none()); - } - - #[tokio::test] - async fn test_evaluate_global_index_multi_field_and() { - // Two fields, each with its own btree index file (same data, different field_id). - // btree_int_100_no_compress.bin: keys 0,2,4,...,198 -> row_ids 0,1,...,99 - let src = format!( - "{}/testdata/btree/btree_int_100_no_compress.bin", - env!("CARGO_MANIFEST_DIR") - ); - let tmp = tempfile::tempdir().unwrap(); - let index_dir = tmp.path().join("index"); - std::fs::create_dir_all(&index_dir).unwrap(); - std::fs::copy(&src, index_dir.join("index_field1.bin")).unwrap(); - std::fs::copy(&src, index_dir.join("index_field2.bin")).unwrap(); - - let table_path = format!("file://{}", tmp.path().display()); - let file_io = crate::io::FileIOBuilder::new("file").build().unwrap(); - - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - - let fields = vec![ - DataField::new( - 1, - "id".to_string(), - DataType::Int(crate::spec::IntType::new()), - ), - DataField::new( - 2, - "value".to_string(), - DataType::Int(crate::spec::IntType::new()), - ), - ]; - - let entries = vec![ - make_global_index_entry("index_field1.bin", 1, 0, 99, &meta), - make_global_index_entry("index_field2.bin", 2, 0, 99, &meta), - ]; - - // id >= 40 AND id <= 60 → keys 40,42,...,60 → row_ids 20..30 - // value >= 44 AND value <= 52 → keys 44,46,48,50,52 → row_ids 22..26 - // AND → intersect [20..30] and [22..26] = [22..26] - let predicates = vec![ - Predicate::Leaf { - column: "id".to_string(), - index: 0, - data_type: DataType::Int(crate::spec::IntType::new()), - op: PredicateOperator::GtEq, - literals: vec![Datum::Int(40)], - }, - Predicate::Leaf { - column: "id".to_string(), - index: 0, - data_type: DataType::Int(crate::spec::IntType::new()), - op: PredicateOperator::LtEq, - literals: vec![Datum::Int(60)], - }, - Predicate::Leaf { - column: "value".to_string(), - index: 1, - data_type: DataType::Int(crate::spec::IntType::new()), - op: PredicateOperator::GtEq, - literals: vec![Datum::Int(44)], - }, - Predicate::Leaf { - column: "value".to_string(), - index: 1, - data_type: DataType::Int(crate::spec::IntType::new()), - op: PredicateOperator::LtEq, - literals: vec![Datum::Int(52)], - }, - ]; - - let result = - evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) - .await - .unwrap(); - let ranges = result.unwrap(); - assert_eq!(ranges, vec![RowRange::new(22, 26)]); - } - - #[tokio::test] - async fn test_multi_field_and_shares_query_concurrency_budget() { - let src = format!( - "{}/testdata/btree/btree_int_100_no_compress.bin", - env!("CARGO_MANIFEST_DIR") - ); - let tmp = tempfile::tempdir().unwrap(); - let index_dir = tmp.path().join("index"); - std::fs::create_dir_all(&index_dir).unwrap(); - let file_names: Vec<_> = (1..=4) - .map(|field_id| { - let file_name = format!("index_field{field_id}.bin"); - std::fs::copy(&src, index_dir.join(&file_name)).unwrap(); - file_name - }) - .collect(); - - let table_path = format!("file://{}", tmp.path().display()); - let file_io = crate::io::FileIOBuilder::new("file").build().unwrap(); - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - let fields: Vec<_> = (0..4) - .map(|index| { - let field_id = index + 1; - DataField::new( - field_id, - format!("field{field_id}"), - DataType::Int(crate::spec::IntType::new()), - ) - }) - .collect(); - let entries: Vec<_> = file_names - .iter() - .enumerate() - .map(|(index, file_name)| { - make_global_index_entry(file_name, index as i32 + 1, 0, 99, &meta) - }) - .collect(); - let predicate = Predicate::and( - (0..4) - .map(|index| int_eq(&format!("field{}", index + 1), index, 50)) - .collect(), - ); - - for (thread_num, expected_peak) in [(1, 1), (2, 2)] { - let mut scanner = GlobalIndexScanner::create( - &file_io, - &table_path, - thread_num, - i64::MAX, - i64::MAX, - &entries, - &fields, - ) - .unwrap() - .unwrap(); - let probe = Arc::new(QueryIoProbe::default()); - scanner.query_io_probe = Some(Arc::clone(&probe)); - - let result = scanner.evaluate(&predicate).await.unwrap().unwrap(); - - assert_eq!(result.row_ranges, vec![RowRange::new(25, 25)]); - assert_eq!(result.evaluated_field_ids, HashSet::from([1, 2, 3, 4])); - assert_eq!(probe.peak(), expected_peak); - } - } - - /// Regression for the Between+remaining bug in `evaluate_leaf`. When a - /// native `Between` leaf is paired with another conjunct (e.g. `id >= 0`), - /// and the file's b-tree key range falls **outside** the Between range - /// but is still matched by the remaining predicate, the whole AND must - /// produce zero rows. Before the fix, `file_result` was initialized from - /// the remaining predicate's bitmap and the Between conjunct was silently - /// dropped — the test would observe the file's full row id set instead of - /// the empty set. - #[tokio::test] - async fn test_between_unmatched_file_drops_remaining_match() { - let (file_io, table_path, file_name, _tmp) = - setup_testdata_table("btree_int_100_no_compress.bin"); - // File covers keys [0, 198] (row_ids 0..99). Pick a Between range - // entirely below 0 so `may_match_between` is false, and a `>= 0` - // conjunct that would otherwise scoop up every row in the file. - let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); - let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; - let fields = int_schema_fields(); - - let predicates = vec![Predicate::and(vec![ - Predicate::Leaf { - column: "id".to_string(), - index: 0, - data_type: DataType::Int(crate::spec::IntType::new()), - op: PredicateOperator::Between, - literals: vec![Datum::Int(-100), Datum::Int(-50)], - }, - Predicate::Leaf { - column: "id".to_string(), - index: 0, - data_type: DataType::Int(crate::spec::IntType::new()), - op: PredicateOperator::GtEq, - literals: vec![Datum::Int(0)], - }, - ])]; - - let result = - evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) - .await - .unwrap(); - let ranges = result.unwrap(); - assert!( - ranges.is_empty(), - "Between(-100..-50) AND id>=0 must produce zero rows on a file \ - whose key range is [0, 198] — got {ranges:?}" - ); - } -} +mod tests; diff --git a/crates/paimon/src/table/global_index_scanner/deletion_vectors.rs b/crates/paimon/src/table/global_index_scanner/deletion_vectors.rs new file mode 100644 index 000000000..3fc85a835 --- /dev/null +++ b/crates/paimon/src/table/global_index_scanner/deletion_vectors.rs @@ -0,0 +1,122 @@ +// 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. + +//! Resolution of deletion-vector positions into global row ranges. + +use super::{DELETION_VECTORS_INDEX_TYPE, INDEX_DIR}; +use crate::deletion_vector::DeletionVectorFactory; +use crate::spec::{FileKind, IndexManifestEntry}; +use crate::table::{merge_row_ranges, DeletionFile, RowRange, Table}; +use crate::Result; +use std::collections::HashMap; + +/// Resolve live deletion-vector index entries into global row-id ranges. +/// +/// Data-evolution DV entries are keyed by the normal anchor data file. The DV +/// bitmap positions are local to that anchor file's `first_row_id`, so this +/// helper joins index metadata with live data-file metadata before converting +/// deleted positions to global row IDs. +pub(crate) async fn deleted_row_ranges_for_data_evolution_dvs( + table: &Table, + index_entries: &[IndexManifestEntry], +) -> Result> { + if !index_entries.iter().any(|entry| { + entry.kind == FileKind::Add && entry.index_file.index_type == DELETION_VECTORS_INDEX_TYPE + }) { + return Ok(Vec::new()); + } + + let plan = table + .new_read_builder() + .new_scan() + .with_scan_all_files() + .plan() + .await?; + + let mut first_row_ids: HashMap<(Vec, i32, String), i64> = HashMap::new(); + for split in plan.splits() { + let partition = split.partition().to_serialized_bytes(); + let bucket = split.bucket(); + for file in split.data_files() { + if let Some(first_row_id) = file.first_row_id { + first_row_ids.insert( + (partition.clone(), bucket, file.file_name.clone()), + first_row_id, + ); + } + } + } + + let mut ranges = Vec::new(); + let table_path = table.location().trim_end_matches('/'); + for entry in index_entries { + if entry.kind != FileKind::Add || entry.index_file.index_type != DELETION_VECTORS_INDEX_TYPE + { + continue; + } + let Some(dv_ranges) = entry.index_file.deletion_vectors_ranges.as_ref() else { + continue; + }; + let index_path = format!("{table_path}/{INDEX_DIR}/{}", entry.index_file.file_name); + for (data_file_name, meta) in dv_ranges { + let key = ( + entry.partition.clone(), + entry.bucket, + data_file_name.clone(), + ); + let first_row_id = first_row_ids.get(&key).copied().ok_or_else(|| { + crate::Error::DataInvalid { + message: format!( + "Deletion vector references data file '{}' but no live row-tracked file was found", + data_file_name + ), + source: None, + } + })?; + let deletion_file = DeletionFile::new( + index_path.clone(), + meta.offset as i64, + meta.length as i64, + meta.cardinality, + ); + let deletion_vector = + DeletionVectorFactory::read(table.file_io(), &deletion_file).await?; + for deleted in deletion_vector.iter() { + let deleted = i64::try_from(deleted).map_err(|_| crate::Error::DataInvalid { + message: format!( + "Deleted position {deleted} for data file '{}' exceeds i64::MAX", + data_file_name + ), + source: None, + })?; + let row_id = + first_row_id + .checked_add(deleted) + .ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "Deleted row id overflows i64 for data file '{}'", + data_file_name + ), + source: None, + })?; + ranges.push(RowRange::new(row_id, row_id)); + } + } + } + + Ok(merge_row_ranges(ranges)) +} diff --git a/crates/paimon/src/table/global_index_scanner/entry.rs b/crates/paimon/src/table/global_index_scanner/entry.rs new file mode 100644 index 000000000..dd32f4f8a --- /dev/null +++ b/crates/paimon/src/table/global_index_scanner/entry.rs @@ -0,0 +1,212 @@ +// 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. + +//! Manifest entry parsing and file-level pruning metadata. + +use crate::btree::BTreeIndexMeta; +use crate::spec::{DataType, PredicateOperator}; +use crate::table::bitmap_global_index_format::is_bitmap_floating_residual_sensitive_op; +use crate::{Error, Result}; +use std::cmp::Ordering; +use std::collections::HashMap; + +/// A resolved global index entry with parsed metadata. +pub(super) struct GlobalIndexEntry { + pub(super) file_name: String, + pub(super) index_type: GlobalIndexFileKind, + pub(super) file_size: i64, + pub(super) row_range_start: i64, + pub(super) row_range_end: i64, + pub(super) meta: GlobalIndexEntryMeta, +} + +pub(super) fn sorted_entry_meta(entry: &GlobalIndexEntry) -> &BTreeIndexMeta { + match &entry.meta { + GlobalIndexEntryMeta::Sorted(meta) => meta, + GlobalIndexEntryMeta::FM { .. } => unreachable!("FM entries do not have sorted metadata"), + } +} + +pub(super) enum GlobalIndexEntryMeta { + Sorted(BTreeIndexMeta), + FM { + bytes: Vec, + first_row_id: u64, + row_count: u64, + }, +} + +struct FMFileRowRange<'a> { + file_name: &'a str, + first_row_id: u64, + row_count: u64, +} + +pub(super) fn validate_fm_file_sets( + entries_by_field: &HashMap>, +) -> Result<()> { + for (field_id, entries) in entries_by_field { + let mut groups: HashMap<(i64, i64), Vec>> = HashMap::new(); + for entry in entries { + let GlobalIndexEntryMeta::FM { + first_row_id, + row_count, + .. + } = &entry.meta + else { + continue; + }; + groups + .entry((entry.row_range_start, entry.row_range_end)) + .or_default() + .push(FMFileRowRange { + file_name: &entry.file_name, + first_row_id: *first_row_id, + row_count: *row_count, + }); + } + + for ((range_start, range_end), mut files) in groups { + let expected_row_count = range_end + .checked_sub(range_start) + .and_then(|count| count.checked_add(1)) + .and_then(|count| u64::try_from(count).ok()) + .ok_or_else(|| Error::DataInvalid { + message: format!( + "Invalid FM global index source row range [{range_start}, {range_end}] for field {field_id}" + ), + source: None, + })?; + files.sort_unstable_by_key(|file| file.first_row_id); + let mut expected_first_row_id = 0u64; + for file in files { + let FMFileRowRange { + file_name, + first_row_id, + row_count, + } = file; + if row_count == 0 || first_row_id != expected_first_row_id { + return Err(Error::DataInvalid { + message: format!( + "FM global index files do not exactly cover source row range [{range_start}, {range_end}] for field {field_id}: expected relative row {expected_first_row_id}, file '{file_name}' starts at {first_row_id}" + ), + source: None, + }); + } + expected_first_row_id = + first_row_id + .checked_add(row_count) + .ok_or_else(|| Error::DataInvalid { + message: format!( + "FM global index row range overflows for file '{file_name}'" + ), + source: None, + })?; + if expected_first_row_id > expected_row_count { + return Err(Error::DataInvalid { + message: format!( + "FM global index file '{file_name}' extends beyond source row range [{range_start}, {range_end}]" + ), + source: None, + }); + } + } + if expected_first_row_id != expected_row_count { + return Err(Error::DataInvalid { + message: format!( + "FM global index files cover {expected_first_row_id} rows, expected {expected_row_count} for source row range [{range_start}, {range_end}] and field {field_id}" + ), + source: None, + }); + } + } + } + Ok(()) +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum GlobalIndexFileKind { + BTree, + Bitmap, + Multivalue, + FM, +} + +pub(super) fn is_floating_point(data_type: &DataType) -> bool { + matches!(data_type, DataType::Float(_) | DataType::Double(_)) +} + +pub(super) fn bitmap_meta_may_match( + meta: &BTreeIndexMeta, + op: PredicateOperator, + data_type: &DataType, + serialized_literals: &[Vec], + cmp: &dyn Fn(&[u8], &[u8]) -> Ordering, +) -> bool { + if is_floating_point(data_type) && is_bitmap_floating_residual_sensitive_op(op) { + !meta.only_nulls() + } else { + meta.may_match(op, serialized_literals, cmp) + } +} + +pub(super) fn bitmap_meta_may_match_between( + meta: &BTreeIndexMeta, + data_type: &DataType, + from_key: &[u8], + to_key: &[u8], + cmp: &dyn Fn(&[u8], &[u8]) -> Ordering, +) -> bool { + if is_floating_point(data_type) + && is_bitmap_floating_residual_sensitive_op(PredicateOperator::Between) + { + !meta.only_nulls() + } else { + meta.may_match_between(from_key, to_key, cmp) + } +} + +pub(super) fn multivalue_meta_may_match( + meta: &BTreeIndexMeta, + op: PredicateOperator, + serialized_literals: &[Vec], + cmp: &dyn Fn(&[u8], &[u8]) -> Ordering, +) -> bool { + match op { + PredicateOperator::ArrayContains => { + meta.may_match(PredicateOperator::Eq, serialized_literals, cmp) + } + PredicateOperator::ArraysOverlap => { + meta.may_match(PredicateOperator::In, serialized_literals, cmp) + } + PredicateOperator::ArrayContainsAll => serialized_literals.iter().all(|literal| { + meta.may_match(PredicateOperator::Eq, std::slice::from_ref(literal), cmp) + }), + _ => false, + } +} + +impl GlobalIndexFileKind { + pub(super) fn name(self) -> &'static str { + match self { + Self::BTree => "BTree", + Self::Bitmap => "bitmap", + Self::Multivalue => "multivalue", + Self::FM => "FM", + } + } +} diff --git a/crates/paimon/src/table/global_index_scanner/evaluator.rs b/crates/paimon/src/table/global_index_scanner/evaluator.rs new file mode 100644 index 000000000..9b2daea1a --- /dev/null +++ b/crates/paimon/src/table/global_index_scanner/evaluator.rs @@ -0,0 +1,609 @@ +// 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. + +//! Predicate orchestration and index-query planning. + +use super::entry::{ + bitmap_meta_may_match, bitmap_meta_may_match_between, multivalue_meta_may_match, + sorted_entry_meta, GlobalIndexEntry, GlobalIndexFileKind, +}; +use super::predicates::{ + entries_support_predicate, is_multivalue_predicate, is_sorted_global_index_supported_op, + select_entries_for_predicates, +}; +use super::query_plan::{ + fallback_plan_evaluates_entry, requires_fallback_scan, EntryQueryPlan, FallbackScanPlan, +}; +#[cfg(test)] +use super::row_ranges::unindexed_ranges_for_coverage; +use super::row_ranges::{bitmap_to_ranges, intersect_sorted_ranges}; +use super::GlobalIndexScanner; +use crate::btree::query::extract_between; +use crate::btree::{make_key_comparator, serialize_datum}; +#[cfg(test)] +use crate::spec::GlobalIndexSearchMode; +use crate::spec::{DataType, Datum, Predicate, PredicateOperator}; +use crate::table::bitmap_global_index_format::{ + make_bitmap_key_comparator, serialize_bitmap_datum, +}; +use crate::table::RowRange; +use crate::{Error, Result}; +use futures::{StreamExt, TryStreamExt}; +use roaring::RoaringTreemap; +use std::collections::HashSet; +use std::future::Future; + +type EvaluateFuture<'a> = std::pin::Pin< + Box>> + Send + 'a>, +>; + +type PredicateTuple<'a> = (PredicateOperator, &'a [Datum], &'a DataType); + +pub(super) struct GlobalIndexScanResult { + pub(super) row_ranges: Vec, + pub(super) evaluated_field_ids: HashSet, + pub(super) indexed_coverage: Vec, +} + +pub(super) async fn try_fold_bounded( + futures: impl IntoIterator, + max_concurrency: usize, + mut accumulator: Acc, + mut fold: Fold, +) -> Result +where + Fut: Future>, + Fold: FnMut(&mut Acc, T), +{ + debug_assert!(max_concurrency > 0); + let stream = futures::stream::iter(futures).buffer_unordered(max_concurrency); + futures::pin_mut!(stream); + while let Some(value) = stream.try_next().await? { + fold(&mut accumulator, value); + } + Ok(accumulator) +} + +impl GlobalIndexScanner { + /// Evaluate a predicate against the global indexes and return matching row ranges. + /// Returns `None` if the predicate cannot be evaluated by the global index. + pub(super) fn evaluate<'a>(&'a self, predicate: &'a Predicate) -> EvaluateFuture<'a> { + Box::pin(async move { + match predicate { + Predicate::Leaf { + column, + op, + literals, + data_type, + .. + } => { + if !is_sorted_global_index_supported_op(*op) { + return Ok(None); + } + let field_id = self.find_field_id_by_name(column)?; + let field_id = match field_id { + Some(id) => id, + None => return Ok(None), + }; + let entries = match self.entries_for_field(field_id) { + Some(e) => e, + None => return Ok(None), + }; + if !entries_support_predicate(entries, *op, literals) { + return Ok(None); + } + let predicates = [(*op, literals.as_slice(), data_type)]; + let selected_entries = select_entries_for_predicates(entries, &predicates); + self.evaluate_leaf(&selected_entries, &predicates) + .await + .map(|result| { + result.map(|(row_ranges, indexed_coverage)| GlobalIndexScanResult { + row_ranges, + evaluated_field_ids: HashSet::from([field_id]), + indexed_coverage, + }) + }) + } + Predicate::And(children) => { + // Group leaf predicates by field_id to reuse readers + let mut leaf_groups: std::collections::HashMap>> = + std::collections::HashMap::new(); + let mut non_leaf_children = Vec::new(); + + for child in children { + if let Predicate::Leaf { + column, + op, + literals, + data_type, + .. + } = child + { + if is_sorted_global_index_supported_op(*op) { + if let Some(field_id) = self.find_field_id_by_name(column)? { + if self.entries_for_field(field_id).is_some_and(|entries| { + entries_support_predicate(entries, *op, literals) + }) { + leaf_groups.entry(field_id).or_default().push(( + *op, + literals.as_slice(), + data_type, + )); + continue; + } + } + } + } + non_leaf_children.push(child); + } + + // Evaluate independent fields concurrently while keeping predicates for the + // same field together so each index file is opened only once. + let mut leaf_futures = Vec::with_capacity(leaf_groups.len()); + for (field_id, predicates) in &leaf_groups { + if let Some(entries) = self.entries_for_field(*field_id) { + let field_id = *field_id; + let mut selected_predicates = predicates.clone(); + let mut selected_entries = + select_entries_for_predicates(entries, &selected_predicates); + if selected_entries.is_empty() { + for predicate in predicates { + let candidate_predicates = vec![*predicate]; + let candidate_entries = select_entries_for_predicates( + entries, + &candidate_predicates, + ); + if !candidate_entries.is_empty() { + selected_predicates = candidate_predicates; + selected_entries = candidate_entries; + break; + } + } + } + leaf_futures.push(async move { + let result = self + .evaluate_leaf(&selected_entries, &selected_predicates) + .await?; + Ok((field_id, result)) + }); + } + } + let leaf_group_count = leaf_futures.len(); + let (mut row_ranges, mut indexed_coverage, mut evaluated_field_ids) = + try_fold_bounded( + leaf_futures, + leaf_group_count.max(1), + (None::>, None::>, HashSet::new()), + |(row_ranges, indexed_coverage, evaluated_field_ids), + (field_id, result)| { + if let Some((ranges, coverage)) = result { + *row_ranges = Some(match row_ranges.take() { + None => ranges, + Some(existing) => { + intersect_sorted_ranges(&existing, &ranges) + } + }); + *indexed_coverage = Some(match indexed_coverage.take() { + None => coverage, + Some(existing) => { + intersect_sorted_ranges(&existing, &coverage) + } + }); + evaluated_field_ids.insert(field_id); + } + }, + ) + .await?; + + // Evaluate non-leaf children recursively + for child in non_leaf_children { + if let Some(child_result) = self.evaluate(child).await? { + row_ranges = Some(match row_ranges { + None => child_result.row_ranges, + Some(existing) => { + intersect_sorted_ranges(&existing, &child_result.row_ranges) + } + }); + evaluated_field_ids.extend(child_result.evaluated_field_ids); + indexed_coverage = Some(match indexed_coverage { + None => child_result.indexed_coverage, + Some(existing) => intersect_sorted_ranges( + &existing, + &child_result.indexed_coverage, + ), + }); + } + } + + Ok(row_ranges.map(|row_ranges| GlobalIndexScanResult { + row_ranges, + evaluated_field_ids, + indexed_coverage: indexed_coverage.unwrap_or_default(), + })) + } + Predicate::Or(children) => { + let mut all_ranges: Vec = Vec::new(); + let mut evaluated_field_ids = HashSet::new(); + let mut indexed_coverage: Option> = None; + for child in children { + match self.evaluate(child).await? { + Some(child_result) => { + all_ranges.extend(child_result.row_ranges); + evaluated_field_ids.extend(child_result.evaluated_field_ids); + indexed_coverage = Some(match indexed_coverage { + None => child_result.indexed_coverage, + Some(existing) => intersect_sorted_ranges( + &existing, + &child_result.indexed_coverage, + ), + }); + } + None => return Ok(None), + } + } + let row_ranges = if all_ranges.is_empty() { + Vec::new() + } else { + crate::table::merge_row_ranges(all_ranges) + }; + Ok(Some(GlobalIndexScanResult { + row_ranges, + evaluated_field_ids, + indexed_coverage: indexed_coverage.unwrap_or_default(), + })) + } + _ => Ok(None), + } + }) + } + + /// Evaluate multiple predicates against the same set of index entries. + /// Opens each file once and evaluates all predicates, intersecting results. + /// Detects between patterns (GtEq/Gt + LtEq/Lt) and merges them into a single range query. + async fn evaluate_leaf( + &self, + entries: &[&GlobalIndexEntry], + predicates: &[(PredicateOperator, &[Datum], &DataType)], + ) -> Result, Vec)>> { + let normalized_predicates = predicates + .iter() + .map(|(op, literals, data_type)| { + let key_type = if is_multivalue_predicate(*op) { + let DataType::Array(array) = data_type else { + return Err(Error::DataInvalid { + message: format!( + "Array global-index predicate {op} requires an ARRAY field type" + ), + source: None, + }); + }; + array.element_type() + } else { + *data_type + }; + Ok((*op, *literals, key_type)) + }) + .collect::>>()?; + let predicates = normalized_predicates.as_slice(); + + // Try to detect between pattern and split into (between, remaining) + let (between, remaining) = extract_between(predicates); + + let effective_predicates = if between.is_some() { + &remaining + } else { + predicates + }; + + // Pre-compute comparators and serialized keys for file-level pruning per predicate + let pruning_info: Vec<_> = effective_predicates + .iter() + .map(|(op, literals, data_type)| { + let btree_cmp = make_key_comparator(data_type); + let btree_serialized = literals + .iter() + .map(|l| serialize_datum(l, data_type)) + .collect::>(); + let bitmap_cmp = make_bitmap_key_comparator(data_type); + let bitmap_serialized = literals + .iter() + .map(|l| serialize_bitmap_datum(l, data_type)) + .collect::>(); + ( + *op, + *data_type, + btree_cmp, + btree_serialized, + bitmap_cmp, + bitmap_serialized, + ) + }) + .collect(); + + let predicate_matches: Vec> = pruning_info + .iter() + .map( + |(op, data_type, btree_cmp, btree_serialized, bitmap_cmp, bitmap_serialized)| { + entries + .iter() + .map(|entry| match entry.index_type { + GlobalIndexFileKind::BTree => { + sorted_entry_meta(entry).may_match(*op, btree_serialized, btree_cmp) + } + GlobalIndexFileKind::Bitmap => bitmap_meta_may_match( + sorted_entry_meta(entry), + *op, + data_type, + bitmap_serialized, + bitmap_cmp.as_ref(), + ), + GlobalIndexFileKind::Multivalue => multivalue_meta_may_match( + sorted_entry_meta(entry), + *op, + bitmap_serialized, + bitmap_cmp.as_ref(), + ), + GlobalIndexFileKind::FM => true, + }) + .collect() + }, + ) + .collect(); + let predicate_fallback_plans: Vec> = effective_predicates + .iter() + .enumerate() + .map(|(i, (op, _, _))| { + requires_fallback_scan(*op) + .then(|| self.fallback_scan_plan(entries, &predicate_matches[i])) + }) + .collect(); + + let between_matches_by_entry: Vec = + match between.as_ref() { + Some(b) => { + let btree_cmp = make_key_comparator(b.data_type); + let btree_from = serialize_datum(b.from, b.data_type); + let btree_to = serialize_datum(b.to, b.data_type); + let bitmap_cmp = make_bitmap_key_comparator(b.data_type); + let bitmap_from = serialize_bitmap_datum(b.from, b.data_type); + let bitmap_to = serialize_bitmap_datum(b.to, b.data_type); + entries + .iter() + .map(|entry| match entry.index_type { + GlobalIndexFileKind::BTree => sorted_entry_meta(entry) + .may_match_between(&btree_from, &btree_to, &btree_cmp), + GlobalIndexFileKind::Bitmap => bitmap_meta_may_match_between( + sorted_entry_meta(entry), + b.data_type, + &bitmap_from, + &bitmap_to, + bitmap_cmp.as_ref(), + ), + GlobalIndexFileKind::Multivalue | GlobalIndexFileKind::FM => false, + }) + .collect() + } + None => Vec::new(), + }; + let between_fallback_plan = between + .as_ref() + .map(|_| self.fallback_scan_plan(entries, &between_matches_by_entry)); + + let mut query_plans = Vec::with_capacity(entries.len()); + for (entry_idx, entry) in entries.iter().enumerate() { + // Also check if between range may match + let between_matches = between + .as_ref() + .is_some_and(|_| between_matches_by_entry[entry_idx]); + let between_evaluated_for_entry = between_fallback_plan.is_some_and(|plan| { + fallback_plan_evaluates_entry(plan, entry.index_type, between_matches) + }); + + // When a Between conjunct exists but the file does not overlap its + // range, the whole AND cannot match — drop the file regardless of + // how the remaining predicates evaluate. Without this guard, a file + // outside the Between range but matched by some remaining predicate + // (e.g. `BETWEEN 10 AND 20 AND id >= 0` on a file [30, 40]) would + // be retained because `file_result` is initialized from the + // remaining bitmap, silently dropping the Between conjunct. + if between_evaluated_for_entry && !between_matches { + continue; + } + + let mut file_evaluated = between_evaluated_for_entry; + let mut file_cannot_match = false; + let mut file_has_unsupported_match = + between_matches && !between_evaluated_for_entry && between_fallback_plan.is_some(); + let matching_predicates: Vec = (0..effective_predicates.len()) + .filter(|&i| { + let predicate_matches_entry = predicate_matches[i][entry_idx]; + let predicate_evaluated_for_entry = + predicate_fallback_plans[i].is_none_or(|plan| { + fallback_plan_evaluates_entry( + plan, + entry.index_type, + predicate_matches_entry, + ) + }); + if !predicate_evaluated_for_entry { + file_has_unsupported_match |= predicate_matches_entry; + return false; + } + file_evaluated = true; + if !predicate_matches_entry { + file_cannot_match = true; + return false; + } + true + }) + .collect(); + if file_cannot_match { + continue; + } + if !file_evaluated { + if file_has_unsupported_match { + return Ok(None); + } + continue; + } + + query_plans.push(EntryQueryPlan { + entry_idx, + between_matches, + between_evaluated: between_evaluated_for_entry, + matching_predicates, + }); + } + + // Complete all pruning and fallback decisions before starting shard I/O. + // A later unsupported shard must fall back to the normal scan without an + // earlier shard racing it with an I/O or query error. + let data_type = between + .as_ref() + .map(|b| b.data_type) + .or_else(|| effective_predicates.first().map(|p| p.2)) + .unwrap_or(predicates[0].2); + let between = between.as_ref(); + let futures = + query_plans.into_iter().map(|plan| async move { + let entry = &entries[plan.entry_idx]; + let _permit = self.query_semaphore.acquire().await.map_err(|error| { + Error::UnexpectedError { + message: "global-index query concurrency budget was closed".to_string(), + source: Some(Box::new(error)), + } + })?; + #[cfg(test)] + let _query_io_probe_guard = match &self.query_io_probe { + Some(probe) => Some(probe.enter().await), + None => None, + }; + let result = self + .query_entry(entry, data_type, between, &plan, effective_predicates) + .await?; + Ok((entry.row_range_start, result)) + }); + let (all_row_ids, declined) = try_fold_bounded( + futures, + self.global_index_thread_num, + (RoaringTreemap::new(), false), + |(all_row_ids, declined), (row_range_start, file_result)| { + if file_result.declined { + *declined = true; + return; + } + if let Some(bitmap) = file_result.bitmap { + for row_id in bitmap.iter() { + all_row_ids.insert(row_id + row_range_start as u64); + } + } + }, + ) + .await?; + + if declined { + return Ok(None); + } + + let coverage = crate::table::merge_row_ranges( + entries + .iter() + .map(|entry| RowRange::new(entry.row_range_start, entry.row_range_end)) + .collect(), + ); + Ok(Some((bitmap_to_ranges(&all_row_ids), coverage))) + } + + fn find_field_id_by_name(&self, column: &str) -> Result> { + Ok(crate::table::find_field_id_by_name( + &self.schema_fields, + column, + )) + } + + fn entries_for_field(&self, field_id: i32) -> Option<&[GlobalIndexEntry]> { + self.entries_by_field + .iter() + .find(|(id, _)| *id == field_id) + .map(|(_, entries)| entries.as_slice()) + } + + /// Return row ranges not covered by global indexes for this predicate. + /// + /// `full` uses `[0, snapshot.next_row_id - 1]`; `detail` uses actual + /// data-file row ranges collected by the scan. The caller unions these + /// ranges with indexed matches, and the normal read filter evaluates the + /// predicate on the raw rows. + #[cfg(test)] + pub(super) fn unindexed_ranges( + &self, + predicate: &Predicate, + search_mode: GlobalIndexSearchMode, + next_row_id: Option, + data_ranges: &[RowRange], + ) -> Result> { + let field_ids = self.collect_field_ids(predicate)?; + Ok(self.unindexed_ranges_for_field_ids(&field_ids, search_mode, next_row_id, data_ranges)) + } + + #[cfg(test)] + fn unindexed_ranges_for_field_ids( + &self, + field_ids: &HashSet, + search_mode: GlobalIndexSearchMode, + next_row_id: Option, + data_ranges: &[RowRange], + ) -> Vec { + unindexed_ranges_for_coverage( + &self.coverage_by_field, + field_ids, + search_mode, + next_row_id, + data_ranges, + ) + } + + #[cfg(test)] + fn collect_field_ids(&self, predicate: &Predicate) -> Result> { + let mut field_ids = HashSet::new(); + self.collect_field_ids_inner(predicate, &mut field_ids)?; + Ok(field_ids) + } + + #[cfg(test)] + fn collect_field_ids_inner( + &self, + predicate: &Predicate, + field_ids: &mut HashSet, + ) -> Result<()> { + match predicate { + Predicate::Leaf { column, .. } => { + if let Some(field_id) = self.find_field_id_by_name(column)? { + field_ids.insert(field_id); + } + } + Predicate::And(children) | Predicate::Or(children) => { + for child in children { + self.collect_field_ids_inner(child, field_ids)?; + } + } + Predicate::Not(inner) => self.collect_field_ids_inner(inner, field_ids)?, + Predicate::AlwaysTrue | Predicate::AlwaysFalse => {} + } + Ok(()) + } +} diff --git a/crates/paimon/src/table/global_index_scanner/predicates.rs b/crates/paimon/src/table/global_index_scanner/predicates.rs new file mode 100644 index 000000000..32f38a717 --- /dev/null +++ b/crates/paimon/src/table/global_index_scanner/predicates.rs @@ -0,0 +1,115 @@ +// 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. + +//! Predicate capability checks and fallback-policy decisions. + +use super::entry::{GlobalIndexEntry, GlobalIndexFileKind}; +use crate::spec::{DataType, Datum, PredicateOperator}; + +/// Whether the sorted global index can evaluate this operator directly. +/// Operators that fall outside this set bypass the index and are evaluated +/// later in the read pipeline (stats prune + parquet row filter). +pub(super) fn is_sorted_global_index_supported_op(op: PredicateOperator) -> bool { + matches!( + op, + PredicateOperator::Eq + | PredicateOperator::NotEq + | PredicateOperator::Lt + | PredicateOperator::LtEq + | PredicateOperator::Gt + | PredicateOperator::GtEq + | PredicateOperator::In + | PredicateOperator::NotIn + | PredicateOperator::IsNull + | PredicateOperator::IsNotNull + | PredicateOperator::Between + | PredicateOperator::NotBetween + | PredicateOperator::StartsWith + | PredicateOperator::EndsWith + | PredicateOperator::Contains + | PredicateOperator::Like + | PredicateOperator::ArrayContains + | PredicateOperator::ArraysOverlap + | PredicateOperator::ArrayContainsAll + ) +} + +pub(super) fn is_multivalue_predicate(op: PredicateOperator) -> bool { + matches!( + op, + PredicateOperator::ArrayContains + | PredicateOperator::ArraysOverlap + | PredicateOperator::ArrayContainsAll + ) +} + +pub(super) fn entry_supports_predicate( + entry: &GlobalIndexEntry, + op: PredicateOperator, + literals: &[Datum], +) -> bool { + match entry.index_type { + GlobalIndexFileKind::Multivalue => { + is_multivalue_predicate(op) + && !(matches!(op, PredicateOperator::ArrayContainsAll) && literals.is_empty()) + } + GlobalIndexFileKind::FM => { + matches!(op, PredicateOperator::IsNull | PredicateOperator::IsNotNull) + || (op == PredicateOperator::Contains && literals.len() == 1) + } + GlobalIndexFileKind::BTree | GlobalIndexFileKind::Bitmap => !is_multivalue_predicate(op), + } +} + +pub(super) fn entries_support_predicate( + entries: &[GlobalIndexEntry], + op: PredicateOperator, + literals: &[Datum], +) -> bool { + entries + .iter() + .any(|entry| entry_supports_predicate(entry, op, literals)) +} + +pub(super) fn select_entries_for_predicates<'a>( + entries: &'a [GlobalIndexEntry], + predicates: &[(PredicateOperator, &[Datum], &DataType)], +) -> Vec<&'a GlobalIndexEntry> { + let compatible = entries + .iter() + .filter(|entry| { + predicates + .iter() + .all(|(op, literals, _)| entry_supports_predicate(entry, *op, literals)) + }) + .collect::>(); + + if predicates + .iter() + .any(|(op, _, _)| *op == PredicateOperator::Contains) + && compatible + .iter() + .any(|entry| entry.index_type == GlobalIndexFileKind::FM) + { + compatible + .into_iter() + .filter(|entry| entry.index_type == GlobalIndexFileKind::FM) + .collect() + } else { + compatible + } +} diff --git a/crates/paimon/src/table/global_index_scanner/query_plan.rs b/crates/paimon/src/table/global_index_scanner/query_plan.rs new file mode 100644 index 000000000..2ff579ff3 --- /dev/null +++ b/crates/paimon/src/table/global_index_scanner/query_plan.rs @@ -0,0 +1,88 @@ +// 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. + +//! Per-entry query plans and fallback-scan policy. + +use super::entry::GlobalIndexFileKind; +use crate::spec::PredicateOperator; +use roaring::RoaringTreemap; + +#[derive(Clone, Copy, Default)] +pub(super) struct FallbackScanPlan { + pub(super) selected_btree: usize, + pub(super) selected_bitmap: usize, + pub(super) allow_btree: bool, + pub(super) allow_bitmap: bool, +} + +pub(super) struct EntryQueryPlan { + pub(super) entry_idx: usize, + pub(super) between_matches: bool, + pub(super) between_evaluated: bool, + pub(super) matching_predicates: Vec, +} + +pub(super) struct EntryQueryResult { + pub(super) bitmap: Option, + pub(super) declined: bool, +} + +impl FallbackScanPlan { + pub(super) fn allowed(self, kind: GlobalIndexFileKind) -> bool { + match kind { + GlobalIndexFileKind::BTree => self.allow_btree, + GlobalIndexFileKind::Bitmap | GlobalIndexFileKind::Multivalue => self.allow_bitmap, + GlobalIndexFileKind::FM => true, + } + } +} + +pub(super) fn requires_fallback_scan(op: PredicateOperator) -> bool { + matches!( + op, + PredicateOperator::Lt + | PredicateOperator::LtEq + | PredicateOperator::Gt + | PredicateOperator::GtEq + | PredicateOperator::Between + | PredicateOperator::NotBetween + | PredicateOperator::EndsWith + | PredicateOperator::Contains + | PredicateOperator::Like + ) +} + +pub(super) fn fallback_plan_evaluates_entry( + plan: FallbackScanPlan, + kind: GlobalIndexFileKind, + selected: bool, +) -> bool { + !selected || plan.allowed(kind) +} + +pub(super) fn add_file_size(total: &mut i64, file_size: i64) -> bool { + if file_size < 0 { + return false; + } + match total.checked_add(file_size) { + Some(next) => { + *total = next; + true + } + None => false, + } +} diff --git a/crates/paimon/src/table/global_index_scanner/reader.rs b/crates/paimon/src/table/global_index_scanner/reader.rs new file mode 100644 index 000000000..acc576e98 --- /dev/null +++ b/crates/paimon/src/table/global_index_scanner/reader.rs @@ -0,0 +1,380 @@ +// 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. + +//! Type-specific index readers and shard I/O execution. + +use super::entry::{ + sorted_entry_meta, GlobalIndexEntry, GlobalIndexEntryMeta, GlobalIndexFileKind, +}; +use super::query_plan::{add_file_size, EntryQueryPlan, EntryQueryResult, FallbackScanPlan}; +use super::{BoxedCmp, GlobalIndexScanner, INDEX_DIR}; +use crate::btree::query::{BetweenInfo, IndexQuery}; +use crate::btree::{make_key_comparator, serialize_datum, BTreeIndexMeta, BTreeIndexReader}; +use crate::fm_index::FMGlobalIndexReader; +use crate::spec::{DataType, Datum, PredicateOperator}; +use crate::table::bitmap_global_index_format::serialize_bitmap_datum; +use crate::table::bitmap_global_index_reader::BitmapGlobalIndexReader; +use crate::{Error, Result}; +use roaring::RoaringTreemap; +use std::sync::Arc; + +pub(super) enum OpenedGlobalIndexReader { + BTree(BTreeIndexReader), + Bitmap(BitmapGlobalIndexReader), + FM(FMGlobalIndexReader), +} + +impl OpenedGlobalIndexReader { + async fn query( + &self, + op: PredicateOperator, + literals: &[Datum], + data_type: &DataType, + ) -> std::io::Result> { + match self { + Self::BTree(reader) => reader.query(op, literals, data_type).await.map(Some), + Self::Bitmap(reader) => reader.query(op, literals, data_type).await.map(Some), + Self::FM(reader) => match op { + PredicateOperator::Contains => { + let literal = literals.first().ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "FM contains requires one literal", + ) + })?; + reader + .contains(&serialize_bitmap_datum(literal, data_type)) + .await + } + PredicateOperator::IsNull => reader.is_null().await.map(Some), + PredicateOperator::IsNotNull => reader.is_not_null().await.map(Some), + _ => Ok(None), + }, + } + } + + async fn range_query( + &self, + from: &[u8], + to: &[u8], + data_type: &DataType, + from_inclusive: bool, + to_inclusive: bool, + ) -> std::io::Result { + match self { + Self::BTree(reader) => { + reader + .range_query(from, to, from_inclusive, to_inclusive) + .await + } + Self::Bitmap(reader) => { + reader + .range_query(from, to, data_type, from_inclusive, to_inclusive) + .await + } + Self::FM(_) => Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "FM index does not support ordered range queries", + )), + } + } +} + +impl GlobalIndexScanner { + pub(super) async fn query_entry( + &self, + entry: &GlobalIndexEntry, + data_type: &DataType, + between: Option<&BetweenInfo<'_>>, + plan: &EntryQueryPlan, + effective_predicates: &[(PredicateOperator, &[Datum], &DataType)], + ) -> Result { + let mut reader = if (plan.between_matches && plan.between_evaluated) + || !plan.matching_predicates.is_empty() + { + Some(self.open_reader_for_entry(entry, data_type).await?) + } else { + None + }; + let mut file_result = None; + + if plan.between_matches && plan.between_evaluated { + let between = between.expect("evaluated between query is present"); + let serialize_key = match entry.index_type { + GlobalIndexFileKind::BTree => serialize_datum, + GlobalIndexFileKind::Bitmap | GlobalIndexFileKind::Multivalue => { + serialize_bitmap_datum + } + GlobalIndexFileKind::FM => unreachable!("FM range query was rejected in planning"), + }; + let from_key = serialize_key(between.from, between.data_type); + let to_key = serialize_key(between.to, between.data_type); + let bitmap = reader + .as_ref() + .expect("reader is opened when between matches") + .range_query( + &from_key, + &to_key, + between.data_type, + between.from_inclusive, + between.to_inclusive, + ) + .await + .map_err(|error| Self::query_error(entry, error))?; + file_result = Some(bitmap); + } + + for &idx in &plan.matching_predicates { + let (op, literals, data_type) = &effective_predicates[idx]; + let Some(bitmap) = reader + .as_ref() + .expect("reader is opened when predicates match") + .query(*op, literals, data_type) + .await + .map_err(|error| Self::query_error(entry, error))? + else { + return Ok(EntryQueryResult { + bitmap: None, + declined: true, + }); + }; + file_result = Some(match file_result { + None => bitmap, + Some(mut existing) => { + existing &= bitmap; + existing + } + }); + } + + // Each concurrent task owns its reader. Only return it to the shared + // cache after all predicates for this shard have completed. + if let Some(OpenedGlobalIndexReader::BTree(reader)) = reader.take() { + self.return_reader(entry.file_name.clone(), reader); + } + Ok(EntryQueryResult { + bitmap: file_result, + declined: false, + }) + } + + fn query_error(entry: &GlobalIndexEntry, error: std::io::Error) -> Error { + Error::DataInvalid { + message: format!( + "Global index query failed for {} file '{}'", + entry.index_type.name(), + entry.file_name + ), + source: Some(Box::new(error)), + } + } + + /// Get a cached reader or open a new one for the given file. + async fn get_or_open_reader( + &self, + entry: &GlobalIndexEntry, + meta: &BTreeIndexMeta, + data_type: &DataType, + ) -> Result { + // Try to take from cache + { + let mut cache = self.reader_cache.lock().unwrap(); + if let Some(reader) = cache.remove(&entry.file_name) { + return Ok(OpenedGlobalIndexReader::BTree(reader)); + } + } + + // Open new reader + let path = format!("{}/{INDEX_DIR}/{}", self.table_path, entry.file_name); + let input = self.file_io.new_input(&path)?; + let file_size = if entry.file_size > 0 { + entry.file_size as u64 + } else { + input.metadata().await?.size + }; + let file_reader = input.reader().await?; + + let cmp = make_key_comparator(data_type); + BTreeIndexReader::open(Box::new(file_reader), file_size, meta, cmp) + .await + .map(OpenedGlobalIndexReader::BTree) + .map_err(|e| crate::Error::DataInvalid { + message: format!("Failed to open BTree index file: {}", entry.file_name), + source: Some(Box::new(e)), + }) + } + + async fn open_reader_for_entry( + &self, + entry: &GlobalIndexEntry, + data_type: &DataType, + ) -> Result { + match entry.index_type { + GlobalIndexFileKind::BTree => { + self.get_or_open_reader(entry, sorted_entry_meta(entry), data_type) + .await + } + GlobalIndexFileKind::Bitmap => self + .open_bitmap_reader(entry) + .await + .map(OpenedGlobalIndexReader::Bitmap) + .map_err(|e| crate::Error::DataInvalid { + message: format!( + "Failed to open bitmap global index file: {}", + entry.file_name + ), + source: Some(Box::new(e)), + }), + GlobalIndexFileKind::Multivalue => self + .open_bitmap_reader(entry) + .await + .map(OpenedGlobalIndexReader::Bitmap) + .map_err(|e| crate::Error::DataInvalid { + message: format!( + "Failed to open multivalue global index file: {}", + entry.file_name + ), + source: Some(Box::new(e)), + }), + GlobalIndexFileKind::FM => self + .open_fm_reader(entry) + .await + .map(OpenedGlobalIndexReader::FM) + .map_err(|e| crate::Error::DataInvalid { + message: format!("Failed to open FM global index file: {}", entry.file_name), + source: Some(Box::new(e)), + }), + } + } + + async fn open_fm_reader( + &self, + entry: &GlobalIndexEntry, + ) -> std::io::Result { + let GlobalIndexEntryMeta::FM { + bytes: manifest_meta, + .. + } = &entry.meta + else { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "FM entry has non-FM manifest metadata", + )); + }; + let path = format!("{}/{INDEX_DIR}/{}", self.table_path, entry.file_name); + let input = self + .file_io + .new_input(&path) + .map_err(|e| std::io::Error::other(e.to_string()))?; + let file_size = if entry.file_size > 0 { + entry.file_size as u64 + } else { + input + .metadata() + .await + .map_err(|e| std::io::Error::other(e.to_string()))? + .size + }; + let file_reader = input + .reader() + .await + .map_err(|e| std::io::Error::other(e.to_string()))?; + FMGlobalIndexReader::open_with_context( + Box::new(file_reader), + file_size, + manifest_meta, + self.fm_read_options, + Arc::clone(&self.fm_read_context), + entry.file_name.clone(), + ) + .await + } + + async fn open_bitmap_reader( + &self, + entry: &GlobalIndexEntry, + ) -> std::io::Result { + let path = format!("{}/{INDEX_DIR}/{}", self.table_path, entry.file_name); + let input = self + .file_io + .new_input(&path) + .map_err(|e| std::io::Error::other(e.to_string()))?; + let file_size = if entry.file_size > 0 { + entry.file_size as u64 + } else { + input + .metadata() + .await + .map_err(|e| std::io::Error::other(e.to_string()))? + .size + }; + let file_reader = input + .reader() + .await + .map_err(|e| std::io::Error::other(e.to_string()))?; + BitmapGlobalIndexReader::open(Box::new(file_reader), file_size).await + } + + pub(super) fn fallback_scan_plan( + &self, + entries: &[&GlobalIndexEntry], + selected: &[bool], + ) -> FallbackScanPlan { + let mut plan = FallbackScanPlan::default(); + let mut btree_total = 0i64; + let mut bitmap_total = 0i64; + let mut btree_valid = true; + let mut bitmap_valid = true; + + for (entry, selected) in entries.iter().zip(selected) { + if !selected { + continue; + } + match entry.index_type { + GlobalIndexFileKind::BTree => { + plan.selected_btree += 1; + btree_valid &= add_file_size(&mut btree_total, entry.file_size); + } + GlobalIndexFileKind::Bitmap => { + plan.selected_bitmap += 1; + bitmap_valid &= add_file_size(&mut bitmap_total, entry.file_size); + } + GlobalIndexFileKind::Multivalue => { + plan.selected_bitmap += 1; + bitmap_valid &= add_file_size(&mut bitmap_total, entry.file_size); + } + GlobalIndexFileKind::FM => {} + } + } + + plan.allow_btree = plan.selected_btree > 0 + && btree_valid + && self.btree_fallback_scan_max_size > 0 + && btree_total <= self.btree_fallback_scan_max_size; + plan.allow_bitmap = plan.selected_bitmap > 0 + && bitmap_valid + && self.bitmap_fallback_scan_max_size > 0 + && bitmap_total <= self.bitmap_fallback_scan_max_size; + plan + } + + /// Return a reader to the cache for future reuse. + fn return_reader(&self, file_name: String, reader: BTreeIndexReader) { + let mut cache = self.reader_cache.lock().unwrap(); + cache.insert(file_name, reader); + } +} diff --git a/crates/paimon/src/table/global_index_scanner/row_ranges.rs b/crates/paimon/src/table/global_index_scanner/row_ranges.rs new file mode 100644 index 000000000..abfe5df73 --- /dev/null +++ b/crates/paimon/src/table/global_index_scanner/row_ranges.rs @@ -0,0 +1,296 @@ +// 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. + +//! Row-range conversion, coverage accounting, and lookup utilities. + +use crate::spec::{FileKind, GlobalIndexSearchMode, IndexFileMeta, IndexManifestEntry}; +use crate::table::{merge_row_ranges, source, RowRange}; +use roaring::RoaringTreemap; +use std::collections::{HashMap, HashSet}; + +/// Convert a RoaringTreemap to merged RowRanges (already sorted and deduplicated). +pub(super) fn bitmap_to_ranges(bitmap: &RoaringTreemap) -> Vec { + if bitmap.is_empty() { + return Vec::new(); + } + let mut ranges = Vec::new(); + let mut iter = bitmap.iter(); + let first = iter.next().unwrap(); + let mut start = first as i64; + let mut end = start; + + for id in iter { + let id = id as i64; + if id == end + 1 { + end = id; + } else { + ranges.push(RowRange::new(start, end)); + start = id; + end = id; + } + } + ranges.push(RowRange::new(start, end)); + ranges +} + +/// Intersect two sorted range lists using RowRangeIndex for efficient binary search. +pub(super) fn intersect_sorted_ranges(a: &[RowRange], b: &[RowRange]) -> Vec { + let idx = RowRangeIndex::create(a.to_vec()); + let mut result = Vec::new(); + for r in b { + result.extend(idx.intersected_ranges(r.from(), r.to())); + } + result +} + +pub(super) fn data_ranges_for_search_mode( + search_mode: GlobalIndexSearchMode, + next_row_id: Option, + data_ranges: &[RowRange], +) -> Option> { + match search_mode { + GlobalIndexSearchMode::Fast => None, + GlobalIndexSearchMode::Full => match next_row_id { + Some(next_row_id) if next_row_id > 0 => Some(vec![RowRange::new(0, next_row_id - 1)]), + _ => None, + }, + GlobalIndexSearchMode::Detail => { + if data_ranges.is_empty() { + None + } else { + Some(data_ranges.to_vec()) + } + } + } +} + +pub(super) fn indexed_ranges_from_coverage( + coverage_by_field: &HashMap>, + field_ids: &HashSet, +) -> Vec { + let mut ranges: Option> = None; + for field_id in field_ids { + let Some(field_ranges) = coverage_by_field.get(field_id) else { + return Vec::new(); + }; + if field_ranges.is_empty() { + return Vec::new(); + } + let field_ranges = merge_row_ranges(field_ranges.clone()); + ranges = Some(match ranges { + None => field_ranges, + Some(existing) => intersect_sorted_ranges(&existing, &field_ranges), + }); + } + ranges.map(merge_row_ranges).unwrap_or_default() +} + +pub(super) fn unindexed_ranges_for_coverage( + coverage_by_field: &HashMap>, + field_ids: &HashSet, + search_mode: GlobalIndexSearchMode, + next_row_id: Option, + data_ranges: &[RowRange], +) -> Vec { + let Some(data_ranges) = data_ranges_for_search_mode(search_mode, next_row_id, data_ranges) + else { + return Vec::new(); + }; + let indexed_ranges = indexed_ranges_from_coverage(coverage_by_field, field_ids); + source::exclude_row_ranges(&data_ranges, &indexed_ranges) +} + +pub(super) fn unindexed_ranges_for_indexed_coverage( + indexed_ranges: &[RowRange], + search_mode: GlobalIndexSearchMode, + next_row_id: Option, + data_ranges: &[RowRange], +) -> Vec { + let Some(data_ranges) = data_ranges_for_search_mode(search_mode, next_row_id, data_ranges) + else { + return Vec::new(); + }; + source::exclude_row_ranges(&data_ranges, &merge_row_ranges(indexed_ranges.to_vec())) +} + +/// Compute row ranges not covered by a family of global index files. +/// +/// This mirrors Java `GlobalIndexCoverage`: `full` compares index coverage +/// against `[0, snapshot.next_row_id - 1]`, while `detail` compares against +/// exact data-file row ranges supplied by the caller. +pub(crate) fn unindexed_ranges_for_global_index_entries( + index_entries: &[IndexManifestEntry], + field_ids: &HashSet, + search_mode: GlobalIndexSearchMode, + next_row_id: Option, + data_ranges: &[RowRange], + index_file_filter: impl Fn(&IndexFileMeta) -> bool, +) -> Vec { + let mut coverage_by_field: HashMap> = HashMap::new(); + for entry in index_entries { + if entry.kind != FileKind::Add || !index_file_filter(&entry.index_file) { + continue; + } + let Some(global_meta) = entry.index_file.global_index_meta.as_ref() else { + continue; + }; + let row_range = RowRange::new(global_meta.row_range_start, global_meta.row_range_end); + coverage_by_field + .entry(global_meta.index_field_id) + .or_default() + .push(row_range.clone()); + if let Some(extra_field_ids) = global_meta.extra_field_ids.as_ref() { + for extra_field_id in extra_field_ids { + coverage_by_field + .entry(*extra_field_id) + .or_default() + .push(row_range.clone()); + } + } + } + unindexed_ranges_for_coverage( + &coverage_by_field, + field_ids, + search_mode, + next_row_id, + data_ranges, + ) +} + +/// Index for row ranges. Stores sorted, non-overlapping ranges and supports +/// efficient intersection queries via binary search. +/// +/// Reference: [org.apache.paimon.utils.RowRangeIndex](https://github.com/apache/paimon/blob/master/paimon-common/src/main/java/org/apache/paimon/utils/RowRangeIndex.java) +#[derive(Debug, Clone)] +pub(crate) struct RowRangeIndex { + ranges: Vec, + starts: Vec, + ends: Vec, +} + +impl RowRangeIndex { + /// Create a new `RowRangeIndex` from a list of ranges. + /// Ranges are sorted and merged to eliminate overlaps. + pub fn create(ranges: Vec) -> Self { + let ranges = merge_row_ranges(ranges); + let starts: Vec = ranges.iter().map(|r| r.from()).collect(); + let ends: Vec = ranges.iter().map(|r| r.to()).collect(); + Self { + ranges, + starts, + ends, + } + } + + /// Returns the underlying ranges. + #[cfg(test)] + pub fn ranges(&self) -> &[RowRange] { + &self.ranges + } + + /// Returns true if the index has any range that intersects `[start, end]`. + pub fn intersects(&self, start: i64, end: i64) -> bool { + let candidate = lower_bound(&self.ends, start); + candidate < self.starts.len() && self.starts[candidate] <= end + } + + /// Counts rows in this index that intersect `[start, end]`. + pub fn intersection_row_count(&self, start: i64, end: i64) -> usize { + if start > end { + return 0; + } + self.intersected_ranges(start, end) + .into_iter() + .fold(0usize, |total, range| { + let len = range.to().saturating_sub(range.from()).saturating_add(1); + total.saturating_add(usize::try_from(len).unwrap_or(usize::MAX)) + }) + } + + /// Returns the sub-ranges of this index that intersect `[start, end]`, + /// clipped to the intersection boundaries. + pub fn intersected_ranges(&self, start: i64, end: i64) -> Vec { + let left = lower_bound(&self.ends, start); + if left >= self.ranges.len() || self.starts[left] > end { + return Vec::new(); + } + + let mut right = lower_bound(&self.ends, end); + if right >= self.ranges.len() { + right = self.ranges.len() - 1; + } + + let mut result = Vec::new(); + + // First range: clip from the left + let first = &self.ranges[left]; + result.push(RowRange::new(start.max(first.from()), end.min(first.to()))); + + // Middle ranges: fully included + if right > left + 1 { + for r in &self.ranges[left + 1..right] { + result.push(r.clone()); + } + } + + // Last range (if different from first): clip from the right + if right != left { + let last = &self.ranges[right]; + if last.from() <= end { + result.push(RowRange::new(start.max(last.from()), end.min(last.to()))); + } + } + + result + } +} + +pub(crate) fn search_limit_with_deleted_rows( + limit: usize, + row_range_start: i64, + row_range_end: i64, + deleted_rows: Option<&RowRangeIndex>, +) -> usize { + let Some(range_len) = row_range_end + .checked_sub(row_range_start) + .and_then(|len| len.checked_add(1)) + .and_then(|len| usize::try_from(len).ok()) + else { + return limit; + }; + + let deleted_count = deleted_rows + .map(|index| index.intersection_row_count(row_range_start, row_range_end)) + .unwrap_or(0) + .min(range_len); + limit.saturating_add(deleted_count).min(range_len) +} + +/// Binary search: find the first index where `sorted[index] >= target`. +fn lower_bound(sorted: &[i64], target: i64) -> usize { + let mut left = 0; + let mut right = sorted.len(); + while left < right { + let mid = left + (right - left) / 2; + if sorted[mid] < target { + left = mid + 1; + } else { + right = mid; + } + } + left +} diff --git a/crates/paimon/src/table/global_index_scanner/tests.rs b/crates/paimon/src/table/global_index_scanner/tests.rs new file mode 100644 index 000000000..29dba260e --- /dev/null +++ b/crates/paimon/src/table/global_index_scanner/tests.rs @@ -0,0 +1,2136 @@ +// 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::entry::*; +use super::evaluator::try_fold_bounded; +use super::predicates::*; +use super::row_ranges::*; +use super::*; +use crate::btree::test_util::VecFileWrite; +use crate::btree::{serialize_datum, BTreeIndexWriter, BlockCompressionType}; +use crate::fm_index::{FMGlobalIndexWriter, FMWriteOptions}; +use crate::spec::{DataType, Datum, IndexFileMeta, PredicateOperator}; +use crate::table::bitmap_global_index_format::{ + make_bitmap_key_comparator, serialize_bitmap_datum, +}; +use crate::table::bitmap_global_index_writer::BitmapGlobalIndexWriter; +use roaring::RoaringTreemap; +use std::collections::HashSet; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; +use std::sync::Arc; + +#[tokio::test] +async fn test_try_fold_bounded_respects_concurrency_limit() { + for limit in [1, 3] { + let active = Arc::new(AtomicUsize::new(0)); + let peak = Arc::new(AtomicUsize::new(0)); + let futures = (0..9usize).map(|value| { + let active = Arc::clone(&active); + let peak = Arc::clone(&peak); + async move { + let current = active.fetch_add(1, AtomicOrdering::SeqCst) + 1; + peak.fetch_max(current, AtomicOrdering::SeqCst); + tokio::task::yield_now().await; + active.fetch_sub(1, AtomicOrdering::SeqCst); + Ok::<_, crate::Error>(value) + } + }); + + let mut values = try_fold_bounded(futures, limit, Vec::new(), |values, value| { + values.push(value) + }) + .await + .unwrap(); + values.sort_unstable(); + + assert_eq!(values, (0..9).collect::>()); + assert_eq!(peak.load(AtomicOrdering::SeqCst), limit); + } +} + +#[test] +fn test_bitmap_to_ranges() { + assert_eq!( + bitmap_to_ranges(&RoaringTreemap::new()), + Vec::::new() + ); + + let mut bm = RoaringTreemap::new(); + bm.insert(5); + assert_eq!(bitmap_to_ranges(&bm), vec![RowRange::new(5, 5)]); + + let mut bm = RoaringTreemap::new(); + for id in [1, 2, 3, 5, 6, 10] { + bm.insert(id); + } + assert_eq!( + bitmap_to_ranges(&bm), + vec![ + RowRange::new(1, 3), + RowRange::new(5, 6), + RowRange::new(10, 10), + ] + ); +} + +#[test] +fn test_intersect_sorted_ranges() { + let a = vec![RowRange::new(0, 10), RowRange::new(20, 30)]; + let b = vec![RowRange::new(5, 25)]; + let result = intersect_sorted_ranges(&a, &b); + assert_eq!(result, vec![RowRange::new(5, 10), RowRange::new(20, 25)]); +} + +#[test] +fn test_intersect_no_overlap() { + let a = vec![RowRange::new(0, 5)]; + let b = vec![RowRange::new(10, 20)]; + assert!(intersect_sorted_ranges(&a, &b).is_empty()); +} + +#[test] +fn test_serialize_datum_int() { + let key = serialize_datum(&Datum::Int(42), &DataType::Int(crate::spec::IntType::new())); + assert_eq!(key, 42i32.to_le_bytes().to_vec()); +} + +#[test] +fn test_serialize_datum_string() { + let key = serialize_datum( + &Datum::String("hello".to_string()), + &DataType::VarChar(crate::spec::VarCharType::new(100).unwrap()), + ); + assert_eq!(key, b"hello".to_vec()); +} + +fn assert_bitmap_floating_meta_policy( + data_type: DataType, + min: Datum, + max: Datum, + outside: Datum, + nan: Datum, +) { + let cmp = make_bitmap_key_comparator(&data_type); + let min_key = serialize_bitmap_datum(&min, &data_type); + let max_key = serialize_bitmap_datum(&max, &data_type); + let outside_key = serialize_bitmap_datum(&outside, &data_type); + let nan_key = serialize_bitmap_datum(&nan, &data_type); + let meta = BTreeIndexMeta::new(Some(min_key.clone()), Some(max_key), false); + + assert!(!bitmap_meta_may_match( + &meta, + PredicateOperator::Eq, + &data_type, + std::slice::from_ref(&outside_key), + cmp.as_ref(), + )); + assert!(!bitmap_meta_may_match( + &meta, + PredicateOperator::In, + &data_type, + std::slice::from_ref(&outside_key), + cmp.as_ref(), + )); + assert!(!bitmap_meta_may_match( + &meta, + PredicateOperator::IsNull, + &data_type, + &[], + cmp.as_ref(), + )); + assert!(bitmap_meta_may_match( + &meta, + PredicateOperator::IsNotNull, + &data_type, + &[], + cmp.as_ref(), + )); + + let nan_meta = BTreeIndexMeta::new(Some(min_key), Some(nan_key.clone()), false); + assert!(bitmap_meta_may_match( + &nan_meta, + PredicateOperator::Eq, + &data_type, + std::slice::from_ref(&nan_key), + cmp.as_ref(), + )); + assert!(bitmap_meta_may_match( + &nan_meta, + PredicateOperator::In, + &data_type, + std::slice::from_ref(&nan_key), + cmp.as_ref(), + )); + + assert!(bitmap_meta_may_match( + &meta, + PredicateOperator::Gt, + &data_type, + std::slice::from_ref(&outside_key), + cmp.as_ref(), + )); + assert!(bitmap_meta_may_match_between( + &meta, + &data_type, + &outside_key, + &outside_key, + cmp.as_ref(), + )); + + let only_nulls = BTreeIndexMeta::new(None, None, true); + assert!(bitmap_meta_may_match( + &only_nulls, + PredicateOperator::IsNull, + &data_type, + &[], + cmp.as_ref(), + )); + assert!(!bitmap_meta_may_match( + &only_nulls, + PredicateOperator::IsNotNull, + &data_type, + &[], + cmp.as_ref(), + )); + assert!(!bitmap_meta_may_match( + &only_nulls, + PredicateOperator::NotEq, + &data_type, + std::slice::from_ref(&outside_key), + cmp.as_ref(), + )); + assert!(!bitmap_meta_may_match_between( + &only_nulls, + &data_type, + &outside_key, + &outside_key, + cmp.as_ref(), + )); +} + +#[test] +fn test_bitmap_floating_meta_prunes_equality_and_fails_open_for_ranges() { + assert_bitmap_floating_meta_policy( + DataType::Float(crate::spec::FloatType::new()), + Datum::Float(-1.0), + Datum::Float(1.0), + Datum::Float(2.0), + Datum::Float(f32::NAN), + ); + assert_bitmap_floating_meta_policy( + DataType::Double(crate::spec::DoubleType::new()), + Datum::Double(-1.0), + Datum::Double(1.0), + Datum::Double(2.0), + Datum::Double(f64::NAN), + ); +} + +#[test] +fn test_row_range_index_merges_overlapping() { + let idx = RowRangeIndex::create(vec![ + RowRange::new(0, 5), + RowRange::new(3, 10), + RowRange::new(20, 30), + ]); + assert_eq!(idx.ranges().len(), 2); + assert_eq!(idx.ranges()[0], RowRange::new(0, 10)); + assert_eq!(idx.ranges()[1], RowRange::new(20, 30)); +} + +#[test] +fn test_row_range_index_merges_adjacent() { + let idx = RowRangeIndex::create(vec![RowRange::new(0, 5), RowRange::new(6, 10)]); + assert_eq!(idx.ranges().len(), 1); + assert_eq!(idx.ranges()[0], RowRange::new(0, 10)); +} + +#[test] +fn test_row_range_index_intersects() { + let idx = RowRangeIndex::create(vec![RowRange::new(10, 20), RowRange::new(30, 40)]); + assert!(idx.intersects(15, 25)); + assert!(idx.intersects(5, 10)); + assert!(idx.intersects(20, 30)); + assert!(!idx.intersects(0, 9)); + assert!(!idx.intersects(21, 29)); + assert!(!idx.intersects(41, 50)); +} + +#[test] +fn test_row_range_index_intersected_ranges() { + let idx = RowRangeIndex::create(vec![ + RowRange::new(10, 20), + RowRange::new(30, 40), + RowRange::new(50, 60), + ]); + let result = idx.intersected_ranges(15, 55); + assert_eq!( + result, + vec![ + RowRange::new(15, 20), + RowRange::new(30, 40), + RowRange::new(50, 55), + ] + ); +} + +#[test] +fn test_row_range_index_intersection_row_count() { + let idx = RowRangeIndex::create(vec![ + RowRange::new(10, 20), + RowRange::new(30, 40), + RowRange::new(50, 60), + ]); + + assert_eq!(idx.intersection_row_count(15, 55), 23); + assert_eq!(idx.intersection_row_count(21, 29), 0); + assert_eq!(idx.intersection_row_count(55, 15), 0); +} + +#[test] +fn test_search_limit_with_deleted_rows_expands_and_caps() { + let idx = RowRangeIndex::create(vec![RowRange::new(2, 4), RowRange::new(8, 10)]); + + assert_eq!(search_limit_with_deleted_rows(5, 0, 19, Some(&idx)), 11); + assert_eq!(search_limit_with_deleted_rows(18, 0, 19, Some(&idx)), 20); + assert_eq!(search_limit_with_deleted_rows(5, 0, 19, None), 5); +} + +#[test] +fn test_row_range_index_empty() { + let idx = RowRangeIndex::create(Vec::new()); + assert!(!idx.intersects(0, 100)); + assert!(idx.intersected_ranges(0, 100).is_empty()); +} + +fn le_int_key(v: i32) -> Vec { + v.to_le_bytes().to_vec() +} + +/// Set up a temp dir with `index/{file_name}` containing the btree testdata file, +/// and return (FileIO, table_path, file_name, _tmp_dir). +fn setup_testdata_table(testdata_name: &str) -> (FileIO, String, String, tempfile::TempDir) { + let src = format!( + "{}/testdata/btree/{testdata_name}", + env!("CARGO_MANIFEST_DIR") + ); + let tmp = tempfile::tempdir().unwrap(); + let index_dir = tmp.path().join("index"); + std::fs::create_dir_all(&index_dir).unwrap(); + std::fs::copy(&src, index_dir.join(testdata_name)).unwrap(); + + let table_path = format!("file://{}", tmp.path().display()); + let file_io = crate::io::FileIOBuilder::new("file").build().unwrap(); + (file_io, table_path, testdata_name.to_string(), tmp) +} + +type BitmapTestdataTable = (FileIO, String, String, BTreeIndexMeta, tempfile::TempDir); + +fn setup_bitmap_testdata_table(file_name: &str) -> BitmapTestdataTable { + let src = format!("{}/testdata/bitmap/{file_name}", env!("CARGO_MANIFEST_DIR")); + let meta_src = format!( + "{}/testdata/bitmap/{file_name}.meta", + env!("CARGO_MANIFEST_DIR") + ); + let tmp = tempfile::tempdir().unwrap(); + let index_dir = tmp.path().join("index"); + std::fs::create_dir_all(&index_dir).unwrap(); + std::fs::copy(&src, index_dir.join(file_name)).unwrap(); + let meta = BTreeIndexMeta::deserialize(&std::fs::read(meta_src).unwrap()).unwrap(); + + let table_path = format!("file://{}", tmp.path().display()); + let file_io = crate::io::FileIOBuilder::new("file").build().unwrap(); + (file_io, table_path, file_name.to_string(), meta, tmp) +} + +fn setup_java_bitmap_testdata_table() -> BitmapTestdataTable { + setup_bitmap_testdata_table("bitmap_varchar_java.index") +} + +fn make_global_index_entry( + file_name: &str, + field_id: i32, + row_range_start: i64, + row_range_end: i64, + meta: &BTreeIndexMeta, +) -> crate::spec::IndexManifestEntry { + make_global_index_entry_with_type( + BTREE_GLOBAL_INDEX_TYPE, + file_name, + field_id, + row_range_start, + row_range_end, + meta, + ) +} + +fn make_global_index_entry_with_type( + index_type: &str, + file_name: &str, + field_id: i32, + row_range_start: i64, + row_range_end: i64, + meta: &BTreeIndexMeta, +) -> crate::spec::IndexManifestEntry { + use crate::spec::{GlobalIndexMeta, IndexFileMeta}; + IndexManifestEntry { + version: 1, + kind: FileKind::Add, + partition: vec![], + bucket: 0, + index_file: IndexFileMeta { + index_type: index_type.to_string(), + file_name: file_name.to_string(), + file_size: 0, + row_count: 0, + deletion_vectors_ranges: None, + global_index_meta: Some(GlobalIndexMeta { + row_range_start, + row_range_end, + index_field_id: field_id, + extra_field_ids: None, + source_meta: None, + index_meta: Some(meta.serialize()), + }), + }, + } +} + +async fn make_fm_index_entry( + file_name: &str, + field_id: i32, + row_range_start: i64, + row_range_end: i64, + first_row_id: u64, + row_count: u64, +) -> IndexManifestEntry { + let output = VecFileWrite::new(); + let mut writer = FMGlobalIndexWriter::new( + Box::new(output.clone()), + FMWriteOptions { + compression: BlockCompressionType::None, + ..FMWriteOptions::default() + }, + ) + .unwrap(); + for row_id in first_row_id..first_row_id + row_count { + writer.write(Some(b"value"), row_id).await.unwrap(); + } + let result = writer.finish().await.unwrap(); + IndexManifestEntry { + version: 1, + kind: FileKind::Add, + partition: vec![], + bucket: 0, + index_file: IndexFileMeta { + index_type: FM_GLOBAL_INDEX_TYPE.to_string(), + file_name: file_name.to_string(), + file_size: output.to_vec().len() as i64, + row_count: result.row_count as i64, + deletion_vectors_ranges: None, + global_index_meta: Some(crate::spec::GlobalIndexMeta { + row_range_start, + row_range_end, + index_field_id: field_id, + extra_field_ids: None, + source_meta: None, + index_meta: Some(result.index_meta), + }), + }, + } +} + +#[tokio::test] +async fn test_fm_file_set_must_exactly_cover_source_range() { + let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + let first = make_fm_index_entry("first.fm", 1, 10, 13, 0, 2).await; + let error = match GlobalIndexScanner::create( + &file_io, + "memory:/table", + 1, + i64::MAX, + i64::MAX, + std::slice::from_ref(&first), + &string_schema_fields(), + ) { + Err(error) => error, + Ok(_) => panic!("a partial Java FM file set must fail closed"), + }; + assert!( + matches!(error, Error::DataInvalid { message, .. } if message.contains("cover 2 rows, expected 4")) + ); + + let second = make_fm_index_entry("second.fm", 1, 10, 13, 2, 2).await; + assert!(GlobalIndexScanner::create( + &file_io, + "memory:/table", + 1, + i64::MAX, + i64::MAX, + &[first, second], + &string_schema_fields(), + ) + .unwrap() + .is_some()); +} + +#[test] +fn test_mixed_fm_and_btree_select_compatible_index_family() { + let btree = GlobalIndexEntry { + file_name: "name.btree".to_string(), + index_type: GlobalIndexFileKind::BTree, + file_size: 1, + row_range_start: 0, + row_range_end: 9, + meta: GlobalIndexEntryMeta::Sorted(BTreeIndexMeta::new(None, None, false)), + }; + let fm = GlobalIndexEntry { + file_name: "name.fm".to_string(), + index_type: GlobalIndexFileKind::FM, + file_size: 1, + row_range_start: 0, + row_range_end: 9, + meta: GlobalIndexEntryMeta::FM { + bytes: Vec::new(), + first_row_id: 0, + row_count: 10, + }, + }; + let entries = [btree, fm]; + let data_type = DataType::VarChar(crate::spec::VarCharType::string_type()); + let literal = [Datum::String("needle".to_string())]; + + let contains = [(PredicateOperator::Contains, literal.as_slice(), &data_type)]; + let selected = select_entries_for_predicates(&entries, &contains); + assert_eq!(selected.len(), 1); + assert_eq!(selected[0].index_type, GlobalIndexFileKind::FM); + + let equals = [(PredicateOperator::Eq, literal.as_slice(), &data_type)]; + let selected = select_entries_for_predicates(&entries, &equals); + assert_eq!(selected.len(), 1); + assert_eq!(selected[0].index_type, GlobalIndexFileKind::BTree); +} + +fn int_schema_fields() -> Vec { + vec![DataField::new( + 1, + "id".to_string(), + DataType::Int(crate::spec::IntType::new()), + )] +} + +fn string_schema_fields() -> Vec { + vec![DataField::new( + 1, + "name".to_string(), + DataType::VarChar(crate::spec::VarCharType::string_type()), + )] +} + +async fn evaluate_global_index_fast( + file_io: &FileIO, + table_path: &str, + entries: &[IndexManifestEntry], + predicates: &[Predicate], + fields: &[DataField], +) -> Result>> { + evaluate_global_index_fast_with_fallback_size( + file_io, + table_path, + entries, + predicates, + fields, + i64::MAX, + i64::MAX, + ) + .await +} + +async fn evaluate_global_index_fast_with_fallback_size( + file_io: &FileIO, + table_path: &str, + entries: &[IndexManifestEntry], + predicates: &[Predicate], + fields: &[DataField], + btree_fallback_scan_max_size: i64, + bitmap_fallback_scan_max_size: i64, +) -> Result>> { + super::evaluate_global_index(super::GlobalIndexEvaluation { + file_io, + table_path, + index_entries: entries, + predicates, + schema_fields: fields, + search_mode: GlobalIndexSearchMode::Fast, + global_index_thread_num: 32, + btree_fallback_scan_max_size, + bitmap_fallback_scan_max_size, + fm_read_options: FMReadOptions::default(), + next_row_id: None, + data_ranges: &[], + }) + .await +} + +fn two_field_schema_fields() -> Vec { + vec![ + DataField::new( + 1, + "id".to_string(), + DataType::Int(crate::spec::IntType::new()), + ), + DataField::new( + 2, + "value".to_string(), + DataType::Int(crate::spec::IntType::new()), + ), + ] +} + +fn int_eq(column: &str, index: usize, value: i32) -> Predicate { + Predicate::Leaf { + column: column.to_string(), + index, + data_type: DataType::Int(crate::spec::IntType::new()), + op: PredicateOperator::Eq, + literals: vec![Datum::Int(value)], + } +} + +#[test] +fn test_unindexed_ranges_fast_mode_empty() { + let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + let meta = BTreeIndexMeta::new(None, None, false); + let entries = vec![make_global_index_entry("idx", 1, 0, 49, &meta)]; + let fields = int_schema_fields(); + let scanner = GlobalIndexScanner::create( + &file_io, + "memory:/t", + 32, + i64::MAX, + i64::MAX, + &entries, + &fields, + ) + .expect("create scanner") + .expect("scanner"); + + let ranges = scanner + .unindexed_ranges( + &int_eq("id", 0, 7), + GlobalIndexSearchMode::Fast, + Some(100), + &[RowRange::new(50, 99)], + ) + .unwrap(); + assert!(ranges.is_empty()); +} + +#[test] +fn test_unindexed_ranges_full_uses_snapshot_next_row_id() { + let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + let meta = BTreeIndexMeta::new(None, None, false); + let entries = vec![make_global_index_entry("idx", 1, 0, 49, &meta)]; + let fields = int_schema_fields(); + let scanner = GlobalIndexScanner::create( + &file_io, + "memory:/t", + 32, + i64::MAX, + i64::MAX, + &entries, + &fields, + ) + .expect("create scanner") + .expect("scanner"); + + let ranges = scanner + .unindexed_ranges( + &int_eq("id", 0, 7), + GlobalIndexSearchMode::Full, + Some(100), + &[], + ) + .unwrap(); + assert_eq!(ranges, vec![RowRange::new(50, 99)]); +} + +#[test] +fn test_unindexed_ranges_detail_uses_data_file_ranges() { + let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + let meta = BTreeIndexMeta::new(None, None, false); + let entries = vec![make_global_index_entry("idx", 1, 0, 49, &meta)]; + let fields = int_schema_fields(); + let scanner = GlobalIndexScanner::create( + &file_io, + "memory:/t", + 32, + i64::MAX, + i64::MAX, + &entries, + &fields, + ) + .expect("create scanner") + .expect("scanner"); + + let ranges = scanner + .unindexed_ranges( + &int_eq("id", 0, 7), + GlobalIndexSearchMode::Detail, + Some(100), + &[ + RowRange::new(0, 10), + RowRange::new(40, 60), + RowRange::new(80, 90), + ], + ) + .unwrap(); + assert_eq!(ranges, vec![RowRange::new(50, 60), RowRange::new(80, 90)]); +} + +#[test] +fn test_unindexed_ranges_uses_all_predicate_field_coverage() { + let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + let meta = BTreeIndexMeta::new(None, None, false); + let entries = vec![ + make_global_index_entry("idx_id", 1, 0, 49, &meta), + make_global_index_entry("idx_value", 2, 0, 99, &meta), + ]; + let fields = two_field_schema_fields(); + let scanner = GlobalIndexScanner::create( + &file_io, + "memory:/t", + 32, + i64::MAX, + i64::MAX, + &entries, + &fields, + ) + .expect("create scanner") + .expect("scanner"); + let predicate = Predicate::and(vec![int_eq("id", 0, 7), int_eq("value", 1, 8)]); + + let ranges = scanner + .unindexed_ranges(&predicate, GlobalIndexSearchMode::Full, Some(100), &[]) + .unwrap(); + assert_eq!(ranges, vec![RowRange::new(50, 99)]); +} + +#[test] +fn test_unindexed_ranges_missing_field_coverage_reads_all_data_ranges() { + let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + let meta = BTreeIndexMeta::new(None, None, false); + let entries = vec![make_global_index_entry("idx_id", 1, 0, 49, &meta)]; + let fields = two_field_schema_fields(); + let scanner = GlobalIndexScanner::create( + &file_io, + "memory:/t", + 32, + i64::MAX, + i64::MAX, + &entries, + &fields, + ) + .expect("create scanner") + .expect("scanner"); + let predicate = Predicate::and(vec![int_eq("id", 0, 7), int_eq("value", 1, 8)]); + + let ranges = scanner + .unindexed_ranges(&predicate, GlobalIndexSearchMode::Full, Some(100), &[]) + .unwrap(); + assert_eq!(ranges, vec![RowRange::new(0, 99)]); +} + +#[test] +fn test_unindexed_ranges_counts_extra_field_coverage() { + let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + let meta = BTreeIndexMeta::new(None, None, false); + let mut entry = make_global_index_entry("idx_id_value", 1, 0, 99, &meta); + entry + .index_file + .global_index_meta + .as_mut() + .unwrap() + .extra_field_ids = Some(vec![2]); + let fields = two_field_schema_fields(); + let scanner = GlobalIndexScanner::create( + &file_io, + "memory:/t", + 32, + i64::MAX, + i64::MAX, + &[entry], + &fields, + ) + .expect("create scanner") + .expect("scanner"); + + let ranges = scanner + .unindexed_ranges( + &int_eq("value", 1, 8), + GlobalIndexSearchMode::Full, + Some(100), + &[], + ) + .unwrap(); + assert!(ranges.is_empty()); +} + +#[tokio::test] +async fn test_evaluate_extra_field_only_without_composite_reader_falls_back() { + let (file_io, table_path, file_name, _tmp) = + setup_testdata_table("btree_int_100_no_compress.bin"); + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + let mut entry = make_global_index_entry(&file_name, 1, 0, 99, &meta); + entry + .index_file + .global_index_meta + .as_mut() + .unwrap() + .extra_field_ids = Some(vec![2]); + let fields = two_field_schema_fields(); + let predicates = vec![int_eq("value", 1, 50)]; + + let result = evaluate_global_index_fast(&file_io, &table_path, &[entry], &predicates, &fields) + .await + .unwrap(); + assert!( + result.is_none(), + "extra-field-only predicates must fall back until composite-key btree reads are supported" + ); +} + +#[tokio::test] +async fn test_evaluate_global_index_eq() { + let (file_io, table_path, file_name, tmp) = + setup_testdata_table("btree_int_100_no_compress.bin"); + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + let mut entry = make_global_index_entry(&file_name, 1, 0, 99, &meta); + entry.index_file.file_size = std::fs::metadata(tmp.path().join("index").join(&file_name)) + .unwrap() + .len() as i64; + let entries = vec![entry]; + let fields = int_schema_fields(); + + // key=50 -> row_id=25, offset by row_range_start=0 -> global row_id=25 + let predicates = vec![Predicate::Leaf { + column: "id".to_string(), + index: 0, + data_type: DataType::Int(crate::spec::IntType::new()), + op: PredicateOperator::Eq, + literals: vec![Datum::Int(50)], + }]; + + let result = evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) + .await + .unwrap(); + let ranges = result.unwrap(); + assert_eq!(ranges, vec![RowRange::new(25, 25)]); +} + +#[tokio::test] +async fn test_evaluate_global_index_uses_known_file_size() { + let (file_io, table_path, file_name, _tmp) = + setup_testdata_table("btree_int_100_no_compress.bin"); + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + let mut entry = make_global_index_entry(&file_name, 1, 0, 99, &meta); + entry.index_file.file_size = 1; + + let error = evaluate_global_index_fast( + &file_io, + &table_path, + &[entry], + &[int_eq("id", 0, 50)], + &int_schema_fields(), + ) + .await + .expect_err("the known file size should be used without a metadata lookup"); + + assert!(matches!( + error, + crate::Error::DataInvalid { message, .. } + if message.contains("Failed to open BTree index file") + )); +} + +#[tokio::test] +async fn test_missing_index_meta_returns_error() { + let (file_io, table_path, file_name, tmp) = + setup_testdata_table("btree_int_100_no_compress.bin"); + let second_file_name = "btree_int_100_no_compress_2.bin"; + std::fs::copy( + tmp.path().join("index").join(&file_name), + tmp.path().join("index").join(second_file_name), + ) + .unwrap(); + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + let valid_entry = make_global_index_entry(&file_name, 1, 0, 99, &meta); + let mut invalid_entry = make_global_index_entry(second_file_name, 1, 100, 199, &meta); + invalid_entry + .index_file + .global_index_meta + .as_mut() + .unwrap() + .index_meta = None; + + let error = evaluate_global_index_fast( + &file_io, + &table_path, + &[valid_entry, invalid_entry], + &[int_eq("id", 0, 50)], + &int_schema_fields(), + ) + .await + .expect_err("missing sorted index metadata must fail the scan"); + + assert!(matches!( + error, + crate::Error::DataInvalid { message, .. } + if message.contains(second_file_name) + )); +} + +#[tokio::test] +async fn test_invalid_index_meta_returns_error() { + let (file_io, table_path, file_name, tmp) = + setup_testdata_table("btree_int_100_no_compress.bin"); + let second_file_name = "btree_int_100_no_compress_2.bin"; + std::fs::copy( + tmp.path().join("index").join(&file_name), + tmp.path().join("index").join(second_file_name), + ) + .unwrap(); + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + let valid_entry = make_global_index_entry(&file_name, 1, 0, 99, &meta); + let mut invalid_entry = make_global_index_entry(second_file_name, 1, 100, 199, &meta); + let mut invalid_meta = vec![0; 9]; + invalid_meta[..4].copy_from_slice(&10i32.to_le_bytes()); + invalid_entry + .index_file + .global_index_meta + .as_mut() + .unwrap() + .index_meta = Some(invalid_meta); + + let error = evaluate_global_index_fast( + &file_io, + &table_path, + &[valid_entry, invalid_entry], + &[int_eq("id", 0, 50)], + &int_schema_fields(), + ) + .await + .expect_err("invalid sorted index metadata must fail the scan"); + + assert!(matches!( + error, + crate::Error::DataInvalid { + message, + source: Some(_), + } if message.contains(second_file_name) + )); +} + +#[tokio::test] +async fn test_missing_global_index_meta_returns_error() { + let (file_io, table_path, file_name, tmp) = + setup_testdata_table("btree_int_100_no_compress.bin"); + let second_file_name = "btree_int_100_no_compress_2.bin"; + std::fs::copy( + tmp.path().join("index").join(&file_name), + tmp.path().join("index").join(second_file_name), + ) + .unwrap(); + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + let valid_entry = make_global_index_entry(&file_name, 1, 0, 99, &meta); + let mut invalid_entry = make_global_index_entry(second_file_name, 1, 100, 199, &meta); + invalid_entry.index_file.global_index_meta = None; + + let error = evaluate_global_index_fast( + &file_io, + &table_path, + &[valid_entry, invalid_entry], + &[int_eq("id", 0, 50)], + &int_schema_fields(), + ) + .await + .expect_err("missing global index metadata must fail the scan"); + + assert!(matches!( + error, + crate::Error::DataInvalid { message, .. } + if message.contains(second_file_name) + )); +} + +#[tokio::test] +async fn test_evaluate_java_bitmap_golden_index_eq_and_null() { + let data_type = DataType::VarChar(crate::spec::VarCharType::string_type()); + let (file_io, table_path, file_name, meta, tmp) = setup_java_bitmap_testdata_table(); + let mut entry = + make_global_index_entry_with_type(BITMAP_GLOBAL_INDEX_TYPE, &file_name, 1, 100, 109, &meta); + entry.index_file.file_size = std::fs::metadata(tmp.path().join("index").join(&file_name)) + .unwrap() + .len() as i64; + let entries = vec![entry]; + let fields = string_schema_fields(); + assert_eq!(meta.first_key, Some(b"alpha".to_vec())); + assert_eq!(meta.last_key, Some(b"office".to_vec())); + assert!(meta.has_nulls); + + let eq_predicates = vec![Predicate::Leaf { + column: "name".to_string(), + index: 0, + data_type: data_type.clone(), + op: PredicateOperator::Eq, + literals: vec![Datum::String("k2".to_string())], + }]; + let eq_result = + evaluate_global_index_fast(&file_io, &table_path, &entries, &eq_predicates, &fields) + .await + .unwrap(); + assert_eq!(eq_result.unwrap(), vec![RowRange::new(105, 106)]); + + let null_predicates = vec![Predicate::Leaf { + column: "name".to_string(), + index: 0, + data_type, + op: PredicateOperator::IsNull, + literals: vec![], + }]; + let null_result = + evaluate_global_index_fast(&file_io, &table_path, &entries, &null_predicates, &fields) + .await + .unwrap(); + assert_eq!(null_result.unwrap(), vec![RowRange::new(104, 104)]); +} + +async fn assert_bitmap_int_fixture(file_name: &str) { + let data_type = DataType::Int(crate::spec::IntType::new()); + let (file_io, table_path, file_name, meta, _tmp) = setup_bitmap_testdata_table(file_name); + let entries = vec![make_global_index_entry_with_type( + BITMAP_GLOBAL_INDEX_TYPE, + &file_name, + 1, + 100, + 105, + &meta, + )]; + let fields = int_schema_fields(); + assert_eq!(meta.first_key, Some(le_int_key(-1))); + assert_eq!(meta.last_key, Some(le_int_key(256))); + assert!(meta.has_nulls); + + let cases = [ + ( + PredicateOperator::Eq, + vec![Datum::Int(0)], + vec![RowRange::new(101, 102)], + ), + ( + PredicateOperator::Eq, + vec![Datum::Int(256)], + vec![RowRange::new(104, 104)], + ), + ( + PredicateOperator::In, + vec![Datum::Int(-1), Datum::Int(1), Datum::Int(256)], + vec![RowRange::new(100, 100), RowRange::new(103, 104)], + ), + ( + PredicateOperator::NotEq, + vec![Datum::Int(0)], + vec![RowRange::new(100, 100), RowRange::new(103, 104)], + ), + ( + PredicateOperator::NotIn, + vec![Datum::Int(-1), Datum::Int(1), Datum::Int(256)], + vec![RowRange::new(101, 102)], + ), + ( + PredicateOperator::IsNull, + vec![], + vec![RowRange::new(105, 105)], + ), + ]; + + for (op, literals, expected) in cases { + let predicates = vec![Predicate::Leaf { + column: "id".to_string(), + index: 0, + data_type: data_type.clone(), + op, + literals, + }]; + let result = + evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) + .await + .unwrap() + .unwrap(); + assert_eq!(result, expected, "{file_name}: {op}"); + } +} + +#[tokio::test] +async fn test_evaluate_java_logical_order_bitmap_int_fixture() { + assert_bitmap_int_fixture("bitmap_int_logical_java.index").await; +} + +async fn assert_bitmap_nan_equality_uses_direct_lookup( + data_type: DataType, + nan_literals: [Datum; 3], + zero: Datum, +) { + let output = VecFileWrite::new(); + let captured = output.clone(); + let mut writer = BitmapGlobalIndexWriter::new( + Box::new(output), + 1, + BlockCompressionType::None, + make_bitmap_key_comparator(&data_type), + ); + for (row_id, literal) in nan_literals.iter().enumerate() { + let key = serialize_bitmap_datum(literal, &data_type); + writer.write(Some(&key), row_id as i64).unwrap(); + } + let zero_key = serialize_bitmap_datum(&zero, &data_type); + writer.write(Some(&zero_key), 3).unwrap(); + let write_result = writer.finish().await.unwrap(); + let bytes = captured.to_vec(); + + let tmp = tempfile::tempdir().unwrap(); + let index_dir = tmp.path().join("index"); + std::fs::create_dir_all(&index_dir).unwrap(); + let file_name = "bitmap-current.index"; + std::fs::write(index_dir.join(file_name), &bytes).unwrap(); + let table_path = format!("file://{}", tmp.path().display()); + let file_io = crate::io::FileIOBuilder::new("file").build().unwrap(); + + let mut entry = make_global_index_entry_with_type( + BITMAP_GLOBAL_INDEX_TYPE, + file_name, + 1, + 100, + 103, + &write_result.meta, + ); + entry.index_file.file_size = bytes.len() as i64; + let entries = vec![entry]; + let fields = vec![DataField::new(1, "id".to_string(), data_type.clone())]; + let cases = [ + (PredicateOperator::Eq, vec![nan_literals[0].clone()]), + ( + PredicateOperator::In, + vec![nan_literals[1].clone(), nan_literals[2].clone()], + ), + ]; + + for (op, literals) in cases { + let predicates = vec![Predicate::Leaf { + column: "id".to_string(), + index: 0, + data_type: data_type.clone(), + op, + literals, + }]; + let result = evaluate_global_index_fast_with_fallback_size( + &file_io, + &table_path, + &entries, + &predicates, + &fields, + i64::MAX, + 0, + ) + .await + .unwrap() + .unwrap(); + assert_eq!(result, vec![RowRange::new(100, 102)], "{data_type:?}: {op}"); + } +} + +#[tokio::test] +async fn test_bitmap_nan_equality_uses_direct_lookup_with_fallback_scan_disabled() { + assert_bitmap_nan_equality_uses_direct_lookup( + DataType::Float(crate::spec::FloatType::new()), + [ + Datum::Float(f32::from_bits(0xffc0_0001)), + Datum::Float(f32::from_bits(0x7fc0_0010)), + Datum::Float(f32::NAN), + ], + Datum::Float(0.0), + ) + .await; + assert_bitmap_nan_equality_uses_direct_lookup( + DataType::Double(crate::spec::DoubleType::new()), + [ + Datum::Double(f64::from_bits(0xfff8_0000_0000_0001)), + Datum::Double(f64::from_bits(0x7ff8_0000_0000_0010)), + Datum::Double(f64::NAN), + ], + Datum::Double(0.0), + ) + .await; +} + +fn legacy_floating_comparator(data_type: &DataType) -> BoxedCmp { + match data_type { + DataType::Float(_) => Box::new(|left, right| { + let left = f32::from_le_bytes(left.try_into().unwrap()); + let right = f32::from_le_bytes(right.try_into().unwrap()); + left.total_cmp(&right) + }), + DataType::Double(_) => Box::new(|left, right| { + let left = f64::from_le_bytes(left.try_into().unwrap()); + let right = f64::from_le_bytes(right.try_into().unwrap()); + left.total_cmp(&right) + }), + _ => unreachable!("legacy floating comparator requires Float or Double"), + } +} + +async fn assert_legacy_floating_btree( + file_name: &str, + data_type: DataType, + nan_keys: Vec>, + nan_literals: Vec, + zero_key: Vec, + zero_literal: Datum, +) { + let mut rows = nan_keys + .into_iter() + .enumerate() + .map(|(row_id, key)| (key, row_id as i64)) + .collect::>(); + rows.push((zero_key, 3)); + let cmp = legacy_floating_comparator(&data_type); + rows.sort_by(|left, right| cmp(&left.0, &right.0)); + let expected_first_key = rows.first().unwrap().0.clone(); + let expected_last_key = rows.last().unwrap().0.clone(); + + let output = VecFileWrite::new(); + let captured = output.clone(); + let mut writer = + BTreeIndexWriter::with_comparator(Box::new(output), 1, BlockCompressionType::None, cmp); + for (key, row_id) in rows { + writer.write(Some(&key), row_id).await.unwrap(); + } + let write_result = writer.finish().await.unwrap(); + assert_eq!(write_result.meta.first_key, Some(expected_first_key)); + assert_eq!(write_result.meta.last_key, Some(expected_last_key)); + + let tmp = tempfile::tempdir().unwrap(); + let index_dir = tmp.path().join("index"); + std::fs::create_dir_all(&index_dir).unwrap(); + std::fs::write(index_dir.join(file_name), captured.to_vec()).unwrap(); + let table_path = format!("file://{}", tmp.path().display()); + let file_io = crate::io::FileIOBuilder::new("file").build().unwrap(); + let entries = vec![make_global_index_entry( + file_name, + 1, + 100, + 103, + &write_result.meta, + )]; + let fields = vec![DataField::new(1, "id".to_string(), data_type.clone())]; + let cases = [ + ( + PredicateOperator::Eq, + vec![zero_literal.clone()], + vec![RowRange::new(103, 103)], + ), + ( + PredicateOperator::Eq, + vec![nan_literals[0].clone()], + vec![RowRange::new(100, 100)], + ), + ( + PredicateOperator::In, + vec![ + nan_literals[0].clone(), + nan_literals[1].clone(), + zero_literal, + ], + vec![RowRange::new(100, 101), RowRange::new(103, 103)], + ), + ]; + + for (op, literals, expected) in cases { + let predicates = vec![Predicate::Leaf { + column: "id".to_string(), + index: 0, + data_type: data_type.clone(), + op, + literals, + }]; + let result = + evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) + .await + .unwrap() + .unwrap(); + assert_eq!(result, expected, "{file_name}: {op}"); + } +} + +#[tokio::test] +async fn test_evaluate_legacy_float_btree() { + let nan_bits = [0xffc0_0001u32, 0xffc0_0010, 0xffff_1234]; + assert_legacy_floating_btree( + "btree_float_legacy_rust.index", + DataType::Float(crate::spec::FloatType::new()), + nan_bits + .iter() + .map(|bits| bits.to_le_bytes().to_vec()) + .collect(), + nan_bits + .iter() + .map(|bits| Datum::Float(f32::from_bits(*bits))) + .collect(), + 0.0f32.to_le_bytes().to_vec(), + Datum::Float(0.0), + ) + .await; +} + +#[tokio::test] +async fn test_evaluate_legacy_double_btree() { + let nan_bits = [ + 0xfff8_0000_0000_0001u64, + 0xfff8_0000_0000_0010, + 0xffff_1234_5678_9abc, + ]; + assert_legacy_floating_btree( + "btree_double_legacy_rust.index", + DataType::Double(crate::spec::DoubleType::new()), + nan_bits + .iter() + .map(|bits| bits.to_le_bytes().to_vec()) + .collect(), + nan_bits + .iter() + .map(|bits| Datum::Double(f64::from_bits(*bits))) + .collect(), + 0.0f64.to_le_bytes().to_vec(), + Datum::Double(0.0), + ) + .await; +} + +#[tokio::test] +async fn test_evaluate_java_bitmap_golden_index_string_fallback_scan() { + let data_type = DataType::VarChar(crate::spec::VarCharType::string_type()); + let (file_io, table_path, file_name, meta, tmp) = setup_java_bitmap_testdata_table(); + let file_size = std::fs::metadata(tmp.path().join("index").join(&file_name)) + .unwrap() + .len() as i64; + let mut entry = + make_global_index_entry_with_type(BITMAP_GLOBAL_INDEX_TYPE, &file_name, 1, 100, 109, &meta); + entry.index_file.file_size = file_size; + let entries = vec![entry]; + let fields = string_schema_fields(); + + let ends_with_predicates = vec![Predicate::Leaf { + column: "name".to_string(), + index: 0, + data_type: data_type.clone(), + op: PredicateOperator::EndsWith, + literals: vec![Datum::String("ta".to_string())], + }]; + let ends_with_result = evaluate_global_index_fast( + &file_io, + &table_path, + &entries, + &ends_with_predicates, + &fields, + ) + .await + .unwrap(); + assert_eq!( + ends_with_result.unwrap(), + vec![RowRange::new(101, 101), RowRange::new(103, 103)] + ); + + let contains_predicates = vec![Predicate::Leaf { + column: "name".to_string(), + index: 0, + data_type: data_type.clone(), + op: PredicateOperator::Contains, + literals: vec![Datum::String("ph".to_string())], + }]; + let contains_result = evaluate_global_index_fast( + &file_io, + &table_path, + &entries, + &contains_predicates, + &fields, + ) + .await + .unwrap(); + assert_eq!( + contains_result.unwrap(), + vec![RowRange::new(100, 100), RowRange::new(102, 102)] + ); + + let like_predicates = vec![Predicate::Leaf { + column: "name".to_string(), + index: 0, + data_type: data_type.clone(), + op: PredicateOperator::Like, + literals: vec![Datum::String("%ha%".to_string())], + }]; + let like_result = + evaluate_global_index_fast(&file_io, &table_path, &entries, &like_predicates, &fields) + .await + .unwrap(); + assert_eq!( + like_result.unwrap(), + vec![RowRange::new(100, 100), RowRange::new(102, 102)] + ); + + let less_than_predicates = vec![Predicate::Leaf { + column: "name".to_string(), + index: 0, + data_type: data_type.clone(), + op: PredicateOperator::Lt, + literals: vec![Datum::String("delta".to_string())], + }]; + let less_than_result = evaluate_global_index_fast( + &file_io, + &table_path, + &entries, + &less_than_predicates, + &fields, + ) + .await + .unwrap(); + assert_eq!(less_than_result.unwrap(), vec![RowRange::new(100, 102)]); + + let mut over_limit_entry = + make_global_index_entry_with_type(BITMAP_GLOBAL_INDEX_TYPE, &file_name, 1, 100, 109, &meta); + over_limit_entry.index_file.file_size = file_size; + let over_limit_entries = vec![over_limit_entry]; + let over_limit_less_than = evaluate_global_index_fast_with_fallback_size( + &file_io, + &table_path, + &over_limit_entries, + &less_than_predicates, + &fields, + i64::MAX, + file_size - 1, + ) + .await + .unwrap(); + assert!( + over_limit_less_than.is_none(), + "range predicates require fallback dictionary scans and should be unsupported over budget" + ); + + let no_match_contains = vec![Predicate::Leaf { + column: "name".to_string(), + index: 0, + data_type: data_type.clone(), + op: PredicateOperator::Contains, + literals: vec![Datum::String("zz".to_string())], + }]; + let over_limit_result = evaluate_global_index_fast_with_fallback_size( + &file_io, + &table_path, + &over_limit_entries, + &no_match_contains, + &fields, + i64::MAX, + file_size - 1, + ) + .await + .unwrap(); + assert!( + over_limit_result.is_none(), + "fallback scans over budget should be unsupported instead of returning full coverage" + ); + + let direct_with_over_limit_fallback = vec![Predicate::and(vec![ + Predicate::Leaf { + column: "name".to_string(), + index: 0, + data_type: data_type.clone(), + op: PredicateOperator::Eq, + literals: vec![Datum::String("k2".to_string())], + }, + Predicate::Leaf { + column: "name".to_string(), + index: 0, + data_type, + op: PredicateOperator::Contains, + literals: vec![Datum::String("zz".to_string())], + }, + ])]; + let direct_result = evaluate_global_index_fast_with_fallback_size( + &file_io, + &table_path, + &over_limit_entries, + &direct_with_over_limit_fallback, + &fields, + i64::MAX, + file_size - 1, + ) + .await + .unwrap(); + assert_eq!(direct_result.unwrap(), vec![RowRange::new(105, 106)]); +} + +#[tokio::test] +async fn test_btree_fallback_scan_over_limit_is_unsupported() { + let (file_io, table_path, file_name, tmp) = + setup_testdata_table("btree_varchar_100_no_compress.bin"); + let meta = BTreeIndexMeta::new(Some(b"a".to_vec()), Some(b"yyyy".to_vec()), false); + let fields = string_schema_fields(); + let data_type = DataType::VarChar(crate::spec::VarCharType::string_type()); + let predicates = vec![Predicate::Leaf { + column: "name".to_string(), + index: 0, + data_type, + op: PredicateOperator::Contains, + literals: vec![Datum::String("not-present".to_string())], + }]; + + let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; + let exact_result = evaluate_global_index_fast_with_fallback_size( + &file_io, + &table_path, + &entries, + &predicates, + &fields, + i64::MAX, + i64::MAX, + ) + .await + .unwrap(); + assert_eq!(exact_result.unwrap(), Vec::::new()); + + let mut over_limit_entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; + over_limit_entries[0].index_file.file_size = 2; + let over_limit_result = evaluate_global_index_fast_with_fallback_size( + &file_io, + &table_path, + &over_limit_entries, + &predicates, + &fields, + 1, + i64::MAX, + ) + .await + .unwrap(); + assert!( + over_limit_result.is_none(), + "fallback scans over budget should be unsupported instead of returning full coverage" + ); + + let second_file_name = "btree_varchar_100_no_compress_2.bin"; + std::fs::copy( + tmp.path().join("index").join(&file_name), + tmp.path().join("index").join(second_file_name), + ) + .unwrap(); + let mut first = make_global_index_entry(&file_name, 1, 0, 99, &meta); + first.index_file.file_size = 1; + let mut second = make_global_index_entry(second_file_name, 1, 100, 199, &meta); + second.index_file.file_size = 1; + let total_over_limit_result = evaluate_global_index_fast_with_fallback_size( + &file_io, + &table_path, + &[first, second], + &predicates, + &fields, + 1, + i64::MAX, + ) + .await + .unwrap(); + assert!( + total_over_limit_result.is_none(), + "fallback budget should use selected files' total size, not per-file size" + ); +} + +#[tokio::test] +async fn test_fallback_scan_over_limit_with_mixed_index_kinds_is_unsupported() { + let (file_io, table_path, file_name, _tmp) = + setup_testdata_table("btree_varchar_100_no_compress.bin"); + let btree_meta = BTreeIndexMeta::new(Some(b"a".to_vec()), Some(b"yyyy".to_vec()), false); + let bitmap_meta = BTreeIndexMeta::new(Some(b"m".to_vec()), Some(b"z".to_vec()), false); + let fields = string_schema_fields(); + let predicates = vec![Predicate::Leaf { + column: "name".to_string(), + index: 0, + data_type: DataType::VarChar(crate::spec::VarCharType::string_type()), + op: PredicateOperator::Lt, + literals: vec![Datum::String("delta".to_string())], + }]; + + let mut btree = make_global_index_entry_with_type( + BTREE_GLOBAL_INDEX_TYPE, + &file_name, + 1, + 0, + 99, + &btree_meta, + ); + btree.index_file.file_size = 2; + let mut bitmap = make_global_index_entry_with_type( + BITMAP_GLOBAL_INDEX_TYPE, + "bitmap-no-match.index", + 1, + 100, + 199, + &bitmap_meta, + ); + bitmap.index_file.file_size = 1; + + let result = evaluate_global_index_fast_with_fallback_size( + &file_io, + &table_path, + &[btree, bitmap], + &predicates, + &fields, + 1, + i64::MAX, + ) + .await + .unwrap(); + assert!( + result.is_none(), + "an over-budget selected BTree file must stay unsupported even if bitmap files are pruned by metadata" + ); +} + +#[tokio::test] +async fn test_fallback_preflight_happens_before_shard_io() { + let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + let table_path = "memory:/missing-index-files"; + let meta = BTreeIndexMeta::new(Some(b"a".to_vec()), Some(b"z".to_vec()), false); + let fields = string_schema_fields(); + let predicates = vec![Predicate::Leaf { + column: "name".to_string(), + index: 0, + data_type: DataType::VarChar(crate::spec::VarCharType::string_type()), + op: PredicateOperator::Contains, + literals: vec![Datum::String("middle".to_string())], + }]; + + let mut btree = make_global_index_entry_with_type( + BTREE_GLOBAL_INDEX_TYPE, + "missing-btree.index", + 1, + 0, + 99, + &meta, + ); + btree.index_file.file_size = 1; + let mut bitmap = make_global_index_entry_with_type( + BITMAP_GLOBAL_INDEX_TYPE, + "missing-bitmap.index", + 1, + 100, + 199, + &meta, + ); + bitmap.index_file.file_size = 1; + + let result = evaluate_global_index_fast_with_fallback_size( + &file_io, + table_path, + &[btree, bitmap], + &predicates, + &fields, + 1, + 0, + ) + .await + .expect("fallback must be decided before opening an earlier shard"); + + assert!(result.is_none()); +} + +#[tokio::test] +async fn test_evaluate_global_index_full_mode_includes_unindexed_tail() { + let (file_io, table_path, file_name, _tmp) = + setup_testdata_table("btree_int_100_no_compress.bin"); + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; + let fields = int_schema_fields(); + let predicates = vec![int_eq("id", 0, 50)]; + + let result = super::evaluate_global_index(super::GlobalIndexEvaluation { + file_io: &file_io, + table_path: &table_path, + index_entries: &entries, + predicates: &predicates, + schema_fields: &fields, + search_mode: GlobalIndexSearchMode::Full, + global_index_thread_num: 32, + btree_fallback_scan_max_size: i64::MAX, + bitmap_fallback_scan_max_size: i64::MAX, + fm_read_options: FMReadOptions::default(), + next_row_id: Some(150), + data_ranges: &[], + }) + .await + .unwrap(); + + assert_eq!( + result.unwrap(), + vec![RowRange::new(25, 25), RowRange::new(100, 149)] + ); +} + +#[tokio::test] +async fn test_evaluate_global_index_and_uses_evaluated_field_coverage_for_raw_fallback() { + let src = format!( + "{}/testdata/btree/btree_int_100_no_compress.bin", + env!("CARGO_MANIFEST_DIR") + ); + let tmp = tempfile::tempdir().unwrap(); + let index_dir = tmp.path().join("index"); + std::fs::create_dir_all(&index_dir).unwrap(); + std::fs::copy(&src, index_dir.join("index_part1.bin")).unwrap(); + std::fs::copy(&src, index_dir.join("index_part2.bin")).unwrap(); + + let table_path = format!("file://{}", tmp.path().display()); + let file_io = crate::io::FileIOBuilder::new("file").build().unwrap(); + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + + let mut first = make_global_index_entry("index_part1.bin", 1, 0, 49, &meta); + first + .index_file + .global_index_meta + .as_mut() + .unwrap() + .extra_field_ids = Some(vec![2]); + let second = make_global_index_entry("index_part2.bin", 1, 50, 99, &meta); + let entries = vec![first, second]; + let fields = two_field_schema_fields(); + + let predicates = vec![Predicate::and(vec![ + int_eq("id", 0, 50), + int_eq("value", 1, 8), + ])]; + let result = super::evaluate_global_index(super::GlobalIndexEvaluation { + file_io: &file_io, + table_path: &table_path, + index_entries: &entries, + predicates: &predicates, + schema_fields: &fields, + search_mode: GlobalIndexSearchMode::Full, + global_index_thread_num: 32, + btree_fallback_scan_max_size: i64::MAX, + bitmap_fallback_scan_max_size: i64::MAX, + fm_read_options: FMReadOptions::default(), + next_row_id: Some(100), + data_ranges: &[], + }) + .await + .unwrap(); + + assert_eq!( + result.unwrap(), + vec![RowRange::new(25, 25), RowRange::new(75, 75)], + "raw fallback should use only the id field that was actually evaluated; \ + the unevaluated extra field must not widen or narrow fallback coverage" + ); +} + +#[tokio::test] +async fn test_evaluate_global_index_detail_mode_uses_data_ranges() { + let (file_io, table_path, file_name, _tmp) = + setup_testdata_table("btree_int_100_no_compress.bin"); + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; + let fields = int_schema_fields(); + let predicates = vec![int_eq("id", 0, 50)]; + + let data_ranges = [RowRange::new(90, 120), RowRange::new(140, 145)]; + let result = super::evaluate_global_index(super::GlobalIndexEvaluation { + file_io: &file_io, + table_path: &table_path, + index_entries: &entries, + predicates: &predicates, + schema_fields: &fields, + search_mode: GlobalIndexSearchMode::Detail, + global_index_thread_num: 32, + btree_fallback_scan_max_size: i64::MAX, + bitmap_fallback_scan_max_size: i64::MAX, + fm_read_options: FMReadOptions::default(), + next_row_id: Some(150), + data_ranges: &data_ranges, + }) + .await + .unwrap(); + + assert_eq!( + result.unwrap(), + vec![ + RowRange::new(25, 25), + RowRange::new(100, 120), + RowRange::new(140, 145), + ] + ); +} + +#[tokio::test] +async fn test_evaluate_global_index_range() { + let (file_io, table_path, file_name, _tmp) = + setup_testdata_table("btree_int_100_no_compress.bin"); + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; + let fields = int_schema_fields(); + + // keys 10..=20 -> keys 10,12,14,16,18,20 -> row_ids 5,6,7,8,9,10 + let predicates = vec![ + Predicate::Leaf { + column: "id".to_string(), + index: 0, + data_type: DataType::Int(crate::spec::IntType::new()), + op: PredicateOperator::GtEq, + literals: vec![Datum::Int(10)], + }, + Predicate::Leaf { + column: "id".to_string(), + index: 0, + data_type: DataType::Int(crate::spec::IntType::new()), + op: PredicateOperator::LtEq, + literals: vec![Datum::Int(20)], + }, + ]; + + let result = evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) + .await + .unwrap(); + let ranges = result.unwrap(); + assert_eq!(ranges, vec![RowRange::new(5, 10)]); + + let mut over_limit_entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; + over_limit_entries[0].index_file.file_size = 2; + let over_limit_result = evaluate_global_index_fast_with_fallback_size( + &file_io, + &table_path, + &over_limit_entries, + &predicates, + &fields, + 1, + i64::MAX, + ) + .await + .unwrap(); + assert!( + over_limit_result.is_none(), + "between/range predicates require fallback scans and should be unsupported over budget" + ); +} + +#[tokio::test] +async fn test_evaluate_global_index_in() { + let (file_io, table_path, file_name, _tmp) = + setup_testdata_table("btree_int_100_no_compress.bin"); + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; + let fields = int_schema_fields(); + + // IN(0, 50, 198) -> row_ids 0, 25, 99 + let predicates = vec![Predicate::Leaf { + column: "id".to_string(), + index: 0, + data_type: DataType::Int(crate::spec::IntType::new()), + op: PredicateOperator::In, + literals: vec![Datum::Int(0), Datum::Int(50), Datum::Int(198)], + }]; + + let result = evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) + .await + .unwrap(); + let ranges = result.unwrap(); + assert_eq!( + ranges, + vec![ + RowRange::new(0, 0), + RowRange::new(25, 25), + RowRange::new(99, 99) + ] + ); +} + +#[tokio::test] +async fn test_evaluate_global_index_no_match() { + let (file_io, table_path, file_name, _tmp) = + setup_testdata_table("btree_int_100_no_compress.bin"); + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; + let fields = int_schema_fields(); + + // key=999 doesn't exist + let predicates = vec![Predicate::Leaf { + column: "id".to_string(), + index: 0, + data_type: DataType::Int(crate::spec::IntType::new()), + op: PredicateOperator::Eq, + literals: vec![Datum::Int(999)], + }]; + + let result = evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) + .await + .unwrap(); + let ranges = result.unwrap(); + assert!(ranges.is_empty()); +} + +#[tokio::test] +async fn test_evaluate_global_index_with_row_offset() { + let (file_io, table_path, file_name, _tmp) = + setup_testdata_table("btree_int_100_no_compress.bin"); + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + // row_range_start=1000 simulates an offset + let entries = vec![make_global_index_entry(&file_name, 1, 1000, 1099, &meta)]; + let fields = int_schema_fields(); + + // key=50 -> local row_id=25, offset -> global row_id=1025 + let predicates = vec![Predicate::Leaf { + column: "id".to_string(), + index: 0, + data_type: DataType::Int(crate::spec::IntType::new()), + op: PredicateOperator::Eq, + literals: vec![Datum::Int(50)], + }]; + + let result = evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) + .await + .unwrap(); + let ranges = result.unwrap(); + assert_eq!(ranges, vec![RowRange::new(1025, 1025)]); +} + +#[tokio::test] +async fn test_evaluate_global_index_unknown_column() { + let (file_io, table_path, file_name, _tmp) = + setup_testdata_table("btree_int_100_no_compress.bin"); + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; + let fields = int_schema_fields(); + + // Column "unknown" not in schema -> None (can't evaluate) + let predicates = vec![Predicate::Leaf { + column: "unknown".to_string(), + index: 0, + data_type: DataType::Int(crate::spec::IntType::new()), + op: PredicateOperator::Eq, + literals: vec![Datum::Int(50)], + }]; + + let result = evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) + .await + .unwrap(); + assert!(result.is_none()); +} + +#[tokio::test] +async fn test_evaluate_global_index_multi_field_and() { + // Two fields, each with its own btree index file (same data, different field_id). + // btree_int_100_no_compress.bin: keys 0,2,4,...,198 -> row_ids 0,1,...,99 + let src = format!( + "{}/testdata/btree/btree_int_100_no_compress.bin", + env!("CARGO_MANIFEST_DIR") + ); + let tmp = tempfile::tempdir().unwrap(); + let index_dir = tmp.path().join("index"); + std::fs::create_dir_all(&index_dir).unwrap(); + std::fs::copy(&src, index_dir.join("index_field1.bin")).unwrap(); + std::fs::copy(&src, index_dir.join("index_field2.bin")).unwrap(); + + let table_path = format!("file://{}", tmp.path().display()); + let file_io = crate::io::FileIOBuilder::new("file").build().unwrap(); + + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + + let fields = vec![ + DataField::new( + 1, + "id".to_string(), + DataType::Int(crate::spec::IntType::new()), + ), + DataField::new( + 2, + "value".to_string(), + DataType::Int(crate::spec::IntType::new()), + ), + ]; + + let entries = vec![ + make_global_index_entry("index_field1.bin", 1, 0, 99, &meta), + make_global_index_entry("index_field2.bin", 2, 0, 99, &meta), + ]; + + // id >= 40 AND id <= 60 → keys 40,42,...,60 → row_ids 20..30 + // value >= 44 AND value <= 52 → keys 44,46,48,50,52 → row_ids 22..26 + // AND → intersect [20..30] and [22..26] = [22..26] + let predicates = vec![ + Predicate::Leaf { + column: "id".to_string(), + index: 0, + data_type: DataType::Int(crate::spec::IntType::new()), + op: PredicateOperator::GtEq, + literals: vec![Datum::Int(40)], + }, + Predicate::Leaf { + column: "id".to_string(), + index: 0, + data_type: DataType::Int(crate::spec::IntType::new()), + op: PredicateOperator::LtEq, + literals: vec![Datum::Int(60)], + }, + Predicate::Leaf { + column: "value".to_string(), + index: 1, + data_type: DataType::Int(crate::spec::IntType::new()), + op: PredicateOperator::GtEq, + literals: vec![Datum::Int(44)], + }, + Predicate::Leaf { + column: "value".to_string(), + index: 1, + data_type: DataType::Int(crate::spec::IntType::new()), + op: PredicateOperator::LtEq, + literals: vec![Datum::Int(52)], + }, + ]; + + let result = evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) + .await + .unwrap(); + let ranges = result.unwrap(); + assert_eq!(ranges, vec![RowRange::new(22, 26)]); +} + +#[tokio::test] +async fn test_multi_field_and_shares_query_concurrency_budget() { + let src = format!( + "{}/testdata/btree/btree_int_100_no_compress.bin", + env!("CARGO_MANIFEST_DIR") + ); + let tmp = tempfile::tempdir().unwrap(); + let index_dir = tmp.path().join("index"); + std::fs::create_dir_all(&index_dir).unwrap(); + let file_names: Vec<_> = (1..=4) + .map(|field_id| { + let file_name = format!("index_field{field_id}.bin"); + std::fs::copy(&src, index_dir.join(&file_name)).unwrap(); + file_name + }) + .collect(); + + let table_path = format!("file://{}", tmp.path().display()); + let file_io = crate::io::FileIOBuilder::new("file").build().unwrap(); + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + let fields: Vec<_> = (0..4) + .map(|index| { + let field_id = index + 1; + DataField::new( + field_id, + format!("field{field_id}"), + DataType::Int(crate::spec::IntType::new()), + ) + }) + .collect(); + let entries: Vec<_> = file_names + .iter() + .enumerate() + .map(|(index, file_name)| { + make_global_index_entry(file_name, index as i32 + 1, 0, 99, &meta) + }) + .collect(); + let predicate = Predicate::and( + (0..4) + .map(|index| int_eq(&format!("field{}", index + 1), index, 50)) + .collect(), + ); + + for (thread_num, expected_peak) in [(1, 1), (2, 2)] { + let mut scanner = GlobalIndexScanner::create( + &file_io, + &table_path, + thread_num, + i64::MAX, + i64::MAX, + &entries, + &fields, + ) + .unwrap() + .unwrap(); + let probe = Arc::new(QueryIoProbe::default()); + scanner.query_io_probe = Some(Arc::clone(&probe)); + + let result = scanner.evaluate(&predicate).await.unwrap().unwrap(); + + assert_eq!(result.row_ranges, vec![RowRange::new(25, 25)]); + assert_eq!(result.evaluated_field_ids, HashSet::from([1, 2, 3, 4])); + assert_eq!(probe.peak(), expected_peak); + } +} + +/// Regression for the Between+remaining bug in `evaluate_leaf`. When a +/// native `Between` leaf is paired with another conjunct (e.g. `id >= 0`), +/// and the file's b-tree key range falls **outside** the Between range +/// but is still matched by the remaining predicate, the whole AND must +/// produce zero rows. Before the fix, `file_result` was initialized from +/// the remaining predicate's bitmap and the Between conjunct was silently +/// dropped — the test would observe the file's full row id set instead of +/// the empty set. +#[tokio::test] +async fn test_between_unmatched_file_drops_remaining_match() { + let (file_io, table_path, file_name, _tmp) = + setup_testdata_table("btree_int_100_no_compress.bin"); + // File covers keys [0, 198] (row_ids 0..99). Pick a Between range + // entirely below 0 so `may_match_between` is false, and a `>= 0` + // conjunct that would otherwise scoop up every row in the file. + let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)), false); + let entries = vec![make_global_index_entry(&file_name, 1, 0, 99, &meta)]; + let fields = int_schema_fields(); + + let predicates = vec![Predicate::and(vec![ + Predicate::Leaf { + column: "id".to_string(), + index: 0, + data_type: DataType::Int(crate::spec::IntType::new()), + op: PredicateOperator::Between, + literals: vec![Datum::Int(-100), Datum::Int(-50)], + }, + Predicate::Leaf { + column: "id".to_string(), + index: 0, + data_type: DataType::Int(crate::spec::IntType::new()), + op: PredicateOperator::GtEq, + literals: vec![Datum::Int(0)], + }, + ])]; + + let result = evaluate_global_index_fast(&file_io, &table_path, &entries, &predicates, &fields) + .await + .unwrap(); + let ranges = result.unwrap(); + assert!( + ranges.is_empty(), + "Between(-100..-50) AND id>=0 must produce zero rows on a file \ + whose key range is [0, 198] — got {ranges:?}" + ); +} diff --git a/crates/paimon/src/table/lumina_index_build_builder.rs b/crates/paimon/src/table/lumina_index_build_builder.rs index d340c97eb..e50a7311c 100644 --- a/crates/paimon/src/table/lumina_index_build_builder.rs +++ b/crates/paimon/src/table/lumina_index_build_builder.rs @@ -15,28 +15,26 @@ // specific language governing permissions and limitations // under the License. -use crate::lumina::ffi::LuminaBuilder; +mod extraction; +mod planning; +mod validation; +mod writer; + +use extraction::extract_vectors; +use planning::plan_lumina_shards; +use validation::{ + effective_lumina_options, find_index_field, resolve_lumina_options, validate_table_options, + validate_vector_field, +}; +use writer::abort_on_build_error; + use crate::lumina::{ is_lumina_index_type, LuminaIndexMeta, LuminaVectorIndexOptions, LUMINA_IDENTIFIER, }; -use crate::spec::{ - bucket_dir_name, BinaryRow, CoreOptions, DataField, DataFileMeta, DataType, FileKind, - GlobalIndexMeta, IndexFileMeta, ROW_ID_FIELD_NAME, -}; -use crate::table::source::exclude_row_ranges; -use crate::table::{ - CommitMessage, DataSplitBuilder, RowRange, SnapshotManager, Table, TableCommit, -}; +use crate::spec::CoreOptions; +use crate::table::{CommitMessage, RowRange, SnapshotManager, Table, TableCommit}; use crate::{Error, Result}; -use arrow_array::{Array, FixedSizeListArray, Float32Array, Int64Array, ListArray, RecordBatch}; -use bytes::Bytes; -use futures::TryStreamExt; use std::collections::HashMap; -use std::path::{Path, PathBuf}; -use tokio::io::AsyncReadExt; - -const INDEX_DIR: &str = "index"; -const COPY_BUFFER_SIZE: usize = 1024 * 1024; pub struct LuminaIndexBuildBuilder<'a> { table: &'a Table, @@ -159,1910 +157,42 @@ impl<'a> LuminaIndexBuildBuilder<'a> { ) .await?; + let commit = TableCommit::new( + self.table.clone(), + format!( + "global-index-{}-create-{}", + LUMINA_IDENTIFIER, + uuid::Uuid::new_v4() + ), + ); let shard_count = shards.len(); let mut messages = Vec::with_capacity(shard_count); for shard in shards { - let vectors = extract_vectors(self.table, &shard, index_column, dimension).await?; - let index_file = self - .build_index_file( + let build_result = async { + let vectors = extract_vectors(self.table, &shard, index_column, dimension).await?; + self.build_index_file( &shard, &vectors, dimension, index_field.id(), index_meta.clone(), ) - .await?; + .await + } + .await; + let index_file = abort_on_build_error(&commit, &messages, build_result).await?; let mut message = CommitMessage::new(shard.partition_bytes.clone(), 0, vec![]); message.new_index_files = vec![index_file]; messages.push(message); } - TableCommit::new( - self.table.clone(), - format!( - "global-index-{}-create-{}", - LUMINA_IDENTIFIER, - uuid::Uuid::new_v4() - ), - ) - .commit_if_latest_snapshot(messages, snapshot.id()) - .await?; - - Ok(shard_count) - } - - async fn build_index_file( - &self, - shard: &LuminaIndexShard, - vectors: &[f32], - dimension: i32, - index_field_id: i32, - index_meta: Vec, - ) -> Result { - let row_count = checked_row_count(shard.row_range_start, shard.row_range_end)?; - // The native Lumina builder counts rows in an i32; the manifest keeps the - // full width. - let native_row_count = i32::try_from(row_count).map_err(|_| Error::DataInvalid { - message: format!( - "Lumina shard row count {row_count} exceeds what the native builder accepts" - ), - source: None, - })?; - validate_vector_buffer(vectors, native_row_count, dimension)?; - let ids = (0..row_count as u64).collect::>(); - let native_options = LuminaIndexMeta::deserialize(&index_meta)?.options().clone(); - - let temp_path = temp_lumina_path(); - let temp_file = TempFileGuard::new(temp_path.clone()); - let temp_path_str = temp_path.to_string_lossy().to_string(); - let builder = LuminaBuilder::create(&native_options)?; - builder.pretrain(vectors, native_row_count, dimension)?; - builder.insert(vectors, &ids, native_row_count, dimension)?; - builder.dump(&temp_path_str)?; - - let file_name = format!("lumina-global-index-{}.index", uuid::Uuid::new_v4()); - self.table - .file_io() - .mkdirs(&format!( - "{}/{INDEX_DIR}/", - self.table.location().trim_end_matches('/') - )) - .await?; - let index_path = format!( - "{}/{INDEX_DIR}/{}", - self.table.location().trim_end_matches('/'), - file_name - ); - copy_local_file_to_output(&temp_path, self.table.file_io().new_output(&index_path)?) + commit + .commit_if_latest_snapshot(messages, snapshot.id()) .await?; - temp_file.cleanup(); - - let status = self.table.file_io().get_status(&index_path).await?; - Ok(IndexFileMeta { - index_type: LUMINA_IDENTIFIER.to_string(), - file_name, - file_size: checked_i64( - status.size, - "Index file is too large for Rust IndexFileMeta", - )?, - row_count, - deletion_vectors_ranges: None, - global_index_meta: Some(GlobalIndexMeta { - row_range_start: shard.row_range_start, - row_range_end: shard.row_range_end, - index_field_id, - extra_field_ids: None, - source_meta: None, - index_meta: Some(index_meta), - }), - }) - } -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct LuminaIndexShard { - pub partition: BinaryRow, - pub partition_bytes: Vec, - pub files: Vec, - pub row_range_start: i64, - pub row_range_end: i64, - snapshot_id: i64, - source_bucket: i32, - total_buckets: i32, - bucket_path: String, -} - -fn validate_table_options(table: &Table, core_options: &CoreOptions) -> Result<()> { - if !table.schema().primary_keys().is_empty() { - return Err(Error::Unsupported { - message: "Lumina index build does not support primary-key tables".to_string(), - }); - } - if !core_options.row_tracking_enabled() { - return Err(Error::DataInvalid { - message: "Lumina index build requires 'row-tracking.enabled' = 'true'".to_string(), - source: None, - }); - } - if !core_options.data_evolution_enabled() { - return Err(Error::DataInvalid { - message: "Lumina index build requires 'data-evolution.enabled' = 'true'".to_string(), - source: None, - }); - } - if !core_options.global_index_enabled() { - return Err(Error::DataInvalid { - message: "Lumina index build requires 'global-index.enabled' = 'true'".to_string(), - source: None, - }); - } - if core_options.deletion_vectors_enabled() { - return Err(Error::Unsupported { - message: - "Lumina index build does not support tables with deletion-vectors.enabled=true" - .to_string(), - }); - } - Ok(()) -} - -fn find_index_field<'a>(table: &'a Table, column: &str) -> Result<&'a DataField> { - table - .schema() - .fields() - .iter() - .find(|field| field.name() == column) - .ok_or_else(|| Error::ColumnNotExist { - full_name: table.identifier().full_name(), - column: column.to_string(), - }) -} - -fn validate_vector_field(field: &DataField) -> Result<()> { - let is_array_float = matches!( - field.data_type(), - DataType::Array(array) if matches!(array.element_type(), DataType::Float(_)) - ); - let is_vector_float = matches!( - field.data_type(), - DataType::Vector(vector) if matches!(vector.element_type(), DataType::Float(_)) - ); - if !is_array_float && !is_vector_float { - return Err(Error::DataInvalid { - message: format!( - "Lumina index requires ARRAY or VECTOR column, got {:?} for column '{}'", - field.data_type(), - field.name() - ), - source: None, - }); - } - Ok(()) -} - -/// For a `VECTOR` column, ensure the effective Lumina options carry -/// `lumina.index.dimension = N`, so the native options and serialized index -/// metadata match the column type. Absent → inject N; present-and-equal → keep; -/// present-and-different → ConfigInvalid. Non-vector columns (e.g. ARRAY) -/// are returned unchanged so the existing option/default behavior is preserved. -fn effective_lumina_options( - field: &DataField, - mut resolved: HashMap, -) -> Result> { - let DataType::Vector(vector) = field.data_type() else { - return Ok(resolved); - }; - let n = vector.length().to_string(); - match resolved.get(crate::lumina::LUMINA_DIMENSION_OPTION) { - None => { - resolved.insert(crate::lumina::LUMINA_DIMENSION_OPTION.to_string(), n); - } - Some(existing) if *existing == n => {} - Some(existing) => { - return Err(Error::ConfigInvalid { - message: format!( - "Vector column '{}' has dimension {} from its type, but '{}' is set to '{}'. \ - Remove the option or set it to {}.", - field.name(), - n, - crate::lumina::LUMINA_DIMENSION_OPTION, - existing, - n - ), - }); - } - } - Ok(resolved) -} - -fn resolve_lumina_options( - table_options: &HashMap, - user_options: &HashMap, -) -> Result> { - let mut options = table_options.clone(); - options.extend(user_options.clone()); - LuminaVectorIndexOptions::new(&options)?; - Ok(options) -} - -#[allow(clippy::too_many_arguments)] -fn plan_lumina_shards( - table_location: &str, - partition_keys: &[String], - schema_fields: &[DataField], - core_options: &CoreOptions, - snapshot_id: i64, - entries: Vec, - rows_per_shard: i64, - indexed: &[RowRange], -) -> Result> { - if rows_per_shard <= 0 { - return Err(Error::DataInvalid { - message: format!( - "Option 'global-index.row-count-per-shard' must be greater than 0, got: {rows_per_shard}" - ), - source: None, - }); - } - - let mut by_partition_bucket: HashMap<(Vec, i32, i32), Vec> = HashMap::new(); - for entry in entries { - if *entry.kind() != FileKind::Add { - continue; - } - if entry.file().first_row_id.is_none() { - return Err(Error::DataInvalid { - message: format!( - "Data file '{}' is missing first_row_id; cannot build a complete Lumina index", - entry.file().file_name - ), - source: None, - }); - } - let (partition, bucket, total_buckets, file) = entry.into_parts(); - by_partition_bucket - .entry((partition, bucket, total_buckets)) - .or_default() - .push(file); - } - - let mut result = Vec::new(); - for ((partition_bytes, source_bucket, total_buckets), files) in by_partition_bucket { - let partition = if partition_keys.is_empty() { - BinaryRow::new(0) - } else { - BinaryRow::from_serialized_bytes(&partition_bytes)? - }; - let bucket_path = bucket_path( - table_location, - partition_keys, - schema_fields, - core_options, - &partition, - source_bucket, - )?; - let mut files_by_shard: HashMap> = HashMap::new(); - for file in files { - let (file_start, file_end) = file.row_id_range().ok_or_else(|| Error::DataInvalid { - message: format!( - "Data file '{}' is missing first_row_id; cannot build a complete Lumina index", - file.file_name - ), - source: None, - })?; - let start_shard = file_start / rows_per_shard; - let end_shard = file_end / rows_per_shard; - for shard_id in start_shard..=end_shard { - files_by_shard - .entry(shard_id * rows_per_shard) - .or_default() - .push(file.clone()); - } - } - - let mut shard_starts = files_by_shard.keys().copied().collect::>(); - shard_starts.sort_unstable(); - for shard_start in shard_starts { - let shard_end = shard_start + rows_per_shard - 1; - let mut shard_files = files_by_shard.remove(&shard_start).unwrap_or_default(); - shard_files.sort_by_key(|file| file.first_row_id); - let groups = group_contiguous_files(shard_files)?; - for group in groups { - let group_start = group - .first() - .and_then(|file| file.first_row_id) - .expect("planned groups are non-empty and row-id assigned"); - let group_end = group - .iter() - .map(|file| file.row_id_range().unwrap().1) - .max() - .unwrap(); - // Coverage of this group clamped to the current shard cell. Then - // subtract the already-indexed ranges so the build only covers - // the gap. Because grid-clamp and gap-subtraction are both range - // intersections, applying the gap here is equivalent to btree's - // "exclude then split" -- and each surviving segment stays inside - // one shard cell, preserving per-shard row-id contiguity. - let coverage_start = group_start.max(shard_start); - let coverage_end = group_end.min(shard_end); - let build_segments = - exclude_row_ranges(&[RowRange::new(coverage_start, coverage_end)], indexed); - for seg in build_segments { - result.push(LuminaIndexShard { - partition: partition.clone(), - partition_bytes: partition_bytes.clone(), - files: group.clone(), - row_range_start: seg.from(), - row_range_end: seg.to(), - snapshot_id, - source_bucket, - total_buckets, - bucket_path: bucket_path.clone(), - }); - } - } - } - } - result.sort_by(|a, b| { - a.partition - .to_serialized_bytes() - .cmp(&b.partition.to_serialized_bytes()) - .then(a.source_bucket.cmp(&b.source_bucket)) - .then(a.row_range_start.cmp(&b.row_range_start)) - }); - Ok(result) -} - -fn group_contiguous_files(mut files: Vec) -> Result>> { - if files.is_empty() { - return Ok(Vec::new()); - } - files.sort_by_key(|file| file.first_row_id); - let mut groups = Vec::new(); - let mut current = Vec::new(); - let mut current_end = None; - for file in files { - let (file_start, file_end) = file.row_id_range().ok_or_else(|| Error::DataInvalid { - message: format!( - "Data file '{}' is missing first_row_id; cannot build a complete Lumina index", - file.file_name - ), - source: None, - })?; - match current_end { - None => { - current.push(file); - current_end = Some(file_end); - } - Some(end) if file_start <= end + 1 => { - current.push(file); - current_end = Some(end.max(file_end)); - } - Some(_) => { - groups.push(std::mem::take(&mut current)); - current.push(file); - current_end = Some(file_end); - } - } - } - if !current.is_empty() { - groups.push(current); - } - Ok(groups) -} - -fn bucket_path( - table_location: &str, - partition_keys: &[String], - schema_fields: &[DataField], - core_options: &CoreOptions, - partition: &BinaryRow, - bucket: i32, -) -> Result { - let base = table_location.trim_end_matches('/'); - if partition_keys.is_empty() { - return Ok(format!("{base}/{}", bucket_dir_name(bucket))); - } - let computer = crate::spec::PartitionComputer::new( - partition_keys, - schema_fields, - core_options.partition_default_name(), - core_options.legacy_partition_name(), - )?; - Ok(format!( - "{base}/{}{}", - computer.generate_partition_path(partition)?, - bucket_dir_name(bucket) - )) -} - -async fn extract_vectors( - table: &Table, - shard: &LuminaIndexShard, - index_column: &str, - dimension: i32, -) -> Result> { - let split = DataSplitBuilder::new() - .with_snapshot(shard.snapshot_id) - .with_partition(shard.partition.clone()) - .with_bucket(shard.source_bucket) - .with_bucket_path(shard.bucket_path.clone()) - .with_total_buckets(shard.total_buckets) - .with_data_files(shard.files.clone()) - .with_row_ranges(vec![RowRange::new( - shard.row_range_start, - shard.row_range_end, - )]) - .build()?; - - let mut read_builder = table.new_read_builder(); - read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; - let read = read_builder.new_read()?; - let batches = read.to_arrow(&[split])?.try_collect::>().await?; - extract_vectors_from_batches( - &batches, - index_column, - dimension, - shard.row_range_start, - checked_row_count(shard.row_range_start, shard.row_range_end)?, - ) -} - -fn extract_vectors_from_batches( - batches: &[RecordBatch], - index_column: &str, - dimension: i32, - row_range_start: i64, - expected_row_count: i64, -) -> Result> { - let dimension = usize::try_from(dimension).map_err(|e| Error::DataInvalid { - message: format!("Invalid Lumina dimension: {dimension}"), - source: Some(Box::new(e)), - })?; - let row_count = batches.iter().map(RecordBatch::num_rows).sum::(); - let mut vectors = Vec::with_capacity(row_count * dimension); - let mut expected_row_id = row_range_start; - for batch in batches { - let vector_index = - batch - .schema() - .index_of(index_column) - .map_err(|e| Error::DataInvalid { - message: format!("Vector column '{index_column}' not found in read batch: {e}"), - source: None, - })?; - let row_id_index = - batch - .schema() - .index_of(ROW_ID_FIELD_NAME) - .map_err(|e| Error::DataInvalid { - message: format!("_ROW_ID column not found in read batch: {e}"), - source: None, - })?; - // Resolve the vector column as either List (ARRAY) or - // FixedSizeList (VECTOR). Both yield a Float32Array of - // values plus a per-row [start, end) slice. - let column = batch.column(vector_index); - enum VectorLayout<'a> { - List(&'a ListArray), - Fixed(&'a FixedSizeListArray), - } - let layout = if let Some(a) = column.as_any().downcast_ref::() { - VectorLayout::List(a) - } else if let Some(a) = column.as_any().downcast_ref::() { - VectorLayout::Fixed(a) - } else { - return Err(Error::DataInvalid { - message: "Lumina vector extraction requires Arrow List or FixedSizeList".to_string(), - source: None, - }); - }; - let values = match layout { - VectorLayout::List(a) => a.values(), - VectorLayout::Fixed(a) => a.values(), - } - .as_any() - .downcast_ref::() - .ok_or_else(|| Error::DataInvalid { - message: "Lumina vector extraction requires Float32 vector elements".to_string(), - source: None, - })?; - let row_ids = batch - .column(row_id_index) - .as_any() - .downcast_ref::() - .ok_or_else(|| Error::DataInvalid { - message: "Lumina vector extraction requires non-null Int64 _ROW_ID".to_string(), - source: None, - })?; - - for row in 0..batch.num_rows() { - if row_ids.is_null(row) { - return Err(Error::DataInvalid { - message: "Lumina vector extraction found null _ROW_ID".to_string(), - source: None, - }); - } - let row_id = row_ids.value(row); - if row_id != expected_row_id { - return Err(Error::DataInvalid { - message: format!( - "Lumina vector extraction expected _ROW_ID {}, got {}", - expected_row_id, row_id - ), - source: None, - }); - } - expected_row_id += 1; - - let is_null = match layout { - VectorLayout::List(a) => a.is_null(row), - VectorLayout::Fixed(a) => a.is_null(row), - }; - if is_null { - return Err(Error::DataInvalid { - message: "Lumina vector extraction found null vector row".to_string(), - source: None, - }); - } - let (start, end) = match layout { - VectorLayout::List(a) => { - let offsets = a.value_offsets(); - (offsets[row] as usize, offsets[row + 1] as usize) - } - VectorLayout::Fixed(a) => { - let len = a.value_length() as usize; - (row * len, (row + 1) * len) - } - }; - if end - start != dimension { - return Err(Error::DataInvalid { - message: format!( - "Lumina vector dimension mismatch: expected {}, got {}", - dimension, - end - start - ), - source: None, - }); - } - for value_index in start..end { - if values.is_null(value_index) { - return Err(Error::DataInvalid { - message: "Lumina vector extraction found null vector element".to_string(), - source: None, - }); - } - vectors.push(values.value(value_index)); - } - } - } - let actual_row_count = expected_row_id - row_range_start; - if actual_row_count != expected_row_count { - return Err(Error::DataInvalid { - message: format!( - "Lumina vector extraction expected {} rows, got {}", - expected_row_count, actual_row_count - ), - source: None, - }); - } - Ok(vectors) -} - -fn checked_i64(value: u64, context: &str) -> Result { - i64::try_from(value).map_err(|_| Error::DataInvalid { - message: format!("{context}: {value}"), - source: None, - }) -} - -fn checked_row_count(row_range_start: i64, row_range_end: i64) -> Result { - if row_range_end < row_range_start { - return Err(Error::DataInvalid { - message: format!("Invalid Lumina row range [{row_range_start}, {row_range_end}]"), - source: None, - }); - } - row_range_end - .checked_sub(row_range_start) - .and_then(|span| span.checked_add(1)) - .ok_or_else(|| Error::DataInvalid { - message: format!( - "Row count overflows for row range [{row_range_start}, {row_range_end}]" - ), - source: None, - }) -} -fn validate_vector_buffer(vectors: &[f32], row_count: i32, dimension: i32) -> Result<()> { - if row_count <= 0 { - return Err(Error::DataInvalid { - message: format!("Lumina shard row count must be positive, got: {row_count}"), - source: None, - }); - } - if dimension <= 0 { - return Err(Error::DataInvalid { - message: format!("Lumina vector dimension must be positive, got: {dimension}"), - source: None, - }); - } - let row_count = row_count as usize; - let dimension = dimension as usize; - let expected_len = row_count - .checked_mul(dimension) - .ok_or_else(|| Error::DataInvalid { - message: format!( - "Lumina vector buffer length overflows: row_count={row_count}, dimension={dimension}" - ), - source: None, - })?; - if vectors.len() != expected_len { - return Err(Error::DataInvalid { - message: format!( - "Lumina vector buffer length {} does not match row_count={} and dimension={}", - vectors.len(), - row_count, - dimension - ), - source: None, - }); - } - Ok(()) -} - -fn temp_lumina_path() -> PathBuf { - std::env::temp_dir().join(format!("lumina-index-{}.index", uuid::Uuid::new_v4())) -} - -struct TempFileGuard { - path: Option, -} - -impl TempFileGuard { - fn new(path: PathBuf) -> Self { - Self { path: Some(path) } - } - - fn cleanup(mut self) { - if let Some(path) = self.path.take() { - let _ = std::fs::remove_file(path); - } - } -} - -impl Drop for TempFileGuard { - fn drop(&mut self) { - if let Some(path) = self.path.take() { - let _ = std::fs::remove_file(path); - } - } -} - -async fn copy_local_file_to_output( - source_path: &Path, - output: crate::io::OutputFile, -) -> Result<()> { - let mut source = - tokio::fs::File::open(source_path) - .await - .map_err(|e| Error::UnexpectedError { - message: format!("Failed to open temporary Lumina index file: {e}"), - source: None, - })?; - let mut writer = output.writer().await?; - let mut buffer = vec![0u8; COPY_BUFFER_SIZE]; - - loop { - let len = source - .read(&mut buffer) - .await - .map_err(|e| Error::UnexpectedError { - message: format!("Failed to read temporary Lumina index file: {e}"), - source: None, - })?; - if len == 0 { - break; - } - writer.write(Bytes::copy_from_slice(&buffer[..len])).await?; + Ok(shard_count) } - writer.close().await } #[cfg(test)] -mod tests { - use super::*; - use crate::catalog::Identifier; - use crate::io::FileIO; - use crate::io::FileIOBuilder; - use crate::lumina::LUMINA_DIMENSION_OPTION; - use crate::spec::stats::BinaryTableStats; - use crate::spec::{ - ArrayType, DoubleType, FloatType, IndexManifest, IntType, ManifestEntry, Schema, - TableSchema, VectorType, - }; - use crate::table::TableWrite; - use arrow_array::builder::{FixedSizeListBuilder, Float32Builder, Int64Builder, ListBuilder}; - use arrow_array::{ArrayRef, Int32Array}; - use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; - use chrono::{DateTime, Utc}; - use std::sync::Arc; - - fn data_file(name: &str, first_row_id: Option, row_count: i64) -> DataFileMeta { - DataFileMeta { - file_name: name.to_string(), - file_size: 128, - row_count, - min_key: vec![], - max_key: vec![], - key_stats: BinaryTableStats::new(vec![], vec![], vec![]), - value_stats: BinaryTableStats::new(vec![], vec![], vec![]), - min_sequence_number: 0, - max_sequence_number: 0, - schema_id: 0, - level: 0, - extra_files: vec![], - creation_time: Some( - "2024-09-06T07:45:55.039+00:00" - .parse::>() - .unwrap(), - ), - delete_row_count: None, - embedded_index: None, - first_row_id, - write_cols: None, - external_path: None, - file_source: None, - value_stats_cols: None, - column_max_sequence_numbers: None, - } - } - - fn manifest_entry(file: DataFileMeta) -> ManifestEntry { - manifest_entry_with_bucket(file, 0, 1) - } - - fn manifest_entry_with_bucket( - file: DataFileMeta, - bucket: i32, - total_buckets: i32, - ) -> ManifestEntry { - ManifestEntry::new(FileKind::Add, vec![], bucket, total_buckets, file, 2) - } - - fn table_options(rows_per_shard: &str) -> HashMap { - HashMap::from([ - ("row-tracking.enabled".to_string(), "true".to_string()), - ("data-evolution.enabled".to_string(), "true".to_string()), - ("global-index.enabled".to_string(), "true".to_string()), - ( - "global-index.row-count-per-shard".to_string(), - rows_per_shard.to_string(), - ), - ]) - } - - fn test_table(options: HashMap) -> Table { - test_table_with_io( - FileIOBuilder::new("memory").build().unwrap(), - "memory:/test_lumina_builder", - Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column( - "embedding", - DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), - ) - .options(options) - .build() - .unwrap(), - ) - } - - fn test_table_with_schema(schema: Schema) -> Table { - test_table_with_io( - FileIOBuilder::new("memory").build().unwrap(), - "memory:/test_lumina_builder", - schema, - ) - } - - fn test_table_with_io(file_io: FileIO, table_path: &str, schema: Schema) -> Table { - Table::new( - file_io, - Identifier::new("default", "test_table"), - table_path.to_string(), - TableSchema::new(0, &schema), - None, - ) - } - - fn vector_schema_builder(options: HashMap) -> crate::spec::SchemaBuilder { - Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column( - "embedding", - DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), - ) - .options(options) - } - - fn plan(entries: Vec, rows_per_shard: i64) -> Result> { - plan_with_indexed(entries, rows_per_shard, &[]) - } - - fn plan_with_indexed( - entries: Vec, - rows_per_shard: i64, - indexed: &[RowRange], - ) -> Result> { - let table = test_table(table_options(&rows_per_shard.to_string())); - let core = CoreOptions::new(table.schema().options()); - plan_lumina_shards( - table.location(), - table.schema().partition_keys(), - table.schema().fields(), - &core, - 1, - entries, - rows_per_shard, - indexed, - ) - } - - #[test] - fn test_planner_splits_single_file_across_shards() { - let shards = plan(vec![manifest_entry(data_file("a", Some(0), 25))], 10).unwrap(); - - assert_eq!( - shards - .iter() - .map(|s| (s.row_range_start, s.row_range_end)) - .collect::>(), - vec![(0, 9), (10, 19), (20, 24)] - ); - } - - #[test] - fn test_planner_merges_contiguous_files() { - let shards = plan( - vec![ - manifest_entry(data_file("a", Some(0), 5)), - manifest_entry(data_file("b", Some(5), 5)), - ], - 20, - ) - .unwrap(); - - assert_eq!(shards.len(), 1); - assert_eq!((shards[0].row_range_start, shards[0].row_range_end), (0, 9)); - assert_eq!( - shards[0] - .files - .iter() - .map(|f| f.file_name.as_str()) - .collect::>(), - vec!["a", "b"] - ); - } - - #[test] - fn test_planner_keeps_source_buckets_separate() { - let shards = plan( - vec![ - manifest_entry_with_bucket(data_file("a", Some(0), 5), 0, 2), - manifest_entry_with_bucket(data_file("b", Some(5), 5), 1, 2), - ], - 20, - ) - .unwrap(); - - assert_eq!(shards.len(), 2); - assert_eq!( - shards - .iter() - .map(|s| ( - s.source_bucket, - s.total_buckets, - s.row_range_start, - s.row_range_end - )) - .collect::>(), - vec![(0, 2, 0, 4), (1, 2, 5, 9)] - ); - } - - #[test] - fn test_planner_splits_gap_into_separate_groups() { - let shards = plan( - vec![ - manifest_entry(data_file("a", Some(0), 5)), - manifest_entry(data_file("b", Some(10), 5)), - ], - 20, - ) - .unwrap(); - - assert_eq!( - shards - .iter() - .map(|s| (s.row_range_start, s.row_range_end)) - .collect::>(), - vec![(0, 4), (10, 14)] - ); - } - - #[test] - fn test_planner_rejects_missing_first_row_id() { - let err = plan(vec![manifest_entry(data_file("a", None, 5))], 10) - .expect_err("missing first_row_id should fail"); - assert!( - matches!(err, Error::DataInvalid { message, .. } if message.contains("missing first_row_id")) - ); - } - - #[test] - fn test_planner_rejects_invalid_rows_per_shard() { - let err = plan(vec![manifest_entry(data_file("a", Some(0), 5))], 0) - .expect_err("invalid rows per shard should fail"); - assert!( - matches!(err, Error::DataInvalid { message, .. } if message.contains("row-count-per-shard")) - ); - } - - #[test] - fn test_validate_vector_field_accepts_array_float() { - let field = DataField::new( - 0, - "embedding".to_string(), - DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), - ); - assert!(validate_vector_field(&field).is_ok()); - } - - #[test] - fn test_validate_vector_field_accepts_vector_float() { - let field = DataField::new( - 0, - "embedding".to_string(), - DataType::Vector( - VectorType::try_new(true, 4, DataType::Float(FloatType::new())).unwrap(), - ), - ); - assert!(validate_vector_field(&field).is_ok()); - } - - #[test] - fn test_validate_vector_field_rejects_vector_double() { - let field = DataField::new( - 0, - "embedding".to_string(), - DataType::Vector( - VectorType::try_new(true, 4, DataType::Double(DoubleType::new())).unwrap(), - ), - ); - let err = validate_vector_field(&field).expect_err("VECTOR must be rejected"); - assert!(matches!(err, Error::DataInvalid { .. })); - } - - #[test] - fn test_effective_options_vector_absent_inserts_length() { - let field = DataField::new( - 0, - "embedding".to_string(), - DataType::Vector( - VectorType::try_new(true, 256, DataType::Float(FloatType::new())).unwrap(), - ), - ); - let opts = effective_lumina_options(&field, HashMap::new()).unwrap(); - assert_eq!( - opts.get(LUMINA_DIMENSION_OPTION).map(String::as_str), - Some("256") - ); - } - - #[test] - fn test_effective_options_vector_matching_option_ok() { - let field = DataField::new( - 0, - "embedding".to_string(), - DataType::Vector( - VectorType::try_new(true, 256, DataType::Float(FloatType::new())).unwrap(), - ), - ); - let resolved = HashMap::from([(LUMINA_DIMENSION_OPTION.to_string(), "256".to_string())]); - let opts = effective_lumina_options(&field, resolved).unwrap(); - assert_eq!( - opts.get(LUMINA_DIMENSION_OPTION).map(String::as_str), - Some("256") - ); - } - - #[test] - fn test_effective_options_vector_mismatch_errors() { - let field = DataField::new( - 0, - "embedding".to_string(), - DataType::Vector( - VectorType::try_new(true, 256, DataType::Float(FloatType::new())).unwrap(), - ), - ); - let resolved = HashMap::from([(LUMINA_DIMENSION_OPTION.to_string(), "128".to_string())]); - let err = - effective_lumina_options(&field, resolved).expect_err("dimension mismatch must error"); - assert!(matches!(err, Error::ConfigInvalid { .. })); - } - - #[test] - fn test_effective_options_array_unchanged() { - let field = DataField::new( - 0, - "embedding".to_string(), - DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), - ); - // No dimension option set: array path must NOT inject one. - let opts = effective_lumina_options(&field, HashMap::new()).unwrap(); - assert!(!opts.contains_key(LUMINA_DIMENSION_OPTION)); - } - - #[test] - fn test_vector_without_option_propagates_dimension_to_native_options() { - use crate::lumina::{LuminaVectorIndexOptions, KEY_DIMENSION}; - let field = DataField::new( - 0, - "embedding".to_string(), - DataType::Vector( - VectorType::try_new(true, 256, DataType::Float(FloatType::new())).unwrap(), - ), - ); - // No lumina.index.dimension set by the user. - let resolved = effective_lumina_options(&field, HashMap::new()).unwrap(); - let opts = LuminaVectorIndexOptions::new(&resolved).unwrap(); - - assert_eq!( - opts.dimension, 256, - "local dimension must be N, not default 128" - ); - assert_eq!( - opts.to_lumina_options() - .get(KEY_DIMENSION) - .map(String::as_str), - Some("256"), - "native index.dimension must be N, not default 128" - ); - } - - #[tokio::test] - async fn test_execute_rejects_primary_key_table() { - let table = test_table_with_schema( - vector_schema_builder(HashMap::new()) - .primary_key(["id"]) - .build() - .unwrap(), - ); - - let err = table - .new_lumina_index_build_builder() - .with_index_column("embedding") - .execute() - .await - .expect_err("primary-key table should fail before native build"); - - assert!( - matches!(err, Error::Unsupported { message } if message.contains("primary-key tables")) - ); - } - - #[tokio::test] - async fn test_execute_rejects_deletion_vectors_table() { - let mut options = table_options("10"); - options.insert("deletion-vectors.enabled".to_string(), "true".to_string()); - let table = test_table(options); - - let err = table - .new_lumina_index_build_builder() - .with_index_column("embedding") - .execute() - .await - .expect_err("deletion vectors table should fail before native build"); - - assert!( - matches!(err, Error::Unsupported { message } if message.contains("deletion-vectors.enabled=true")) - ); - } - - fn vector_batch(rows: Vec>>>, row_ids: Vec>) -> RecordBatch { - let mut vector_builder = ListBuilder::new(Float32Builder::new()); - for row in rows { - match row { - Some(values) => { - for value in values { - match value { - Some(value) => vector_builder.values().append_value(value), - None => vector_builder.values().append_null(), - } - } - vector_builder.append(true); - } - None => vector_builder.append(false), - } - } - let mut row_id_builder = Int64Builder::new(); - for row_id in row_ids { - match row_id { - Some(value) => row_id_builder.append_value(value), - None => row_id_builder.append_null(), - } - } - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new( - "embedding", - ArrowDataType::List(Arc::new(ArrowField::new( - "item", - ArrowDataType::Float32, - true, - ))), - true, - ), - ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), - ])); - RecordBatch::try_new( - schema, - vec![ - Arc::new(vector_builder.finish()) as ArrayRef, - Arc::new(row_id_builder.finish()) as ArrayRef, - ], - ) - .unwrap() - } - - #[test] - fn test_extract_vectors_accepts_list_float32_and_row_ids() { - let batch = vector_batch( - vec![ - Some(vec![Some(1.0), Some(2.0)]), - Some(vec![Some(3.0), Some(4.0)]), - ], - vec![Some(10), Some(11)], - ); - - let vectors = extract_vectors_from_batches(&[batch], "embedding", 2, 10, 2).unwrap(); - - assert_eq!(vectors, vec![1.0, 2.0, 3.0, 4.0]); - } - - #[test] - fn test_extract_vectors_rejects_null_vector() { - let batch = vector_batch(vec![None], vec![Some(0)]); - - let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) - .expect_err("null vector should fail"); - - assert!( - matches!(err, Error::DataInvalid { message, .. } if message.contains("null vector")) - ); - } - - #[test] - fn test_extract_vectors_rejects_null_element() { - let batch = vector_batch(vec![Some(vec![Some(1.0), None])], vec![Some(0)]); - - let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) - .expect_err("null element should fail"); - - assert!( - matches!(err, Error::DataInvalid { message, .. } if message.contains("null vector element")) - ); - } - - #[test] - fn test_extract_vectors_rejects_dimension_mismatch() { - let batch = vector_batch(vec![Some(vec![Some(1.0)])], vec![Some(0)]); - - let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) - .expect_err("dimension mismatch should fail"); - - assert!( - matches!(err, Error::DataInvalid { message, .. } if message.contains("dimension mismatch")) - ); - } - - #[test] - fn test_extract_vectors_rejects_row_id_gap() { - let batch = vector_batch( - vec![ - Some(vec![Some(1.0), Some(2.0)]), - Some(vec![Some(3.0), Some(4.0)]), - ], - vec![Some(0), Some(2)], - ); - - let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 2) - .expect_err("row id gap should fail"); - - assert!( - matches!(err, Error::DataInvalid { message, .. } if message.contains("expected _ROW_ID")) - ); - } - - #[test] - fn test_extract_vectors_rejects_non_list_float32() { - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("embedding", ArrowDataType::Int32, false), - ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from(vec![1])) as ArrayRef, - Arc::new(Int64Array::from(vec![Some(0)])) as ArrayRef, - ], - ) - .unwrap(); - - let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) - .expect_err("non-list vector should fail"); - - assert!( - matches!(err, Error::DataInvalid { message, .. } if message.contains("List")) - ); - } - - fn fixed_size_vector_batch( - rows: Vec>>, - row_ids: Vec>, - len: i32, - ) -> RecordBatch { - let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); - let mut builder = - FixedSizeListBuilder::new(Float32Builder::new(), len).with_field(element_field); - for row in rows { - match row { - Some(values) => { - for v in values { - builder.values().append_value(v); - } - builder.append(true); - } - None => { - for _ in 0..len { - builder.values().append_value(0.0); - } - builder.append(false); - } - } - } - let mut row_id_builder = Int64Builder::new(); - for row_id in row_ids { - match row_id { - Some(value) => row_id_builder.append_value(value), - None => row_id_builder.append_null(), - } - } - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new( - "embedding", - ArrowDataType::FixedSizeList( - Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)), - len, - ), - true, - ), - ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), - ])); - RecordBatch::try_new( - schema, - vec![ - Arc::new(builder.finish()) as ArrayRef, - Arc::new(row_id_builder.finish()) as ArrayRef, - ], - ) - .unwrap() - } - - #[test] - fn test_extract_vectors_accepts_fixed_size_list_float32() { - let batch = fixed_size_vector_batch( - vec![Some(vec![1.0, 2.0]), Some(vec![3.0, 4.0])], - vec![Some(10), Some(11)], - 2, - ); - let vectors = extract_vectors_from_batches(&[batch], "embedding", 2, 10, 2).unwrap(); - assert_eq!(vectors, vec![1.0, 2.0, 3.0, 4.0]); - } - - #[test] - fn test_extract_vectors_fixed_size_list_rejects_null_vector() { - let batch = fixed_size_vector_batch(vec![None], vec![Some(0)], 2); - let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) - .expect_err("null vector should fail"); - assert!( - matches!(err, Error::DataInvalid { message, .. } if message.contains("null vector")) - ); - } - - #[test] - fn test_extract_vectors_fixed_size_list_dimension_mismatch() { - // Column is FixedSizeList of length 3, but caller expects dimension 2. - let batch = fixed_size_vector_batch(vec![Some(vec![1.0, 2.0, 3.0])], vec![Some(0)], 3); - let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) - .expect_err("dimension mismatch should fail"); - assert!( - matches!(err, Error::DataInvalid { message, .. } if message.contains("dimension mismatch")) - ); - } - - #[test] - fn test_extract_vectors_fixed_size_list_rejects_null_element() { - // A non-null vector row whose second child element is null. Mirrors the - // List path's test_extract_vectors_rejects_null_element so both layouts - // reject null elements identically. - let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); - let mut builder = - FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(element_field); - builder.values().append_value(1.0); - builder.values().append_null(); - builder.append(true); - let row_ids = Arc::new(Int64Array::from(vec![Some(0)])) as ArrayRef; - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new( - "embedding", - ArrowDataType::FixedSizeList( - Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)), - 2, - ), - true, - ), - ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), - ])); - let batch = RecordBatch::try_new( - schema, - vec![Arc::new(builder.finish()) as ArrayRef, row_ids], - ) - .unwrap(); - - let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) - .expect_err("null element should fail"); - - assert!( - matches!(err, Error::DataInvalid { message, .. } if message.contains("null vector element")) - ); - } - - #[test] - fn test_checked_metadata_conversion_supports_i64_file_size() { - let above_i32_max = i32::MAX as u64 + 1; - assert_eq!( - checked_i64(above_i32_max, "Index file is too large").unwrap(), - i64::from(i32::MAX) + 1 - ); - - let err = checked_i64(i64::MAX as u64 + 1, "Index file is too large") - .expect_err("file size above i64::MAX should fail"); - assert!(matches!(err, Error::DataInvalid { message, .. } if message.contains("too large"))); - } - - #[test] - fn test_temp_file_guard_cleans_up_on_drop() { - let path = temp_lumina_path(); - std::fs::write(&path, b"temporary lumina data").unwrap(); - { - let _guard = TempFileGuard::new(path.clone()); - assert!(path.exists()); - } - assert!(!path.exists()); - } - - async fn setup_dirs(file_io: &FileIO, table_path: &str) { - file_io - .mkdirs(&format!("{table_path}/snapshot/")) - .await - .unwrap(); - file_io - .mkdirs(&format!("{table_path}/manifest/")) - .await - .unwrap(); - } - - fn build_vector_batch(ids: Vec, vectors: Vec>) -> RecordBatch { - let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); - let mut vector_builder = - ListBuilder::new(Float32Builder::new()).with_field(element_field.clone()); - for vector in vectors { - for value in vector { - vector_builder.values().append_value(value); - } - vector_builder.append(true); - } - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", ArrowDataType::Int32, false), - ArrowField::new("embedding", ArrowDataType::List(element_field), true), - ])); - RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from(ids)) as ArrayRef, - Arc::new(vector_builder.finish()) as ArrayRef, - ], - ) - .unwrap() - } - - // Manual run with a local Lumina native library: - // LUMINA_LIB_PATH=/path/to/liblumina_py.so cargo test -p paimon \ - // table::lumina_index_build_builder::tests::test_execute_writes_lumina_index_manifest \ - // --features fulltext,vortex -- --ignored --exact - #[tokio::test] - #[ignore = "requires LUMINA_LIB_PATH; see manual run command above"] - async fn test_execute_writes_lumina_index_manifest() { - let file_io = FileIOBuilder::new("memory").build().unwrap(); - let table_path = "memory:/test_lumina_builder_e2e"; - setup_dirs(&file_io, table_path).await; - - let mut options = table_options("10"); - options.insert("lumina.index.dimension".to_string(), "2".to_string()); - options.insert("lumina.encoding.type".to_string(), "rawf32".to_string()); - let table = test_table_with_io( - file_io.clone(), - table_path, - vector_schema_builder(options).build().unwrap(), - ); - - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&build_vector_batch( - vec![1, 2], - vec![vec![1.0, 0.0], vec![0.0, 1.0]], - )) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - let shard_count = table - .new_lumina_index_build_builder() - .with_index_column("embedding") - .execute() - .await - .unwrap(); - assert_eq!(shard_count, 1); - - let snapshot_manager = SnapshotManager::new(file_io.clone(), table_path.to_string()); - let snapshot = snapshot_manager - .get_latest_snapshot() - .await - .unwrap() - .unwrap(); - assert_eq!(snapshot.id(), 2); - assert_eq!(snapshot.total_record_count(), Some(2)); - assert_eq!(snapshot.delta_record_count(), Some(0)); - assert_eq!(snapshot.next_row_id(), Some(2)); - - let index_manifest = snapshot.index_manifest().expect("index manifest"); - let index_entries = - IndexManifest::read(&file_io, &format!("{table_path}/manifest/{index_manifest}")) - .await - .unwrap(); - assert_eq!(index_entries.len(), 1); - - let index_file = &index_entries[0].index_file; - assert_eq!(index_file.index_type, LUMINA_IDENTIFIER); - assert!(index_file.file_name.starts_with("lumina-global-index-")); - assert_eq!(index_file.row_count, 2); - assert!(index_file.file_size > 0); - - let global_meta = index_file - .global_index_meta - .as_ref() - .expect("global index meta"); - assert_eq!(global_meta.row_range_start, 0); - assert_eq!(global_meta.row_range_end, 1); - assert_eq!(global_meta.index_field_id, 1); - assert!(global_meta - .index_meta - .as_ref() - .is_some_and(|m| !m.is_empty())); - - let index_path = format!("{table_path}/index/{}", index_file.file_name); - let status = file_io.get_status(&index_path).await.unwrap(); - assert_eq!(index_file.file_size as u64, status.size); - } - - fn lumina_e2e_options(rows_per_shard: &str) -> HashMap { - let mut options = table_options(rows_per_shard); - options.insert("lumina.index.dimension".to_string(), "2".to_string()); - options.insert("lumina.encoding.type".to_string(), "rawf32".to_string()); - options - } - - fn lumina_e2e_table(table_path: &str, rows_per_shard: &str) -> Table { - test_table_with_io( - FileIOBuilder::new("memory").build().unwrap(), - table_path, - vector_schema_builder(lumina_e2e_options(rows_per_shard)) - .build() - .unwrap(), - ) - } - - async fn write_vectors(table: &Table, ids: Vec, vectors: Vec>) { - let mut table_write = TableWrite::new(table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&build_vector_batch(ids, vectors)) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - } - - /// Commit a synthetic Lumina `IndexFileMeta` covering `[start, end]` for - /// `field_id` directly into the index manifest, without invoking the native - /// builder. Mirrors the btree mid-hole test so the incremental gap logic can - /// be exercised in CI where the native Lumina library is unavailable. - async fn commit_synthetic_lumina_index(table: &Table, field_id: i32, start: i64, end: i64) { - let synthetic = IndexFileMeta { - index_type: LUMINA_IDENTIFIER.to_string(), - file_name: format!("lumina-synthetic-{start}-{end}.index"), - file_size: 1, - row_count: (end - start + 1), - deletion_vectors_ranges: None, - global_index_meta: Some(GlobalIndexMeta { - row_range_start: start, - row_range_end: end, - index_field_id: field_id, - extra_field_ids: None, - source_meta: None, - index_meta: None, - }), - }; - let mut message = CommitMessage::new(BinaryRow::new(0).to_serialized_bytes(), 0, vec![]); - message.new_index_files = vec![synthetic]; - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(vec![message]) - .await - .unwrap(); - } - - async fn latest_lumina_index_files(table: &Table) -> Vec { - let snapshot_manager = - SnapshotManager::new(table.file_io().clone(), table.location().to_string()); - let snapshot = snapshot_manager - .get_latest_snapshot() - .await - .unwrap() - .unwrap(); - let Some(index_manifest_name) = snapshot.index_manifest() else { - return Vec::new(); - }; - IndexManifest::read( - table.file_io(), - &snapshot_manager.manifest_path(index_manifest_name), - ) - .await - .unwrap() - .into_iter() - .filter(|entry| { - entry.kind == FileKind::Add && entry.index_file.index_type == LUMINA_IDENTIFIER - }) - .map(|entry| entry.index_file) - .collect() - } - - /// Row-id coverage of the committed data files, read back from the data - /// manifest (never hard-coded) and merged into contiguous ranges. Mirrors - /// how `execute` gathers `manifest_entries`. - async fn data_row_id_coverage(table: &Table) -> Vec { - let snapshot_manager = - SnapshotManager::new(table.file_io().clone(), table.location().to_string()); - let snapshot = snapshot_manager - .get_latest_snapshot() - .await - .unwrap() - .unwrap(); - let entries = table - .new_read_builder() - .new_scan() - .with_scan_all_files() - .plan_manifest_entries(&snapshot) - .await - .unwrap(); - let ranges = entries - .iter() - .filter(|entry| *entry.kind() == FileKind::Add) - .filter_map(|entry| { - entry - .file() - .row_id_range() - .map(|(start, end)| RowRange::new(start, end)) - }) - .collect::>(); - crate::table::merge_row_ranges(ranges) - } - - /// Second build with the whole coverage already indexed must be a clean - /// no-op (returns 0), not an overlap error. Reaches `Ok(0)` before the - /// native build, so it runs in CI without the Lumina library. This is the - /// core bug fix: today the second call errors with the overlap message. - #[tokio::test] - async fn lumina_second_build_without_new_data_is_noop() { - let table_path = "memory:/test_lumina_second_build_noop"; - let table = lumina_e2e_table(table_path, "10"); - setup_dirs(table.file_io(), table_path).await; - - write_vectors( - &table, - vec![1, 2, 3], - vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]], - ) - .await; - - // Fully index the coverage via a synthetic manifest entry. - let coverage = data_row_id_coverage(&table).await; - assert_eq!(coverage.len(), 1, "data must be one contiguous range"); - let field_id = find_index_field(&table, "embedding").unwrap().id(); - commit_synthetic_lumina_index(&table, field_id, coverage[0].from(), coverage[0].to()).await; - - let names_before = latest_lumina_index_files(&table) - .await - .iter() - .map(|f| f.file_name.clone()) - .collect::>(); - assert!(!names_before.is_empty()); - - let built = table - .new_lumina_index_build_builder() - .with_index_column("embedding") - .execute() - .await - .unwrap(); - assert_eq!(built, 0, "fully-indexed table must build nothing on re-run"); - - let names_after = latest_lumina_index_files(&table) - .await - .iter() - .map(|f| f.file_name.clone()) - .collect::>(); - assert_eq!( - names_before, names_after, - "re-run must not add or remove index manifest entries" - ); - } - - /// Build over an already-indexed prefix, then append new rows: the second - /// build must target only the appended gap and must NOT fail with the - /// overlap error. Without the native Lumina library the gap build surfaces a - /// library-load error (not the overlap error); with it present it succeeds. - #[tokio::test] - async fn lumina_incremental_build_indexes_only_new_rows() { - let table_path = "memory:/test_lumina_incremental"; - let table = lumina_e2e_table(table_path, "10"); - setup_dirs(table.file_io(), table_path).await; - - // Initial batch, then mark it fully indexed via a synthetic entry. - write_vectors( - &table, - vec![1, 2, 3], - vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]], - ) - .await; - let indexed_coverage = data_row_id_coverage(&table).await; - assert_eq!(indexed_coverage.len(), 1); - let n = indexed_coverage[0].to() + 1; - let field_id = find_index_field(&table, "embedding").unwrap().id(); - commit_synthetic_lumina_index( - &table, - field_id, - indexed_coverage[0].from(), - indexed_coverage[0].to(), - ) - .await; - - // Append a second batch (new row-ids [n..]). - write_vectors( - &table, - vec![4, 5, 6], - vec![vec![2.0, 0.0], vec![0.0, 2.0], vec![2.0, 2.0]], - ) - .await; - - // White-box: fed the real indexed ranges from the manifest, the planner - // must target only the appended gap [n, ..], never the already-indexed - // prefix. Computed before `execute` so it is independent of whether the - // native build (which needs the Lumina library) runs. - let snapshot_manager = - SnapshotManager::new(table.file_io().clone(), table.location().to_string()); - let snapshot = snapshot_manager - .get_latest_snapshot() - .await - .unwrap() - .unwrap(); - let manifest_entries = table - .new_read_builder() - .new_scan() - .with_scan_all_files() - .plan_manifest_entries(&snapshot) - .await - .unwrap(); - let indexed = crate::table::global_index_build_common::indexed_row_ranges( - &table, - snapshot.index_manifest(), - LUMINA_IDENTIFIER, - field_id, - None, - ) - .await - .unwrap(); - let core = CoreOptions::new(table.schema().options()); - let shards = plan_lumina_shards( - table.location(), - table.schema().partition_keys(), - table.schema().fields(), - &core, - snapshot.id(), - manifest_entries, - 10, - &indexed, - ) - .unwrap(); - assert!(!shards.is_empty(), "appended gap must produce build shards"); - for shard in &shards { - assert!( - shard.row_range_start >= n, - "shard [{}, {}] must start at or after the indexed prefix end {n}", - shard.row_range_start, - shard.row_range_end - ); - } - - // End-to-end: the incremental build must no longer fail with the overlap - // error. Without the native Lumina library the gap build surfaces a - // library-load error instead; with it present it succeeds. - let result = table - .new_lumina_index_build_builder() - .with_index_column("embedding") - .execute() - .await; - match result { - Ok(_) => {} - Err(Error::DataInvalid { message, .. }) => { - assert!( - !message.contains("overlaps requested row range"), - "incremental build must not fail with the overlap error; got: {message}" - ); - } - Err(other) => panic!("unexpected error from incremental build: {other:?}"), - } - } - - /// Regression: a first build (no existing index) must equal the pre-change - /// full build -- subtracting an empty `indexed` yields full coverage. - #[test] - fn lumina_first_build_indexes_full_coverage() { - let full = plan(vec![manifest_entry(data_file("a", Some(0), 25))], 10).unwrap(); - let gapped = - plan_with_indexed(vec![manifest_entry(data_file("a", Some(0), 25))], 10, &[]).unwrap(); - // Empty `indexed` must not alter the shard layout. - assert_eq!( - full.iter() - .map(|s| (s.row_range_start, s.row_range_end)) - .collect::>(), - gapped - .iter() - .map(|s| (s.row_range_start, s.row_range_end)) - .collect::>() - ); - assert_eq!( - full.iter() - .map(|s| (s.row_range_start, s.row_range_end)) - .collect::>(), - vec![(0, 9), (10, 19), (20, 24)], - "first build must cover the full row range across shards" - ); - } - - /// Planner-level mid-coverage hole, mirroring btree's - /// `incremental_build_splits_gap_around_mid_coverage_indexed_hole`: with a - /// single shard cell (rows_per_shard large enough to hold all data) the grid - /// never splits, so the only split is the indexed hole itself. An indexed - /// range strictly inside the data coverage must carve the build into exactly - /// the two contiguous segments on either side of the hole -- both bounds - /// pinned, and neither segment may span or touch the hole. - #[test] - fn lumina_plan_splits_gap_around_mid_coverage_indexed_hole() { - // Data row-ids [0, 9]; one shard cell [0, 99] so the grid never splits. - let n = 9; - let hole_start = 4; - let hole_end = 6; - let shards = plan_with_indexed( - vec![manifest_entry(data_file("a", Some(0), n + 1))], - 100, - &[RowRange::new(hole_start, hole_end)], - ) - .unwrap(); - - let ranges = shards - .iter() - .map(|s| (s.row_range_start, s.row_range_end)) - .collect::>(); - // Exactly the two contiguous segments around the hole. - assert_eq!( - ranges, - vec![(0, hole_start - 1), (hole_end + 1, n)], - "mid-coverage hole must split into exactly the two segments around it" - ); - // Every emitted range is contiguous and none spans or touches the hole. - for (start, end) in &ranges { - assert!(end >= start, "range must be non-empty: [{start}, {end}]"); - assert!( - *end < hole_start || *start > hole_end, - "shard [{start}, {end}] must not overlap indexed hole [{hole_start}, {hole_end}]" - ); - } - // Together the shards cover exactly coverage - indexed. - let expected = exclude_row_ranges( - &[RowRange::new(0, n)], - &[RowRange::new(hole_start, hole_end)], - ) - .into_iter() - .map(|r| (r.from(), r.to())) - .collect::>(); - assert_eq!( - ranges, expected, - "shards must cover exactly coverage minus the indexed hole" - ); - } - - /// Planner-level incremental prefix. Strengthens - /// `lumina_incremental_build_indexes_only_new_rows`, which asserted only a - /// one-sided lower bound (`row_range_start >= n`): an indexed prefix [0, k] - /// must leave EXACTLY the suffix [k+1, N] on both bounds, split along the - /// shard grid, with nothing re-indexed inside the prefix. - #[test] - fn lumina_plan_incremental_prefix_leaves_suffix() { - // Data row-ids [0, 24], rows_per_shard = 10 -> cells [0,9],[10,19],[20,29]. - // Indexed prefix [0, 9] fully fills the first cell, so the build must be - // exactly [10, 19] and [20, 24] (the suffix split along the grid). - let n = 24; - let k = 9; // prefix [0, k] == the first full shard cell - let shards = plan_with_indexed( - vec![manifest_entry(data_file("a", Some(0), n + 1))], - 10, - &[RowRange::new(0, k)], - ) - .unwrap(); - - let ranges = shards - .iter() - .map(|s| (s.row_range_start, s.row_range_end)) - .collect::>(); - assert_eq!( - ranges, - vec![(k + 1, 19), (20, n)], - "indexed prefix must leave exactly the suffix, split along the shard grid" - ); - // Both bounds pinned (this is what the one-sided existing check omits). - assert_eq!(ranges.first().unwrap().0, k + 1, "suffix must start at k+1"); - assert_eq!(ranges.last().unwrap().1, n, "suffix must end at N"); - // Contiguous, and no shard reaches back into the indexed prefix. - for pair in ranges.windows(2) { - assert_eq!( - pair[1].0, - pair[0].1 + 1, - "ranges must be contiguous: {:?} then {:?}", - pair[0], - pair[1] - ); - } - for (start, end) in &ranges { - assert!( - *start > k, - "shard [{start}, {end}] must not re-index the prefix [0, {k}]" - ); - } - } -} +mod tests; diff --git a/crates/paimon/src/table/lumina_index_build_builder/extraction.rs b/crates/paimon/src/table/lumina_index_build_builder/extraction.rs new file mode 100644 index 000000000..1076507ac --- /dev/null +++ b/crates/paimon/src/table/lumina_index_build_builder/extraction.rs @@ -0,0 +1,199 @@ +// 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. + +//! Arrow vector extraction for Lumina shards. + +use super::planning::LuminaIndexShard; +use super::validation::checked_row_count; +use crate::spec::ROW_ID_FIELD_NAME; +use crate::table::{DataSplitBuilder, RowRange, Table}; +use crate::{Error, Result}; +use arrow_array::{Array, FixedSizeListArray, Float32Array, Int64Array, ListArray, RecordBatch}; +use futures::TryStreamExt; + +pub(super) async fn extract_vectors( + table: &Table, + shard: &LuminaIndexShard, + index_column: &str, + dimension: i32, +) -> Result> { + let split = DataSplitBuilder::new() + .with_snapshot(shard.snapshot_id) + .with_partition(shard.partition.clone()) + .with_bucket(shard.source_bucket) + .with_bucket_path(shard.bucket_path.clone()) + .with_total_buckets(shard.total_buckets) + .with_data_files(shard.files.clone()) + .with_row_ranges(vec![RowRange::new( + shard.row_range_start, + shard.row_range_end, + )]) + .build()?; + + let mut read_builder = table.new_read_builder(); + read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; + let read = read_builder.new_read()?; + let batches = read.to_arrow(&[split])?.try_collect::>().await?; + extract_vectors_from_batches( + &batches, + index_column, + dimension, + shard.row_range_start, + checked_row_count(shard.row_range_start, shard.row_range_end)?, + ) +} + +pub(super) fn extract_vectors_from_batches( + batches: &[RecordBatch], + index_column: &str, + dimension: i32, + row_range_start: i64, + expected_row_count: i64, +) -> Result> { + let dimension = usize::try_from(dimension).map_err(|e| Error::DataInvalid { + message: format!("Invalid Lumina dimension: {dimension}"), + source: Some(Box::new(e)), + })?; + let row_count = batches.iter().map(RecordBatch::num_rows).sum::(); + let mut vectors = Vec::with_capacity(row_count * dimension); + let mut expected_row_id = row_range_start; + for batch in batches { + let vector_index = + batch + .schema() + .index_of(index_column) + .map_err(|e| Error::DataInvalid { + message: format!("Vector column '{index_column}' not found in read batch: {e}"), + source: None, + })?; + let row_id_index = + batch + .schema() + .index_of(ROW_ID_FIELD_NAME) + .map_err(|e| Error::DataInvalid { + message: format!("_ROW_ID column not found in read batch: {e}"), + source: None, + })?; + // Resolve the vector column as either List (ARRAY) or + // FixedSizeList (VECTOR). Both yield a Float32Array of + // values plus a per-row [start, end) slice. + let column = batch.column(vector_index); + enum VectorLayout<'a> { + List(&'a ListArray), + Fixed(&'a FixedSizeListArray), + } + let layout = if let Some(a) = column.as_any().downcast_ref::() { + VectorLayout::List(a) + } else if let Some(a) = column.as_any().downcast_ref::() { + VectorLayout::Fixed(a) + } else { + return Err(Error::DataInvalid { + message: "Lumina vector extraction requires Arrow List or FixedSizeList".to_string(), + source: None, + }); + }; + let values = match layout { + VectorLayout::List(a) => a.values(), + VectorLayout::Fixed(a) => a.values(), + } + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::DataInvalid { + message: "Lumina vector extraction requires Float32 vector elements".to_string(), + source: None, + })?; + let row_ids = batch + .column(row_id_index) + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::DataInvalid { + message: "Lumina vector extraction requires non-null Int64 _ROW_ID".to_string(), + source: None, + })?; + + for row in 0..batch.num_rows() { + if row_ids.is_null(row) { + return Err(Error::DataInvalid { + message: "Lumina vector extraction found null _ROW_ID".to_string(), + source: None, + }); + } + let row_id = row_ids.value(row); + if row_id != expected_row_id { + return Err(Error::DataInvalid { + message: format!( + "Lumina vector extraction expected _ROW_ID {}, got {}", + expected_row_id, row_id + ), + source: None, + }); + } + expected_row_id += 1; + + let is_null = match layout { + VectorLayout::List(a) => a.is_null(row), + VectorLayout::Fixed(a) => a.is_null(row), + }; + if is_null { + return Err(Error::DataInvalid { + message: "Lumina vector extraction found null vector row".to_string(), + source: None, + }); + } + let (start, end) = match layout { + VectorLayout::List(a) => { + let offsets = a.value_offsets(); + (offsets[row] as usize, offsets[row + 1] as usize) + } + VectorLayout::Fixed(a) => { + let len = a.value_length() as usize; + (row * len, (row + 1) * len) + } + }; + if end - start != dimension { + return Err(Error::DataInvalid { + message: format!( + "Lumina vector dimension mismatch: expected {}, got {}", + dimension, + end - start + ), + source: None, + }); + } + for value_index in start..end { + if values.is_null(value_index) { + return Err(Error::DataInvalid { + message: "Lumina vector extraction found null vector element".to_string(), + source: None, + }); + } + vectors.push(values.value(value_index)); + } + } + } + let actual_row_count = expected_row_id - row_range_start; + if actual_row_count != expected_row_count { + return Err(Error::DataInvalid { + message: format!( + "Lumina vector extraction expected {} rows, got {}", + expected_row_count, actual_row_count + ), + source: None, + }); + } + Ok(vectors) +} diff --git a/crates/paimon/src/table/lumina_index_build_builder/planning.rs b/crates/paimon/src/table/lumina_index_build_builder/planning.rs new file mode 100644 index 000000000..0eafcbf88 --- /dev/null +++ b/crates/paimon/src/table/lumina_index_build_builder/planning.rs @@ -0,0 +1,49 @@ +// 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. + +//! Lumina adapter for shared vector-index shard planning. + +use crate::spec::{CoreOptions, DataField, ManifestEntry}; +use crate::table::global_index_build_common::vector::{plan_vector_index_shards, VectorIndexShard}; +use crate::table::RowRange; +use crate::Result; + +pub(crate) type LuminaIndexShard = VectorIndexShard; + +#[allow(clippy::too_many_arguments)] +pub(super) fn plan_lumina_shards( + table_location: &str, + partition_keys: &[String], + schema_fields: &[DataField], + core_options: &CoreOptions, + snapshot_id: i64, + entries: Vec, + rows_per_shard: i64, + indexed: &[RowRange], +) -> Result> { + plan_vector_index_shards( + table_location, + partition_keys, + schema_fields, + core_options, + snapshot_id, + entries, + rows_per_shard, + indexed, + "Lumina", + ) +} diff --git a/crates/paimon/src/table/lumina_index_build_builder/tests.rs b/crates/paimon/src/table/lumina_index_build_builder/tests.rs new file mode 100644 index 000000000..6aaff6f51 --- /dev/null +++ b/crates/paimon/src/table/lumina_index_build_builder/tests.rs @@ -0,0 +1,1257 @@ +// 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::extract_vectors_from_batches; +use super::planning::{plan_lumina_shards, LuminaIndexShard}; +use super::validation::{ + checked_i64, effective_lumina_options, find_index_field, validate_vector_field, +}; +use super::writer::{abort_on_build_error, temp_lumina_path, TempFileGuard}; +use crate::catalog::Identifier; +use crate::io::FileIO; +use crate::io::FileIOBuilder; +use crate::lumina::{LUMINA_DIMENSION_OPTION, LUMINA_IDENTIFIER}; +use crate::spec::stats::BinaryTableStats; +use crate::spec::{ + ArrayType, BinaryRow, CoreOptions, DataField, DataFileMeta, DataType, DoubleType, FileKind, + FloatType, GlobalIndexMeta, IndexFileMeta, IndexManifest, IntType, ManifestEntry, Schema, + TableSchema, VectorType, ROW_ID_FIELD_NAME, +}; +use crate::table::source::exclude_row_ranges; +use crate::table::{CommitMessage, RowRange, SnapshotManager, Table, TableCommit, TableWrite}; +use crate::{Error, Result}; +use arrow_array::builder::{FixedSizeListBuilder, Float32Builder, Int64Builder, ListBuilder}; +use arrow_array::{ArrayRef, Int32Array, Int64Array, RecordBatch}; +use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; +use bytes::Bytes; +use chrono::{DateTime, Utc}; +use std::collections::HashMap; +use std::sync::Arc; + +fn data_file(name: &str, first_row_id: Option, row_count: i64) -> DataFileMeta { + DataFileMeta { + file_name: name.to_string(), + file_size: 128, + row_count, + min_key: vec![], + max_key: vec![], + key_stats: BinaryTableStats::new(vec![], vec![], vec![]), + value_stats: BinaryTableStats::new(vec![], vec![], vec![]), + min_sequence_number: 0, + max_sequence_number: 0, + schema_id: 0, + level: 0, + extra_files: vec![], + creation_time: Some( + "2024-09-06T07:45:55.039+00:00" + .parse::>() + .unwrap(), + ), + delete_row_count: None, + embedded_index: None, + first_row_id, + write_cols: None, + external_path: None, + file_source: None, + value_stats_cols: None, + column_max_sequence_numbers: None, + } +} + +fn manifest_entry(file: DataFileMeta) -> ManifestEntry { + manifest_entry_with_bucket(file, 0, 1) +} + +fn manifest_entry_with_bucket( + file: DataFileMeta, + bucket: i32, + total_buckets: i32, +) -> ManifestEntry { + ManifestEntry::new(FileKind::Add, vec![], bucket, total_buckets, file, 2) +} + +fn table_options(rows_per_shard: &str) -> HashMap { + HashMap::from([ + ("row-tracking.enabled".to_string(), "true".to_string()), + ("data-evolution.enabled".to_string(), "true".to_string()), + ("global-index.enabled".to_string(), "true".to_string()), + ( + "global-index.row-count-per-shard".to_string(), + rows_per_shard.to_string(), + ), + ]) +} + +fn test_table(options: HashMap) -> Table { + test_table_with_io( + FileIOBuilder::new("memory").build().unwrap(), + "memory:/test_lumina_builder", + Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ) + .options(options) + .build() + .unwrap(), + ) +} + +fn test_table_with_schema(schema: Schema) -> Table { + test_table_with_io( + FileIOBuilder::new("memory").build().unwrap(), + "memory:/test_lumina_builder", + schema, + ) +} + +fn test_table_with_io(file_io: FileIO, table_path: &str, schema: Schema) -> Table { + Table::new( + file_io, + Identifier::new("default", "test_table"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ) +} + +fn vector_schema_builder(options: HashMap) -> crate::spec::SchemaBuilder { + Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ) + .options(options) +} + +fn plan(entries: Vec, rows_per_shard: i64) -> Result> { + plan_with_indexed(entries, rows_per_shard, &[]) +} + +fn plan_with_indexed( + entries: Vec, + rows_per_shard: i64, + indexed: &[RowRange], +) -> Result> { + let table = test_table(table_options(&rows_per_shard.to_string())); + let core = CoreOptions::new(table.schema().options()); + plan_lumina_shards( + table.location(), + table.schema().partition_keys(), + table.schema().fields(), + &core, + 1, + entries, + rows_per_shard, + indexed, + ) +} + +#[test] +fn test_planner_splits_single_file_across_shards() { + let shards = plan(vec![manifest_entry(data_file("a", Some(0), 25))], 10).unwrap(); + + assert_eq!( + shards + .iter() + .map(|s| (s.row_range_start, s.row_range_end)) + .collect::>(), + vec![(0, 9), (10, 19), (20, 24)] + ); +} + +#[test] +fn test_planner_merges_contiguous_files() { + let shards = plan( + vec![ + manifest_entry(data_file("a", Some(0), 5)), + manifest_entry(data_file("b", Some(5), 5)), + ], + 20, + ) + .unwrap(); + + assert_eq!(shards.len(), 1); + assert_eq!((shards[0].row_range_start, shards[0].row_range_end), (0, 9)); + assert_eq!( + shards[0] + .files + .iter() + .map(|f| f.file_name.as_str()) + .collect::>(), + vec!["a", "b"] + ); +} + +#[test] +fn test_planner_keeps_source_buckets_separate() { + let shards = plan( + vec![ + manifest_entry_with_bucket(data_file("a", Some(0), 5), 0, 2), + manifest_entry_with_bucket(data_file("b", Some(5), 5), 1, 2), + ], + 20, + ) + .unwrap(); + + assert_eq!(shards.len(), 2); + assert_eq!( + shards + .iter() + .map(|s| ( + s.source_bucket, + s.total_buckets, + s.row_range_start, + s.row_range_end + )) + .collect::>(), + vec![(0, 2, 0, 4), (1, 2, 5, 9)] + ); +} + +#[test] +fn test_planner_splits_gap_into_separate_groups() { + let shards = plan( + vec![ + manifest_entry(data_file("a", Some(0), 5)), + manifest_entry(data_file("b", Some(10), 5)), + ], + 20, + ) + .unwrap(); + + assert_eq!( + shards + .iter() + .map(|s| (s.row_range_start, s.row_range_end)) + .collect::>(), + vec![(0, 4), (10, 14)] + ); +} + +#[test] +fn test_planner_rejects_missing_first_row_id() { + let err = plan(vec![manifest_entry(data_file("a", None, 5))], 10) + .expect_err("missing first_row_id should fail"); + assert!( + matches!(err, Error::DataInvalid { message, .. } if message.contains("missing first_row_id")) + ); +} + +#[test] +fn test_planner_rejects_invalid_rows_per_shard() { + let err = plan(vec![manifest_entry(data_file("a", Some(0), 5))], 0) + .expect_err("invalid rows per shard should fail"); + assert!( + matches!(err, Error::DataInvalid { message, .. } if message.contains("row-count-per-shard")) + ); +} + +#[test] +fn test_validate_vector_field_accepts_array_float() { + let field = DataField::new( + 0, + "embedding".to_string(), + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ); + assert!(validate_vector_field(&field).is_ok()); +} + +#[test] +fn test_validate_vector_field_accepts_vector_float() { + let field = DataField::new( + 0, + "embedding".to_string(), + DataType::Vector(VectorType::try_new(true, 4, DataType::Float(FloatType::new())).unwrap()), + ); + assert!(validate_vector_field(&field).is_ok()); +} + +#[test] +fn test_validate_vector_field_rejects_vector_double() { + let field = DataField::new( + 0, + "embedding".to_string(), + DataType::Vector( + VectorType::try_new(true, 4, DataType::Double(DoubleType::new())).unwrap(), + ), + ); + let err = validate_vector_field(&field).expect_err("VECTOR must be rejected"); + assert!(matches!(err, Error::DataInvalid { .. })); +} + +#[test] +fn test_effective_options_vector_absent_inserts_length() { + let field = DataField::new( + 0, + "embedding".to_string(), + DataType::Vector( + VectorType::try_new(true, 256, DataType::Float(FloatType::new())).unwrap(), + ), + ); + let opts = effective_lumina_options(&field, HashMap::new()).unwrap(); + assert_eq!( + opts.get(LUMINA_DIMENSION_OPTION).map(String::as_str), + Some("256") + ); +} + +#[test] +fn test_effective_options_vector_matching_option_ok() { + let field = DataField::new( + 0, + "embedding".to_string(), + DataType::Vector( + VectorType::try_new(true, 256, DataType::Float(FloatType::new())).unwrap(), + ), + ); + let resolved = HashMap::from([(LUMINA_DIMENSION_OPTION.to_string(), "256".to_string())]); + let opts = effective_lumina_options(&field, resolved).unwrap(); + assert_eq!( + opts.get(LUMINA_DIMENSION_OPTION).map(String::as_str), + Some("256") + ); +} + +#[test] +fn test_effective_options_vector_mismatch_errors() { + let field = DataField::new( + 0, + "embedding".to_string(), + DataType::Vector( + VectorType::try_new(true, 256, DataType::Float(FloatType::new())).unwrap(), + ), + ); + let resolved = HashMap::from([(LUMINA_DIMENSION_OPTION.to_string(), "128".to_string())]); + let err = + effective_lumina_options(&field, resolved).expect_err("dimension mismatch must error"); + assert!(matches!(err, Error::ConfigInvalid { .. })); +} + +#[test] +fn test_effective_options_array_unchanged() { + let field = DataField::new( + 0, + "embedding".to_string(), + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ); + // No dimension option set: array path must NOT inject one. + let opts = effective_lumina_options(&field, HashMap::new()).unwrap(); + assert!(!opts.contains_key(LUMINA_DIMENSION_OPTION)); +} + +#[test] +fn test_vector_without_option_propagates_dimension_to_native_options() { + use crate::lumina::{LuminaVectorIndexOptions, KEY_DIMENSION}; + let field = DataField::new( + 0, + "embedding".to_string(), + DataType::Vector( + VectorType::try_new(true, 256, DataType::Float(FloatType::new())).unwrap(), + ), + ); + // No lumina.index.dimension set by the user. + let resolved = effective_lumina_options(&field, HashMap::new()).unwrap(); + let opts = LuminaVectorIndexOptions::new(&resolved).unwrap(); + + assert_eq!( + opts.dimension, 256, + "local dimension must be N, not default 128" + ); + assert_eq!( + opts.to_lumina_options() + .get(KEY_DIMENSION) + .map(String::as_str), + Some("256"), + "native index.dimension must be N, not default 128" + ); +} + +#[tokio::test] +async fn test_execute_rejects_primary_key_table() { + let table = test_table_with_schema( + vector_schema_builder(HashMap::new()) + .primary_key(["id"]) + .build() + .unwrap(), + ); + + let err = table + .new_lumina_index_build_builder() + .with_index_column("embedding") + .execute() + .await + .expect_err("primary-key table should fail before native build"); + + assert!( + matches!(err, Error::Unsupported { message } if message.contains("primary-key tables")) + ); +} + +#[tokio::test] +async fn test_execute_rejects_deletion_vectors_table() { + let mut options = table_options("10"); + options.insert("deletion-vectors.enabled".to_string(), "true".to_string()); + let table = test_table(options); + + let err = table + .new_lumina_index_build_builder() + .with_index_column("embedding") + .execute() + .await + .expect_err("deletion vectors table should fail before native build"); + + assert!( + matches!(err, Error::Unsupported { message } if message.contains("deletion-vectors.enabled=true")) + ); +} + +fn vector_batch(rows: Vec>>>, row_ids: Vec>) -> RecordBatch { + let mut vector_builder = ListBuilder::new(Float32Builder::new()); + for row in rows { + match row { + Some(values) => { + for value in values { + match value { + Some(value) => vector_builder.values().append_value(value), + None => vector_builder.values().append_null(), + } + } + vector_builder.append(true); + } + None => vector_builder.append(false), + } + } + let mut row_id_builder = Int64Builder::new(); + for row_id in row_ids { + match row_id { + Some(value) => row_id_builder.append_value(value), + None => row_id_builder.append_null(), + } + } + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "embedding", + ArrowDataType::List(Arc::new(ArrowField::new( + "item", + ArrowDataType::Float32, + true, + ))), + true, + ), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(vector_builder.finish()) as ArrayRef, + Arc::new(row_id_builder.finish()) as ArrayRef, + ], + ) + .unwrap() +} + +#[test] +fn test_extract_vectors_accepts_list_float32_and_row_ids() { + let batch = vector_batch( + vec![ + Some(vec![Some(1.0), Some(2.0)]), + Some(vec![Some(3.0), Some(4.0)]), + ], + vec![Some(10), Some(11)], + ); + + let vectors = extract_vectors_from_batches(&[batch], "embedding", 2, 10, 2).unwrap(); + + assert_eq!(vectors, vec![1.0, 2.0, 3.0, 4.0]); +} + +#[test] +fn test_extract_vectors_rejects_null_vector() { + let batch = vector_batch(vec![None], vec![Some(0)]); + + let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) + .expect_err("null vector should fail"); + + assert!(matches!(err, Error::DataInvalid { message, .. } if message.contains("null vector"))); +} + +#[test] +fn test_extract_vectors_rejects_null_element() { + let batch = vector_batch(vec![Some(vec![Some(1.0), None])], vec![Some(0)]); + + let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) + .expect_err("null element should fail"); + + assert!( + matches!(err, Error::DataInvalid { message, .. } if message.contains("null vector element")) + ); +} + +#[test] +fn test_extract_vectors_rejects_dimension_mismatch() { + let batch = vector_batch(vec![Some(vec![Some(1.0)])], vec![Some(0)]); + + let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) + .expect_err("dimension mismatch should fail"); + + assert!( + matches!(err, Error::DataInvalid { message, .. } if message.contains("dimension mismatch")) + ); +} + +#[test] +fn test_extract_vectors_rejects_row_id_gap() { + let batch = vector_batch( + vec![ + Some(vec![Some(1.0), Some(2.0)]), + Some(vec![Some(3.0), Some(4.0)]), + ], + vec![Some(0), Some(2)], + ); + + let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 2) + .expect_err("row id gap should fail"); + + assert!( + matches!(err, Error::DataInvalid { message, .. } if message.contains("expected _ROW_ID")) + ); +} + +#[test] +fn test_extract_vectors_rejects_non_list_float32() { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("embedding", ArrowDataType::Int32, false), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + Arc::new(Int64Array::from(vec![Some(0)])) as ArrayRef, + ], + ) + .unwrap(); + + let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) + .expect_err("non-list vector should fail"); + + assert!(matches!(err, Error::DataInvalid { message, .. } if message.contains("List"))); +} + +fn fixed_size_vector_batch( + rows: Vec>>, + row_ids: Vec>, + len: i32, +) -> RecordBatch { + let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut builder = + FixedSizeListBuilder::new(Float32Builder::new(), len).with_field(element_field); + for row in rows { + match row { + Some(values) => { + for v in values { + builder.values().append_value(v); + } + builder.append(true); + } + None => { + for _ in 0..len { + builder.values().append_value(0.0); + } + builder.append(false); + } + } + } + let mut row_id_builder = Int64Builder::new(); + for row_id in row_ids { + match row_id { + Some(value) => row_id_builder.append_value(value), + None => row_id_builder.append_null(), + } + } + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "embedding", + ArrowDataType::FixedSizeList( + Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)), + len, + ), + true, + ), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(builder.finish()) as ArrayRef, + Arc::new(row_id_builder.finish()) as ArrayRef, + ], + ) + .unwrap() +} + +#[test] +fn test_extract_vectors_accepts_fixed_size_list_float32() { + let batch = fixed_size_vector_batch( + vec![Some(vec![1.0, 2.0]), Some(vec![3.0, 4.0])], + vec![Some(10), Some(11)], + 2, + ); + let vectors = extract_vectors_from_batches(&[batch], "embedding", 2, 10, 2).unwrap(); + assert_eq!(vectors, vec![1.0, 2.0, 3.0, 4.0]); +} + +#[test] +fn test_extract_vectors_fixed_size_list_rejects_null_vector() { + let batch = fixed_size_vector_batch(vec![None], vec![Some(0)], 2); + let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) + .expect_err("null vector should fail"); + assert!(matches!(err, Error::DataInvalid { message, .. } if message.contains("null vector"))); +} + +#[test] +fn test_extract_vectors_fixed_size_list_dimension_mismatch() { + // Column is FixedSizeList of length 3, but caller expects dimension 2. + let batch = fixed_size_vector_batch(vec![Some(vec![1.0, 2.0, 3.0])], vec![Some(0)], 3); + let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) + .expect_err("dimension mismatch should fail"); + assert!( + matches!(err, Error::DataInvalid { message, .. } if message.contains("dimension mismatch")) + ); +} + +#[test] +fn test_extract_vectors_fixed_size_list_rejects_null_element() { + // A non-null vector row whose second child element is null. Mirrors the + // List path's test_extract_vectors_rejects_null_element so both layouts + // reject null elements identically. + let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(element_field); + builder.values().append_value(1.0); + builder.values().append_null(); + builder.append(true); + let row_ids = Arc::new(Int64Array::from(vec![Some(0)])) as ArrayRef; + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "embedding", + ArrowDataType::FixedSizeList( + Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)), + 2, + ), + true, + ), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(builder.finish()) as ArrayRef, row_ids], + ) + .unwrap(); + + let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) + .expect_err("null element should fail"); + + assert!( + matches!(err, Error::DataInvalid { message, .. } if message.contains("null vector element")) + ); +} + +#[test] +fn test_checked_metadata_conversion_supports_i64_file_size() { + let above_i32_max = i32::MAX as u64 + 1; + assert_eq!( + checked_i64(above_i32_max, "Index file is too large").unwrap(), + i64::from(i32::MAX) + 1 + ); + + let err = checked_i64(i64::MAX as u64 + 1, "Index file is too large") + .expect_err("file size above i64::MAX should fail"); + assert!(matches!(err, Error::DataInvalid { message, .. } if message.contains("too large"))); +} + +#[test] +fn test_temp_file_guard_cleans_up_on_drop() { + let path = temp_lumina_path(); + std::fs::write(&path, b"temporary lumina data").unwrap(); + { + let _guard = TempFileGuard::new(path.clone()); + assert!(path.exists()); + } + assert!(!path.exists()); +} + +#[tokio::test] +async fn test_abort_on_build_error_removes_completed_index_files() { + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let table_path = "memory:/test_lumina_abort_cleanup"; + let table = test_table_with_io( + file_io.clone(), + table_path, + vector_schema_builder(table_options("1")).build().unwrap(), + ); + file_io + .mkdirs(&format!("{table_path}/index/")) + .await + .unwrap(); + + let file_name = "completed-lumina-shard.index"; + let index_path = format!("{table_path}/index/{file_name}"); + file_io + .new_output(&index_path) + .unwrap() + .write(Bytes::from_static(b"completed shard")) + .await + .unwrap(); + let mut message = CommitMessage::new(vec![], 0, vec![]); + message.new_index_files = vec![IndexFileMeta { + index_type: LUMINA_IDENTIFIER.to_string(), + file_name: file_name.to_string(), + file_size: 15, + row_count: 1, + deletion_vectors_ranges: None, + global_index_meta: None, + }]; + + let commit = TableCommit::new(table, "test-lumina-abort".to_string()); + let result = abort_on_build_error::<()>( + &commit, + &[message], + Err(Error::UnexpectedError { + message: "second shard failed".to_string(), + source: None, + }), + ) + .await; + + assert!(result.is_err()); + assert!(file_io.get_status(&index_path).await.is_err()); +} + +async fn setup_dirs(file_io: &FileIO, table_path: &str) { + file_io + .mkdirs(&format!("{table_path}/snapshot/")) + .await + .unwrap(); + file_io + .mkdirs(&format!("{table_path}/manifest/")) + .await + .unwrap(); +} + +fn build_vector_batch(ids: Vec, vectors: Vec>) -> RecordBatch { + let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut vector_builder = + ListBuilder::new(Float32Builder::new()).with_field(element_field.clone()); + for vector in vectors { + for value in vector { + vector_builder.values().append_value(value); + } + vector_builder.append(true); + } + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("embedding", ArrowDataType::List(element_field), true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(ids)) as ArrayRef, + Arc::new(vector_builder.finish()) as ArrayRef, + ], + ) + .unwrap() +} + +// Manual run with a local Lumina native library: +// LUMINA_LIB_PATH=/path/to/liblumina_py.so cargo test -p paimon \ +// table::lumina_index_build_builder::tests::test_execute_writes_lumina_index_manifest \ +// --features fulltext,vortex -- --ignored --exact +#[tokio::test] +#[ignore = "requires LUMINA_LIB_PATH; see manual run command above"] +async fn test_execute_writes_lumina_index_manifest() { + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let table_path = "memory:/test_lumina_builder_e2e"; + setup_dirs(&file_io, table_path).await; + + let mut options = table_options("10"); + options.insert("lumina.index.dimension".to_string(), "2".to_string()); + options.insert("lumina.encoding.type".to_string(), "rawf32".to_string()); + let table = test_table_with_io( + file_io.clone(), + table_path, + vector_schema_builder(options).build().unwrap(), + ); + + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&build_vector_batch( + vec![1, 2], + vec![vec![1.0, 0.0], vec![0.0, 1.0]], + )) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let shard_count = table + .new_lumina_index_build_builder() + .with_index_column("embedding") + .execute() + .await + .unwrap(); + assert_eq!(shard_count, 1); + + let snapshot_manager = SnapshotManager::new(file_io.clone(), table_path.to_string()); + let snapshot = snapshot_manager + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + assert_eq!(snapshot.id(), 2); + assert_eq!(snapshot.total_record_count(), Some(2)); + assert_eq!(snapshot.delta_record_count(), Some(0)); + assert_eq!(snapshot.next_row_id(), Some(2)); + + let index_manifest = snapshot.index_manifest().expect("index manifest"); + let index_entries = + IndexManifest::read(&file_io, &format!("{table_path}/manifest/{index_manifest}")) + .await + .unwrap(); + assert_eq!(index_entries.len(), 1); + + let index_file = &index_entries[0].index_file; + assert_eq!(index_file.index_type, LUMINA_IDENTIFIER); + assert!(index_file.file_name.starts_with("lumina-global-index-")); + assert_eq!(index_file.row_count, 2); + assert!(index_file.file_size > 0); + + let global_meta = index_file + .global_index_meta + .as_ref() + .expect("global index meta"); + assert_eq!(global_meta.row_range_start, 0); + assert_eq!(global_meta.row_range_end, 1); + assert_eq!(global_meta.index_field_id, 1); + assert!(global_meta + .index_meta + .as_ref() + .is_some_and(|m| !m.is_empty())); + + let index_path = format!("{table_path}/index/{}", index_file.file_name); + let status = file_io.get_status(&index_path).await.unwrap(); + assert_eq!(index_file.file_size as u64, status.size); +} + +fn lumina_e2e_options(rows_per_shard: &str) -> HashMap { + let mut options = table_options(rows_per_shard); + options.insert("lumina.index.dimension".to_string(), "2".to_string()); + options.insert("lumina.encoding.type".to_string(), "rawf32".to_string()); + options +} + +fn lumina_e2e_table(table_path: &str, rows_per_shard: &str) -> Table { + test_table_with_io( + FileIOBuilder::new("memory").build().unwrap(), + table_path, + vector_schema_builder(lumina_e2e_options(rows_per_shard)) + .build() + .unwrap(), + ) +} + +async fn write_vectors(table: &Table, ids: Vec, vectors: Vec>) { + let mut table_write = TableWrite::new(table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&build_vector_batch(ids, vectors)) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); +} + +/// Commit a synthetic Lumina `IndexFileMeta` covering `[start, end]` for +/// `field_id` directly into the index manifest, without invoking the native +/// builder. Mirrors the btree mid-hole test so the incremental gap logic can +/// be exercised in CI where the native Lumina library is unavailable. +async fn commit_synthetic_lumina_index(table: &Table, field_id: i32, start: i64, end: i64) { + let synthetic = IndexFileMeta { + index_type: LUMINA_IDENTIFIER.to_string(), + file_name: format!("lumina-synthetic-{start}-{end}.index"), + file_size: 1, + row_count: (end - start + 1), + deletion_vectors_ranges: None, + global_index_meta: Some(GlobalIndexMeta { + row_range_start: start, + row_range_end: end, + index_field_id: field_id, + extra_field_ids: None, + source_meta: None, + index_meta: None, + }), + }; + let mut message = CommitMessage::new(BinaryRow::new(0).to_serialized_bytes(), 0, vec![]); + message.new_index_files = vec![synthetic]; + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(vec![message]) + .await + .unwrap(); +} + +async fn latest_lumina_index_files(table: &Table) -> Vec { + let snapshot_manager = + SnapshotManager::new(table.file_io().clone(), table.location().to_string()); + let snapshot = snapshot_manager + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let Some(index_manifest_name) = snapshot.index_manifest() else { + return Vec::new(); + }; + IndexManifest::read( + table.file_io(), + &snapshot_manager.manifest_path(index_manifest_name), + ) + .await + .unwrap() + .into_iter() + .filter(|entry| entry.kind == FileKind::Add && entry.index_file.index_type == LUMINA_IDENTIFIER) + .map(|entry| entry.index_file) + .collect() +} + +/// Row-id coverage of the committed data files, read back from the data +/// manifest (never hard-coded) and merged into contiguous ranges. Mirrors +/// how `execute` gathers `manifest_entries`. +async fn data_row_id_coverage(table: &Table) -> Vec { + let snapshot_manager = + SnapshotManager::new(table.file_io().clone(), table.location().to_string()); + let snapshot = snapshot_manager + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let entries = table + .new_read_builder() + .new_scan() + .with_scan_all_files() + .plan_manifest_entries(&snapshot) + .await + .unwrap(); + let ranges = entries + .iter() + .filter(|entry| *entry.kind() == FileKind::Add) + .filter_map(|entry| { + entry + .file() + .row_id_range() + .map(|(start, end)| RowRange::new(start, end)) + }) + .collect::>(); + crate::table::merge_row_ranges(ranges) +} + +/// Second build with the whole coverage already indexed must be a clean +/// no-op (returns 0), not an overlap error. Reaches `Ok(0)` before the +/// native build, so it runs in CI without the Lumina library. This is the +/// core bug fix: today the second call errors with the overlap message. +#[tokio::test] +async fn lumina_second_build_without_new_data_is_noop() { + let table_path = "memory:/test_lumina_second_build_noop"; + let table = lumina_e2e_table(table_path, "10"); + setup_dirs(table.file_io(), table_path).await; + + write_vectors( + &table, + vec![1, 2, 3], + vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]], + ) + .await; + + // Fully index the coverage via a synthetic manifest entry. + let coverage = data_row_id_coverage(&table).await; + assert_eq!(coverage.len(), 1, "data must be one contiguous range"); + let field_id = find_index_field(&table, "embedding").unwrap().id(); + commit_synthetic_lumina_index(&table, field_id, coverage[0].from(), coverage[0].to()).await; + + let names_before = latest_lumina_index_files(&table) + .await + .iter() + .map(|f| f.file_name.clone()) + .collect::>(); + assert!(!names_before.is_empty()); + + let built = table + .new_lumina_index_build_builder() + .with_index_column("embedding") + .execute() + .await + .unwrap(); + assert_eq!(built, 0, "fully-indexed table must build nothing on re-run"); + + let names_after = latest_lumina_index_files(&table) + .await + .iter() + .map(|f| f.file_name.clone()) + .collect::>(); + assert_eq!( + names_before, names_after, + "re-run must not add or remove index manifest entries" + ); +} + +/// Build over an already-indexed prefix, then append new rows: the second +/// build must target only the appended gap and must NOT fail with the +/// overlap error. Without the native Lumina library the gap build surfaces a +/// library-load error (not the overlap error); with it present it succeeds. +#[tokio::test] +async fn lumina_incremental_build_indexes_only_new_rows() { + let table_path = "memory:/test_lumina_incremental"; + let table = lumina_e2e_table(table_path, "10"); + setup_dirs(table.file_io(), table_path).await; + + // Initial batch, then mark it fully indexed via a synthetic entry. + write_vectors( + &table, + vec![1, 2, 3], + vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]], + ) + .await; + let indexed_coverage = data_row_id_coverage(&table).await; + assert_eq!(indexed_coverage.len(), 1); + let n = indexed_coverage[0].to() + 1; + let field_id = find_index_field(&table, "embedding").unwrap().id(); + commit_synthetic_lumina_index( + &table, + field_id, + indexed_coverage[0].from(), + indexed_coverage[0].to(), + ) + .await; + + // Append a second batch (new row-ids [n..]). + write_vectors( + &table, + vec![4, 5, 6], + vec![vec![2.0, 0.0], vec![0.0, 2.0], vec![2.0, 2.0]], + ) + .await; + + // White-box: fed the real indexed ranges from the manifest, the planner + // must target only the appended gap [n, ..], never the already-indexed + // prefix. Computed before `execute` so it is independent of whether the + // native build (which needs the Lumina library) runs. + let snapshot_manager = + SnapshotManager::new(table.file_io().clone(), table.location().to_string()); + let snapshot = snapshot_manager + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let manifest_entries = table + .new_read_builder() + .new_scan() + .with_scan_all_files() + .plan_manifest_entries(&snapshot) + .await + .unwrap(); + let indexed = crate::table::global_index_build_common::indexed_row_ranges( + &table, + snapshot.index_manifest(), + LUMINA_IDENTIFIER, + field_id, + None, + ) + .await + .unwrap(); + let core = CoreOptions::new(table.schema().options()); + let shards = plan_lumina_shards( + table.location(), + table.schema().partition_keys(), + table.schema().fields(), + &core, + snapshot.id(), + manifest_entries, + 10, + &indexed, + ) + .unwrap(); + assert!(!shards.is_empty(), "appended gap must produce build shards"); + for shard in &shards { + assert!( + shard.row_range_start >= n, + "shard [{}, {}] must start at or after the indexed prefix end {n}", + shard.row_range_start, + shard.row_range_end + ); + } + + // End-to-end: the incremental build must no longer fail with the overlap + // error. Without the native Lumina library the gap build surfaces a + // library-load error instead; with it present it succeeds. + let result = table + .new_lumina_index_build_builder() + .with_index_column("embedding") + .execute() + .await; + match result { + Ok(_) => {} + Err(Error::DataInvalid { message, .. }) => { + assert!( + !message.contains("overlaps requested row range"), + "incremental build must not fail with the overlap error; got: {message}" + ); + } + Err(other) => panic!("unexpected error from incremental build: {other:?}"), + } +} + +/// Regression: a first build (no existing index) must equal the pre-change +/// full build -- subtracting an empty `indexed` yields full coverage. +#[test] +fn lumina_first_build_indexes_full_coverage() { + let full = plan(vec![manifest_entry(data_file("a", Some(0), 25))], 10).unwrap(); + let gapped = + plan_with_indexed(vec![manifest_entry(data_file("a", Some(0), 25))], 10, &[]).unwrap(); + // Empty `indexed` must not alter the shard layout. + assert_eq!( + full.iter() + .map(|s| (s.row_range_start, s.row_range_end)) + .collect::>(), + gapped + .iter() + .map(|s| (s.row_range_start, s.row_range_end)) + .collect::>() + ); + assert_eq!( + full.iter() + .map(|s| (s.row_range_start, s.row_range_end)) + .collect::>(), + vec![(0, 9), (10, 19), (20, 24)], + "first build must cover the full row range across shards" + ); +} + +/// Planner-level mid-coverage hole, mirroring btree's +/// `incremental_build_splits_gap_around_mid_coverage_indexed_hole`: with a +/// single shard cell (rows_per_shard large enough to hold all data) the grid +/// never splits, so the only split is the indexed hole itself. An indexed +/// range strictly inside the data coverage must carve the build into exactly +/// the two contiguous segments on either side of the hole -- both bounds +/// pinned, and neither segment may span or touch the hole. +#[test] +fn lumina_plan_splits_gap_around_mid_coverage_indexed_hole() { + // Data row-ids [0, 9]; one shard cell [0, 99] so the grid never splits. + let n = 9; + let hole_start = 4; + let hole_end = 6; + let shards = plan_with_indexed( + vec![manifest_entry(data_file("a", Some(0), n + 1))], + 100, + &[RowRange::new(hole_start, hole_end)], + ) + .unwrap(); + + let ranges = shards + .iter() + .map(|s| (s.row_range_start, s.row_range_end)) + .collect::>(); + // Exactly the two contiguous segments around the hole. + assert_eq!( + ranges, + vec![(0, hole_start - 1), (hole_end + 1, n)], + "mid-coverage hole must split into exactly the two segments around it" + ); + // Every emitted range is contiguous and none spans or touches the hole. + for (start, end) in &ranges { + assert!(end >= start, "range must be non-empty: [{start}, {end}]"); + assert!( + *end < hole_start || *start > hole_end, + "shard [{start}, {end}] must not overlap indexed hole [{hole_start}, {hole_end}]" + ); + } + // Together the shards cover exactly coverage - indexed. + let expected = exclude_row_ranges( + &[RowRange::new(0, n)], + &[RowRange::new(hole_start, hole_end)], + ) + .into_iter() + .map(|r| (r.from(), r.to())) + .collect::>(); + assert_eq!( + ranges, expected, + "shards must cover exactly coverage minus the indexed hole" + ); +} + +/// Planner-level incremental prefix. Strengthens +/// `lumina_incremental_build_indexes_only_new_rows`, which asserted only a +/// one-sided lower bound (`row_range_start >= n`): an indexed prefix [0, k] +/// must leave EXACTLY the suffix [k+1, N] on both bounds, split along the +/// shard grid, with nothing re-indexed inside the prefix. +#[test] +fn lumina_plan_incremental_prefix_leaves_suffix() { + // Data row-ids [0, 24], rows_per_shard = 10 -> cells [0,9],[10,19],[20,29]. + // Indexed prefix [0, 9] fully fills the first cell, so the build must be + // exactly [10, 19] and [20, 24] (the suffix split along the grid). + let n = 24; + let k = 9; // prefix [0, k] == the first full shard cell + let shards = plan_with_indexed( + vec![manifest_entry(data_file("a", Some(0), n + 1))], + 10, + &[RowRange::new(0, k)], + ) + .unwrap(); + + let ranges = shards + .iter() + .map(|s| (s.row_range_start, s.row_range_end)) + .collect::>(); + assert_eq!( + ranges, + vec![(k + 1, 19), (20, n)], + "indexed prefix must leave exactly the suffix, split along the shard grid" + ); + // Both bounds pinned (this is what the one-sided existing check omits). + assert_eq!(ranges.first().unwrap().0, k + 1, "suffix must start at k+1"); + assert_eq!(ranges.last().unwrap().1, n, "suffix must end at N"); + // Contiguous, and no shard reaches back into the indexed prefix. + for pair in ranges.windows(2) { + assert_eq!( + pair[1].0, + pair[0].1 + 1, + "ranges must be contiguous: {:?} then {:?}", + pair[0], + pair[1] + ); + } + for (start, end) in &ranges { + assert!( + *start > k, + "shard [{start}, {end}] must not re-index the prefix [0, {k}]" + ); + } +} diff --git a/crates/paimon/src/table/lumina_index_build_builder/validation.rs b/crates/paimon/src/table/lumina_index_build_builder/validation.rs new file mode 100644 index 000000000..ec563571c --- /dev/null +++ b/crates/paimon/src/table/lumina_index_build_builder/validation.rs @@ -0,0 +1,135 @@ +// 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. + +//! Lumina table, field, option, and numeric validation. + +use crate::lumina::{LuminaVectorIndexOptions, LUMINA_DIMENSION_OPTION}; +use crate::spec::{CoreOptions, DataField, DataType}; +use crate::table::global_index_build_common::vector::{ + find_index_field as find_vector_index_field, + validate_vector_field as validate_common_vector_field, +}; +use crate::table::Table; +use crate::{Error, Result}; +use std::collections::HashMap; + +pub(super) fn validate_table_options(table: &Table, core_options: &CoreOptions) -> Result<()> { + if !table.schema().primary_keys().is_empty() { + return Err(Error::Unsupported { + message: "Lumina index build does not support primary-key tables".to_string(), + }); + } + if !core_options.row_tracking_enabled() { + return Err(Error::DataInvalid { + message: "Lumina index build requires 'row-tracking.enabled' = 'true'".to_string(), + source: None, + }); + } + if !core_options.data_evolution_enabled() { + return Err(Error::DataInvalid { + message: "Lumina index build requires 'data-evolution.enabled' = 'true'".to_string(), + source: None, + }); + } + if !core_options.global_index_enabled() { + return Err(Error::DataInvalid { + message: "Lumina index build requires 'global-index.enabled' = 'true'".to_string(), + source: None, + }); + } + if core_options.deletion_vectors_enabled() { + return Err(Error::Unsupported { + message: + "Lumina index build does not support tables with deletion-vectors.enabled=true" + .to_string(), + }); + } + Ok(()) +} + +pub(super) fn find_index_field<'a>(table: &'a Table, column: &str) -> Result<&'a DataField> { + find_vector_index_field(table, column) +} + +pub(super) fn validate_vector_field(field: &DataField) -> Result<()> { + validate_common_vector_field(field, "Lumina") +} + +pub(super) fn effective_lumina_options( + field: &DataField, + mut resolved: HashMap, +) -> Result> { + let DataType::Vector(vector) = field.data_type() else { + return Ok(resolved); + }; + let dimension = vector.length().to_string(); + match resolved.get(LUMINA_DIMENSION_OPTION) { + None => { + resolved.insert(LUMINA_DIMENSION_OPTION.to_string(), dimension); + } + Some(existing) if *existing == dimension => {} + Some(existing) => { + return Err(Error::ConfigInvalid { + message: format!( + "Vector column '{}' has dimension {} from its type, but '{}' is set to '{}'. \ + Remove the option or set it to {}.", + field.name(), + dimension, + LUMINA_DIMENSION_OPTION, + existing, + dimension + ), + }); + } + } + Ok(resolved) +} + +pub(super) fn resolve_lumina_options( + table_options: &HashMap, + user_options: &HashMap, +) -> Result> { + let mut options = table_options.clone(); + options.extend(user_options.clone()); + LuminaVectorIndexOptions::new(&options)?; + Ok(options) +} + +pub(super) fn checked_i64(value: u64, context: &str) -> Result { + i64::try_from(value).map_err(|_| Error::DataInvalid { + message: format!("{context}: {value}"), + source: None, + }) +} + +pub(super) fn checked_row_count(row_range_start: i64, row_range_end: i64) -> Result { + if row_range_end < row_range_start { + return Err(Error::DataInvalid { + message: format!("Invalid Lumina row range [{row_range_start}, {row_range_end}]"), + source: None, + }); + } + row_range_end + .checked_sub(row_range_start) + .and_then(|span| span.checked_add(1)) + .ok_or_else(|| Error::DataInvalid { + message: format!( + "Row count overflows for row range [{row_range_start}, {row_range_end}]" + ), + source: None, + }) +} diff --git a/crates/paimon/src/table/lumina_index_build_builder/writer.rs b/crates/paimon/src/table/lumina_index_build_builder/writer.rs new file mode 100644 index 000000000..a3bd41366 --- /dev/null +++ b/crates/paimon/src/table/lumina_index_build_builder/writer.rs @@ -0,0 +1,221 @@ +// 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. + +//! Lumina index-file construction, upload, and cleanup. + +use super::planning::LuminaIndexShard; +use super::validation::{checked_i64, checked_row_count}; +use super::LuminaIndexBuildBuilder; +use crate::lumina::ffi::LuminaBuilder; +use crate::lumina::{LuminaIndexMeta, LUMINA_IDENTIFIER}; +use crate::spec::{GlobalIndexMeta, IndexFileMeta}; +use crate::table::{CommitMessage, TableCommit}; +use crate::{Error, Result}; +use bytes::Bytes; +use std::path::{Path, PathBuf}; +use tokio::io::AsyncReadExt; + +const INDEX_DIR: &str = "index"; +const COPY_BUFFER_SIZE: usize = 1024 * 1024; + +impl LuminaIndexBuildBuilder<'_> { + pub(super) async fn build_index_file( + &self, + shard: &LuminaIndexShard, + vectors: &[f32], + dimension: i32, + index_field_id: i32, + index_meta: Vec, + ) -> Result { + let row_count = checked_row_count(shard.row_range_start, shard.row_range_end)?; + // The native Lumina builder counts rows in an i32; the manifest keeps the + // full width. + let native_row_count = i32::try_from(row_count).map_err(|_| Error::DataInvalid { + message: format!( + "Lumina shard row count {row_count} exceeds what the native builder accepts" + ), + source: None, + })?; + validate_vector_buffer(vectors, native_row_count, dimension)?; + let ids = (0..row_count as u64).collect::>(); + let native_options = LuminaIndexMeta::deserialize(&index_meta)?.options().clone(); + + let temp_path = temp_lumina_path(); + let temp_file = TempFileGuard::new(temp_path.clone()); + let temp_path_str = temp_path.to_string_lossy().to_string(); + let builder = LuminaBuilder::create(&native_options)?; + builder.pretrain(vectors, native_row_count, dimension)?; + builder.insert(vectors, &ids, native_row_count, dimension)?; + builder.dump(&temp_path_str)?; + + let file_name = format!("lumina-global-index-{}.index", uuid::Uuid::new_v4()); + self.table + .file_io() + .mkdirs(&format!( + "{}/{INDEX_DIR}/", + self.table.location().trim_end_matches('/') + )) + .await?; + let index_path = format!( + "{}/{INDEX_DIR}/{}", + self.table.location().trim_end_matches('/'), + file_name + ); + let write_result: Result = async { + copy_local_file_to_output(&temp_path, self.table.file_io().new_output(&index_path)?) + .await?; + temp_file.cleanup(); + let status = self.table.file_io().get_status(&index_path).await?; + checked_i64( + status.size, + "Index file is too large for Rust IndexFileMeta", + ) + } + .await; + let file_size = match write_result { + Ok(file_size) => file_size, + Err(error) => { + let _ = self.table.file_io().delete_file(&index_path).await; + return Err(error); + } + }; + Ok(IndexFileMeta { + index_type: LUMINA_IDENTIFIER.to_string(), + file_name, + file_size, + row_count, + deletion_vectors_ranges: None, + global_index_meta: Some(GlobalIndexMeta { + row_range_start: shard.row_range_start, + row_range_end: shard.row_range_end, + index_field_id, + extra_field_ids: None, + source_meta: None, + index_meta: Some(index_meta), + }), + }) + } +} + +pub(super) async fn abort_on_build_error( + commit: &TableCommit, + messages: &[CommitMessage], + result: Result, +) -> Result { + match result { + Ok(value) => Ok(value), + Err(error) => { + let _ = commit.abort(messages).await; + Err(error) + } + } +} + +fn validate_vector_buffer(vectors: &[f32], row_count: i32, dimension: i32) -> Result<()> { + if row_count <= 0 { + return Err(Error::DataInvalid { + message: format!("Lumina shard row count must be positive, got: {row_count}"), + source: None, + }); + } + if dimension <= 0 { + return Err(Error::DataInvalid { + message: format!("Lumina vector dimension must be positive, got: {dimension}"), + source: None, + }); + } + let row_count = row_count as usize; + let dimension = dimension as usize; + let expected_len = row_count + .checked_mul(dimension) + .ok_or_else(|| Error::DataInvalid { + message: format!( + "Lumina vector buffer length overflows: row_count={row_count}, dimension={dimension}" + ), + source: None, + })?; + if vectors.len() != expected_len { + return Err(Error::DataInvalid { + message: format!( + "Lumina vector buffer length {} does not match row_count={} and dimension={}", + vectors.len(), + row_count, + dimension + ), + source: None, + }); + } + Ok(()) +} + +pub(super) fn temp_lumina_path() -> PathBuf { + std::env::temp_dir().join(format!("lumina-index-{}.index", uuid::Uuid::new_v4())) +} + +pub(super) struct TempFileGuard { + path: Option, +} + +impl TempFileGuard { + pub(super) fn new(path: PathBuf) -> Self { + Self { path: Some(path) } + } + + fn cleanup(mut self) { + if let Some(path) = self.path.take() { + let _ = std::fs::remove_file(path); + } + } +} + +impl Drop for TempFileGuard { + fn drop(&mut self) { + if let Some(path) = self.path.take() { + let _ = std::fs::remove_file(path); + } + } +} + +async fn copy_local_file_to_output( + source_path: &Path, + output: crate::io::OutputFile, +) -> Result<()> { + let mut source = + tokio::fs::File::open(source_path) + .await + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to open temporary Lumina index file: {e}"), + source: None, + })?; + let mut writer = output.writer().await?; + let mut buffer = vec![0u8; COPY_BUFFER_SIZE]; + + loop { + let len = source + .read(&mut buffer) + .await + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to read temporary Lumina index file: {e}"), + source: None, + })?; + if len == 0 { + break; + } + writer.write(Bytes::copy_from_slice(&buffer[..len])).await?; + } + writer.close().await +} diff --git a/crates/paimon/src/table/sorted_global_index_build_builder.rs b/crates/paimon/src/table/sorted_global_index_build_builder.rs index f785767f1..b83761f66 100644 --- a/crates/paimon/src/table/sorted_global_index_build_builder.rs +++ b/crates/paimon/src/table/sorted_global_index_build_builder.rs @@ -15,32 +15,30 @@ // specific language governing permissions and limitations // under the License. +mod extraction; +mod planning; +mod validation; +mod writer; + +#[cfg(test)] +use extraction::*; +use planning::plan_sorted_index_shards; +#[cfg(test)] +use planning::SortedGlobalIndexShard; +use validation::*; + use super::bitmap_global_index_format::{make_bitmap_key_comparator, serialize_bitmap_datum}; -use super::bitmap_global_index_writer::{BitmapGlobalIndexWriter, BitmapWriteResult}; use super::global_index_types::{ normalize_queryable_global_index_type, BITMAP_GLOBAL_INDEX_TYPE, BTREE_GLOBAL_INDEX_TYPE, - FM_GLOBAL_INDEX_TYPE, MULTIVALUE_GLOBAL_INDEX_TYPE, + FM_GLOBAL_INDEX_TYPE, }; use super::sorted_global_index_options::SortedIndexWriteOptions; use crate::btree::key_serde::KeyComparator; -use crate::btree::{make_key_comparator, serialize_datum, BTreeIndexWriter}; -use crate::fm_index::{FMGlobalIndexWriter, FMOptions, FMWriteOptions}; -use crate::io::FileWrite; -use crate::spec::{ - bucket_dir_name, extract_datum_from_array, extract_datum_from_arrow, BinaryRow, CoreOptions, - DataField, DataFileMeta, DataType, Datum, FileKind, GlobalIndexMeta, IndexFileMeta, - ROW_ID_FIELD_NAME, -}; -use crate::table::source::exclude_row_ranges; -use crate::table::source::is_data_evolution_normal_file; -use crate::table::stats_filter::group_by_overlapping_row_id; -use crate::table::{ - CommitMessage, DataSplit, DataSplitBuilder, RowRange, SnapshotManager, Table, TableCommit, -}; +use crate::btree::{make_key_comparator, serialize_datum}; +use crate::fm_index::{FMOptions, FMWriteOptions}; +use crate::spec::{CoreOptions, DataType, Datum}; +use crate::table::{CommitMessage, RowRange, SnapshotManager, Table, TableCommit}; use crate::{Error, Result}; -use arrow_array::{Array, FixedSizeListArray, Int64Array, LargeListArray, ListArray, RecordBatch}; -use futures::TryStreamExt; -use std::cmp::Ordering; use std::collections::HashMap; const INDEX_DIR: &str = "index"; @@ -227,3317 +225,7 @@ impl<'a> SortedGlobalIndexBuildBuilder<'a> { Ok(shard_count) } - - async fn build_index_file( - &self, - shard: &SortedGlobalIndexShard, - index_field: &DataField, - index_column: &str, - write_options: &GlobalIndexWriteOptions, - ) -> Result { - let index_type = normalize_queryable_global_index_type(&self.index_type).ok_or_else(|| { - Error::Unsupported { - message: format!( - "Scalar global index build only supports index_type => 'btree', 'bitmap', 'multivalue', or 'fm', got '{}'", - self.index_type - ), - } - })?; - let row_count = checked_row_count(shard.row_range_start, shard.row_range_end)?; - let key_type = index_key_type(index_type, index_field)?; - let codec_type = if matches!( - index_type, - MULTIVALUE_GLOBAL_INDEX_TYPE | FM_GLOBAL_INDEX_TYPE - ) { - BITMAP_GLOBAL_INDEX_TYPE - } else { - index_type - }; - let (cmp, serialize_key) = make_index_key_codec(codec_type, key_type); - let mut rows = if index_type == FM_GLOBAL_INDEX_TYPE { - Vec::new() - } else { - extract_index_rows( - self.table, - shard, - index_column, - index_field, - index_type, - serialize_key, - ) - .await? - }; - if !rows.is_empty() { - sort_index_rows(&mut rows, &cmp); - } - - self.table - .file_io() - .mkdirs(&format!( - "{}/{INDEX_DIR}/", - self.table.location().trim_end_matches('/') - )) - .await?; - let file_name = format!("{index_type}-global-index-{}.index", uuid::Uuid::new_v4()); - let index_path = format!( - "{}/{INDEX_DIR}/{}", - self.table.location().trim_end_matches('/'), - file_name - ); - let write_result: Result<(u64, Vec, i64)> = async { - let output = self.table.file_io().new_output(&index_path)?; - let writer = output.writer().await?; - let (written_row_count, index_meta) = match index_type { - BTREE_GLOBAL_INDEX_TYPE => { - let GlobalIndexWriteOptions::Sorted(write_options) = write_options else { - unreachable!("BTree uses sorted write options") - }; - let mut writer = BTreeIndexWriter::with_comparator_and_compression_level( - writer, - write_options.block_size, - write_options.compression_type, - write_options.compression_level, - cmp, - ); - for (key, local_row_id) in &rows { - writer - .write(key.as_deref(), *local_row_id) - .await - .map_err(|e| Error::DataInvalid { - message: format!( - "Failed to write BTree global index file '{file_name}'" - ), - source: Some(Box::new(e)), - })?; - } - let write_result = writer.finish().await.map_err(|e| Error::DataInvalid { - message: format!("Failed to finish BTree global index file '{file_name}'"), - source: Some(Box::new(e)), - })?; - (write_result.row_count, write_result.meta.serialize()) - } - BITMAP_GLOBAL_INDEX_TYPE => { - let GlobalIndexWriteOptions::Sorted(write_options) = write_options else { - unreachable!("bitmap uses sorted write options") - }; - let mut writer = BitmapGlobalIndexWriter::with_compression_level( - writer, - write_options.block_size, - write_options.compression_type, - write_options.compression_level, - cmp, - ); - for (key, local_row_id) in &rows { - writer.write(key.as_deref(), *local_row_id).map_err(|e| { - Error::DataInvalid { - message: format!( - "Failed to write bitmap global index file '{file_name}'" - ), - source: Some(Box::new(e)), - } - })?; - } - let BitmapWriteResult { row_count, meta } = - writer.finish().await.map_err(|e| Error::DataInvalid { - message: format!( - "Failed to finish bitmap global index file '{file_name}'" - ), - source: Some(Box::new(e)), - })?; - (row_count, meta.serialize()) - } - MULTIVALUE_GLOBAL_INDEX_TYPE => { - let GlobalIndexWriteOptions::Sorted(write_options) = write_options else { - unreachable!("multivalue uses sorted write options") - }; - let mut writer = BitmapGlobalIndexWriter::with_compression_level( - writer, - write_options.block_size, - write_options.compression_type, - write_options.compression_level, - cmp, - ); - for (key, local_row_id) in &rows { - let key = key - .as_deref() - .expect("multivalue extraction skips null keys"); - writer.write_posting(key, *local_row_id).map_err(|e| { - Error::DataInvalid { - message: format!( - "Failed to write multivalue global index file '{file_name}'" - ), - source: Some(Box::new(e)), - } - })?; - } - let BitmapWriteResult { row_count, meta } = writer - .finish_with_source_row_count(u64::try_from(row_count).unwrap()) - .await - .map_err(|e| Error::DataInvalid { - message: format!( - "Failed to finish multivalue global index file '{file_name}'" - ), - source: Some(Box::new(e)), - })?; - (row_count, meta.serialize()) - } - FM_GLOBAL_INDEX_TYPE => { - let GlobalIndexWriteOptions::FM(write_options) = write_options else { - unreachable!("FM uses FM write options") - }; - write_fm_index_streaming( - self.table, - shard, - index_column, - index_field, - serialize_key, - writer, - *write_options, - &file_name, - ) - .await? - } - _ => unreachable!("normalized queryable global index type"), - }; - - if written_row_count != u64::try_from(row_count).unwrap() { - return Err(Error::DataInvalid { - message: format!( - "Sorted global index expected {} rows, wrote {}", - row_count, written_row_count - ), - source: None, - }); - } - - let status = self.table.file_io().get_status(&index_path).await?; - let file_size = checked_i64( - status.size, - "Index file is too large for Rust IndexFileMeta", - )?; - Ok((written_row_count, index_meta, file_size)) - } - .await; - let (_, index_meta, file_size) = match write_result { - Ok(result) => result, - Err(error) => { - let _ = self.table.file_io().delete_file(&index_path).await; - return Err(error); - } - }; - Ok(IndexFileMeta { - index_type: index_type.to_string(), - file_name, - file_size, - row_count, - deletion_vectors_ranges: None, - global_index_meta: Some(GlobalIndexMeta { - row_range_start: shard.row_range_start, - row_range_end: shard.row_range_end, - index_field_id: index_field.id(), - extra_field_ids: None, - source_meta: None, - index_meta: Some(index_meta), - }), - }) - } -} - -#[allow(clippy::too_many_arguments)] -async fn write_fm_index_streaming( - table: &Table, - shard: &SortedGlobalIndexShard, - index_column: &str, - index_field: &DataField, - serialize_key: SerializeKeyFn, - output: Box, - write_options: FMWriteOptions, - file_name: &str, -) -> Result<(u64, Vec)> { - let mut writer = - FMGlobalIndexWriter::new(output, write_options).map_err(|error| Error::DataInvalid { - message: format!("Failed to create FM global index file '{file_name}'"), - source: Some(Box::new(error)), - })?; - let splits = build_read_splits_for_shard(shard)?; - let mut read_builder = table.new_read_builder(); - read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; - let read = read_builder.new_read()?; - let mut batches = read.to_arrow(&splits)?; - let expected_row_count = checked_row_count(shard.row_range_start, shard.row_range_end)?; - let mut expected_row_id = shard.row_range_start; - - while let Some(batch) = batches.try_next().await? { - let value_index = - batch - .schema() - .index_of(index_column) - .map_err(|error| Error::DataInvalid { - message: format!( - "Index column '{index_column}' not found in FM read batch: {error}" - ), - source: None, - })?; - let row_id_index = batch - .schema() - .index_of(ROW_ID_FIELD_NAME) - .map_err(|error| Error::DataInvalid { - message: format!("_ROW_ID column not found in FM read batch: {error}"), - source: None, - })?; - let row_ids = batch - .column(row_id_index) - .as_any() - .downcast_ref::() - .ok_or_else(|| Error::DataInvalid { - message: "FM global index build requires non-null Int64 _ROW_ID".to_string(), - source: None, - })?; - - for row in 0..batch.num_rows() { - if row_ids.is_null(row) { - return Err(Error::DataInvalid { - message: "FM global index build found null _ROW_ID".to_string(), - source: None, - }); - } - let row_id = row_ids.value(row); - if row_id != expected_row_id { - return Err(Error::DataInvalid { - message: format!( - "FM global index build expected _ROW_ID {expected_row_id}, got {row_id}" - ), - source: None, - }); - } - expected_row_id = expected_row_id - .checked_add(1) - .ok_or_else(|| Error::DataInvalid { - message: "FM global index row ID overflow".to_string(), - source: None, - })?; - let local_row_id = - u64::try_from(row_id - shard.row_range_start).map_err(|_| Error::DataInvalid { - message: format!( - "FM global index file '{file_name}' has a negative local row ID" - ), - source: None, - })?; - let key = extract_datum_from_arrow(&batch, row, value_index, index_field.data_type())? - .map(|datum| serialize_key(&datum, index_field.data_type())); - writer - .write(key.as_deref(), local_row_id) - .await - .map_err(|error| Error::DataInvalid { - message: format!("Failed to write FM global index file '{file_name}'"), - source: Some(Box::new(error)), - })?; - } - } - - let actual_row_count = expected_row_id - shard.row_range_start; - if actual_row_count != expected_row_count { - return Err(Error::DataInvalid { - message: format!( - "FM global index build expected {expected_row_count} rows, got {actual_row_count}" - ), - source: None, - }); - } - let result = writer.finish().await.map_err(|error| Error::DataInvalid { - message: format!("Failed to finish FM global index file '{file_name}'"), - source: Some(Box::new(error)), - })?; - Ok((result.row_count, result.index_meta)) -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct SortedGlobalIndexShard { - pub partition: BinaryRow, - pub partition_bytes: Vec, - pub files: Vec, - pub row_range_start: i64, - pub row_range_end: i64, - snapshot_id: i64, - source_bucket: i32, - total_buckets: i32, - bucket_path: String, -} - -fn validate_table_options(table: &Table, core_options: &CoreOptions) -> Result<()> { - if !core_options.row_tracking_enabled() { - return Err(Error::DataInvalid { - message: "Sorted global index build requires 'row-tracking.enabled' = 'true'" - .to_string(), - source: None, - }); - } - if !core_options.data_evolution_enabled() { - return Err(Error::DataInvalid { - message: "Sorted global index build requires 'data-evolution.enabled' = 'true'" - .to_string(), - source: None, - }); - } - if !core_options.global_index_enabled() { - return Err(Error::DataInvalid { - message: "Sorted global index build requires 'global-index.enabled' = 'true'" - .to_string(), - source: None, - }); - } - if !table.schema().primary_keys().is_empty() { - return Err(Error::Unsupported { - message: "Sorted global index build does not support primary-key tables".to_string(), - }); - } - if core_options.deletion_vectors_enabled() { - return Err(Error::Unsupported { - message: - "Sorted global index build does not support tables with deletion-vectors.enabled=true" - .to_string(), - }); - } - Ok(()) -} - -fn find_index_field<'a>(table: &'a Table, column: &str) -> Result<&'a DataField> { - table - .schema() - .fields() - .iter() - .find(|field| field.name() == column) - .ok_or_else(|| Error::ColumnNotExist { - full_name: table.identifier().full_name(), - column: column.to_string(), - }) -} - -fn validate_btree_field(field: &DataField) -> Result<()> { - if !is_btree_supported_data_type(field.data_type()) { - return Err(Error::Unsupported { - message: format!( - "Sorted global index only supports scalar columns, got {:?} for column '{}'", - field.data_type(), - field.name() - ), - }); - } - Ok(()) -} - -fn index_key_type<'a>(index_type: &str, field: &'a DataField) -> Result<&'a DataType> { - if index_type == MULTIVALUE_GLOBAL_INDEX_TYPE { - let DataType::Array(array_type) = field.data_type() else { - return Err(Error::Unsupported { - message: format!( - "Multivalue global index requires an ARRAY column, got {:?} for column '{}'", - field.data_type(), - field.name() - ), - }); - }; - if !is_btree_supported_data_type(array_type.element_type()) { - return Err(Error::Unsupported { - message: format!( - "Multivalue global index does not support array element type {:?} for column '{}'", - array_type.element_type(), - field.name() - ), - }); - } - Ok(array_type.element_type()) - } else if index_type == FM_GLOBAL_INDEX_TYPE { - if !matches!(field.data_type(), DataType::Char(_) | DataType::VarChar(_)) { - return Err(Error::Unsupported { - message: format!( - "FM global index requires a character string column, got {:?} for column '{}'", - field.data_type(), - field.name() - ), - }); - } - Ok(field.data_type()) - } else { - validate_btree_field(field)?; - Ok(field.data_type()) - } -} - -fn is_btree_supported_data_type(data_type: &DataType) -> bool { - matches!( - data_type, - DataType::Boolean(_) - | DataType::TinyInt(_) - | DataType::SmallInt(_) - | DataType::Int(_) - | DataType::BigInt(_) - | DataType::Decimal(_) - | DataType::Double(_) - | DataType::Float(_) - | DataType::Char(_) - | DataType::VarChar(_) - | DataType::Date(_) - | DataType::LocalZonedTimestamp(_) - | DataType::Time(_) - | DataType::Timestamp(_) - ) -} - -#[allow(clippy::too_many_arguments)] -fn plan_sorted_index_shards( - table_location: &str, - partition_keys: &[String], - schema_fields: &[DataField], - core_options: &CoreOptions, - snapshot_id: i64, - entries: Vec, - records_per_range: i64, - indexed: &[RowRange], -) -> Result> { - if records_per_range <= 0 { - return Err(Error::DataInvalid { - message: format!( - "Option 'sorted-index.records-per-range' must be greater than 0, got: {records_per_range}" - ), - source: None, - }); - } - - let mut by_partition_bucket: HashMap<(Vec, i32, i32), Vec> = HashMap::new(); - for entry in entries { - if *entry.kind() != FileKind::Add { - continue; - } - if entry.file().first_row_id.is_none() { - return Err(Error::DataInvalid { - message: format!( - "Data file '{}' is missing first_row_id; cannot build a complete sorted global index", - entry.file().file_name - ), - source: None, - }); - } - let (partition, bucket, total_buckets, file) = entry.into_parts(); - by_partition_bucket - .entry((partition, bucket, total_buckets)) - .or_default() - .push(file); - } - - let mut result = Vec::new(); - for ((partition_bytes, source_bucket, total_buckets), files) in by_partition_bucket { - let partition = if partition_keys.is_empty() { - BinaryRow::new(0) - } else { - BinaryRow::from_serialized_bytes(&partition_bytes)? - }; - let bucket_path = bucket_path( - table_location, - partition_keys, - schema_fields, - core_options, - &partition, - source_bucket, - )?; - let normal_groups = group_normal_file_ranges(files)?; - for group in normal_groups { - let (coverage_start, coverage_end) = normal_coverage_range(&group.files)?; - let build_segments = - exclude_row_ranges(&[RowRange::new(coverage_start, coverage_end)], indexed); - for seg in build_segments { - let seg_start = seg.from(); - let seg_end = seg.to(); - let start_range = seg_start / records_per_range; - let end_range = seg_end / records_per_range; - for range_id in start_range..=end_range { - let range_start = range_id * records_per_range; - let range_end = range_start + records_per_range - 1; - let row_range_start = seg_start.max(range_start); - let row_range_end = seg_end.min(range_end); - result.push(SortedGlobalIndexShard { - partition: partition.clone(), - partition_bytes: partition_bytes.clone(), - files: group.files.clone(), - row_range_start, - row_range_end, - snapshot_id, - source_bucket, - total_buckets, - bucket_path: bucket_path.clone(), - }); - } - } - } - } - result.sort_by(|a, b| { - a.partition - .to_serialized_bytes() - .cmp(&b.partition.to_serialized_bytes()) - .then(a.source_bucket.cmp(&b.source_bucket)) - .then(a.row_range_start.cmp(&b.row_range_start)) - }); - Ok(result) -} - -#[derive(Debug)] -struct PlannedFileGroup { - files: Vec, -} - -fn group_normal_file_ranges(files: Vec) -> Result> { - if files.is_empty() { - return Ok(Vec::new()); - } - for file in &files { - file.row_id_range().ok_or_else(|| Error::DataInvalid { - message: format!( - "Data file '{}' is missing first_row_id; cannot build a complete sorted global index", - file.file_name - ), - source: None, - })?; - } - - let mut normal_ranges = files - .iter() - .filter(|file| is_data_evolution_normal_file(file)) - .filter_map(DataFileMeta::row_id_range) - .collect::>(); - normal_ranges.sort_by_key(|(start, _)| *start); - - let mut coverage_ranges: Vec<(i64, i64)> = Vec::new(); - for (file_start, file_end) in normal_ranges { - match coverage_ranges.last_mut() { - Some((_, end)) if file_start <= *end + 1 => { - *end = (*end).max(file_end); - } - _ => coverage_ranges.push((file_start, file_end)), - } - } - - coverage_ranges - .into_iter() - .map(|(start, end)| { - let mut group_files = files - .iter() - .filter(|file| { - file.row_id_range().is_some_and(|(file_start, file_end)| { - ranges_overlap(start, end, file_start, file_end) - }) - }) - .cloned() - .collect::>(); - group_files.sort_by_key(|file| { - ( - file.first_row_id.unwrap_or(i64::MAX), - !is_data_evolution_normal_file(file), - file.file_name.clone(), - ) - }); - Ok(PlannedFileGroup { files: group_files }) - }) - .collect() -} - -fn normal_coverage_range(files: &[DataFileMeta]) -> Result<(i64, i64)> { - let mut start = None; - let mut end = None; - for file in files - .iter() - .filter(|file| is_data_evolution_normal_file(file)) - { - let (file_start, file_end) = file.row_id_range().ok_or_else(|| Error::DataInvalid { - message: format!( - "Data file '{}' is missing first_row_id; cannot build a complete sorted global index", - file.file_name - ), - source: None, - })?; - start = Some(start.map_or(file_start, |value: i64| value.min(file_start))); - end = Some(end.map_or(file_end, |value: i64| value.max(file_end))); - } - start.zip(end).ok_or_else(|| Error::DataInvalid { - message: "Sorted global index shard has no normal data files".to_string(), - source: None, - }) -} - -fn bucket_path( - table_location: &str, - partition_keys: &[String], - schema_fields: &[DataField], - core_options: &CoreOptions, - partition: &BinaryRow, - bucket: i32, -) -> Result { - let base = table_location.trim_end_matches('/'); - if partition_keys.is_empty() { - return Ok(format!("{base}/{}", bucket_dir_name(bucket))); - } - let computer = crate::spec::PartitionComputer::new( - partition_keys, - schema_fields, - core_options.partition_default_name(), - core_options.legacy_partition_name(), - )?; - Ok(format!( - "{base}/{}{}", - computer.generate_partition_path(partition)?, - bucket_dir_name(bucket) - )) -} - -async fn extract_index_rows( - table: &Table, - shard: &SortedGlobalIndexShard, - index_column: &str, - index_field: &DataField, - index_type: &str, - serialize_key: SerializeKeyFn, -) -> Result> { - let splits = build_read_splits_for_shard(shard)?; - - let mut read_builder = table.new_read_builder(); - read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; - let read = read_builder.new_read()?; - let batches = read.to_arrow(&splits)?.try_collect::>().await?; - let expected_row_count = checked_row_count(shard.row_range_start, shard.row_range_end)?; - if index_type == MULTIVALUE_GLOBAL_INDEX_TYPE { - let DataType::Array(array_type) = index_field.data_type() else { - unreachable!("multivalue field was validated before extraction") - }; - extract_multivalue_index_rows_from_batches( - &batches, - index_column, - array_type.element_type(), - shard.row_range_start, - expected_row_count, - serialize_key, - ) - } else { - extract_index_rows_from_batches( - &batches, - index_column, - index_field.data_type(), - shard.row_range_start, - expected_row_count, - serialize_key, - ) - } -} - -fn build_read_splits_for_shard(shard: &SortedGlobalIndexShard) -> Result> { - let shard_range = RowRange::new(shard.row_range_start, shard.row_range_end); - group_by_overlapping_row_id(shard.files.clone()) - .into_iter() - .filter_map(|files| { - let ranges = files - .iter() - .filter_map(|file| { - file.row_id_range() - .and_then(|(start, end)| shard_range.intersect_inclusive(start, end)) - }) - .collect::>(); - let ranges = crate::table::merge_row_ranges(ranges); - if ranges.is_empty() { - return None; - } - let raw_convertible = files.len() == 1; - Some( - DataSplitBuilder::new() - .with_snapshot(shard.snapshot_id) - .with_partition(shard.partition.clone()) - .with_bucket(shard.source_bucket) - .with_bucket_path(shard.bucket_path.clone()) - .with_total_buckets(shard.total_buckets) - .with_data_files(files) - .with_row_ranges(ranges) - .with_raw_convertible(raw_convertible) - .build(), - ) - }) - .collect() -} - -fn extract_index_rows_from_batches( - batches: &[RecordBatch], - index_column: &str, - data_type: &DataType, - row_range_start: i64, - expected_row_count: i64, - serialize_key: SerializeKeyFn, -) -> Result> { - let row_count = batches.iter().map(RecordBatch::num_rows).sum::(); - let mut rows = Vec::with_capacity(row_count); - let mut expected_row_id = row_range_start; - for batch in batches { - let value_index = - batch - .schema() - .index_of(index_column) - .map_err(|e| Error::DataInvalid { - message: format!("Index column '{index_column}' not found in read batch: {e}"), - source: None, - })?; - let row_id_index = - batch - .schema() - .index_of(ROW_ID_FIELD_NAME) - .map_err(|e| Error::DataInvalid { - message: format!("_ROW_ID column not found in read batch: {e}"), - source: None, - })?; - let row_ids = batch - .column(row_id_index) - .as_any() - .downcast_ref::() - .ok_or_else(|| Error::DataInvalid { - message: "Sorted global index build requires non-null Int64 _ROW_ID".to_string(), - source: None, - })?; - - for row in 0..batch.num_rows() { - if row_ids.is_null(row) { - return Err(Error::DataInvalid { - message: "Sorted global index build found null _ROW_ID".to_string(), - source: None, - }); - } - let row_id = row_ids.value(row); - if row_id != expected_row_id { - return Err(Error::DataInvalid { - message: format!( - "Sorted global index build expected _ROW_ID {}, got {}", - expected_row_id, row_id - ), - source: None, - }); - } - expected_row_id += 1; - - let key = extract_datum_from_arrow(batch, row, value_index, data_type)? - .map(|datum| serialize_key(&datum, data_type)); - rows.push((key, row_id - row_range_start)); - } - } - let actual_row_count = expected_row_id - row_range_start; - if actual_row_count != expected_row_count { - return Err(Error::DataInvalid { - message: format!( - "Sorted global index build expected {} rows, got {}", - expected_row_count, actual_row_count - ), - source: None, - }); - } - Ok(rows) -} - -fn extract_multivalue_index_rows_from_batches( - batches: &[RecordBatch], - index_column: &str, - element_type: &DataType, - row_range_start: i64, - expected_row_count: i64, - serialize_key: SerializeKeyFn, -) -> Result> { - let mut rows = Vec::new(); - let mut expected_row_id = row_range_start; - for batch in batches { - let value_index = - batch - .schema() - .index_of(index_column) - .map_err(|e| Error::DataInvalid { - message: format!("Index column '{index_column}' not found in read batch: {e}"), - source: None, - })?; - let row_id_index = - batch - .schema() - .index_of(ROW_ID_FIELD_NAME) - .map_err(|e| Error::DataInvalid { - message: format!("_ROW_ID column not found in read batch: {e}"), - source: None, - })?; - let row_ids = batch - .column(row_id_index) - .as_any() - .downcast_ref::() - .ok_or_else(|| Error::DataInvalid { - message: "Multivalue global index build requires non-null Int64 _ROW_ID" - .to_string(), - source: None, - })?; - - #[derive(Clone, Copy)] - enum ArrayLayout<'a> { - List(&'a ListArray), - LargeList(&'a LargeListArray), - Fixed(&'a FixedSizeListArray), - } - let column = batch.column(value_index); - let layout = if let Some(array) = column.as_any().downcast_ref::() { - ArrayLayout::List(array) - } else if let Some(array) = column.as_any().downcast_ref::() { - ArrayLayout::LargeList(array) - } else if let Some(array) = column.as_any().downcast_ref::() { - ArrayLayout::Fixed(array) - } else { - return Err(Error::DataInvalid { - message: format!( - "Multivalue global index extraction requires an Arrow list column, got {:?}", - column.data_type() - ), - source: None, - }); - }; - let values = match layout { - ArrayLayout::List(array) => array.values(), - ArrayLayout::LargeList(array) => array.values(), - ArrayLayout::Fixed(array) => array.values(), - }; - - for row in 0..batch.num_rows() { - if row_ids.is_null(row) { - return Err(Error::DataInvalid { - message: "Multivalue global index build found null _ROW_ID".to_string(), - source: None, - }); - } - let row_id = row_ids.value(row); - if row_id != expected_row_id { - return Err(Error::DataInvalid { - message: format!( - "Multivalue global index build expected _ROW_ID {}, got {}", - expected_row_id, row_id - ), - source: None, - }); - } - expected_row_id += 1; - - let is_null = match layout { - ArrayLayout::List(array) => array.is_null(row), - ArrayLayout::LargeList(array) => array.is_null(row), - ArrayLayout::Fixed(array) => array.is_null(row), - }; - if is_null { - continue; - } - let (start, end) = match layout { - ArrayLayout::List(array) => { - let offsets = array.value_offsets(); - ( - usize::try_from(offsets[row]), - usize::try_from(offsets[row + 1]), - ) - } - ArrayLayout::LargeList(array) => { - let offsets = array.value_offsets(); - ( - usize::try_from(offsets[row]), - usize::try_from(offsets[row + 1]), - ) - } - ArrayLayout::Fixed(array) => { - let start = usize::try_from(array.value_offset(row)); - let end = usize::try_from(array.value_offset(row) + array.value_length()); - (start, end) - } - }; - let (start, end) = match (start, end) { - (Ok(start), Ok(end)) => (start, end), - _ => { - return Err(Error::DataInvalid { - message: "Multivalue global index found a negative array offset" - .to_string(), - source: None, - }) - } - }; - for element_index in start..end { - if let Some(datum) = - extract_datum_from_array(values, element_index, value_index, element_type)? - { - rows.push(( - Some(serialize_key(&datum, element_type)), - row_id - row_range_start, - )); - } - } - } - } - - let actual_row_count = expected_row_id - row_range_start; - if actual_row_count != expected_row_count { - return Err(Error::DataInvalid { - message: format!( - "Multivalue global index build expected {} source rows, got {}", - expected_row_count, actual_row_count - ), - source: None, - }); - } - Ok(rows) -} - -fn sort_index_rows(rows: &mut [SortedIndexKeyRow], cmp: &dyn Fn(&[u8], &[u8]) -> Ordering) { - rows.sort_by(|left, right| match (&left.0, &right.0) { - (None, None) => left.1.cmp(&right.1), - (None, Some(_)) => Ordering::Less, - (Some(_), None) => Ordering::Greater, - (Some(left_key), Some(right_key)) => { - cmp(left_key, right_key).then_with(|| left.1.cmp(&right.1)) - } - }); -} - -fn checked_i64(value: u64, context: &str) -> Result { - i64::try_from(value).map_err(|_| Error::DataInvalid { - message: format!("{context}: {value}"), - source: None, - }) -} - -fn checked_row_count(row_range_start: i64, row_range_end: i64) -> Result { - if row_range_end < row_range_start { - return Err(Error::DataInvalid { - message: format!( - "Invalid sorted global index row range [{row_range_start}, {row_range_end}]" - ), - source: None, - }); - } - row_range_end - .checked_sub(row_range_start) - .and_then(|span| span.checked_add(1)) - .ok_or_else(|| Error::DataInvalid { - message: format!( - "Row count overflows for row range [{row_range_start}, {row_range_end}]" - ), - source: None, - }) -} - -fn ranges_overlap(left_start: i64, left_end: i64, right_start: i64, right_end: i64) -> bool { - left_start <= right_end && right_start <= left_end } #[cfg(test)] -mod tests { - - /// A row range wider than `i32::MAX` yields the full count instead of being - /// rejected, and an inverted or overflowing range is still an error. - #[test] - fn checked_row_count_spans_beyond_i32() { - let start = 0; - let end = i64::from(i32::MAX) + 10; - assert_eq!( - super::checked_row_count(start, end).unwrap(), - i64::from(i32::MAX) + 11 - ); - assert!(super::checked_row_count(5, 4).is_err()); - assert!(super::checked_row_count(i64::MIN, i64::MAX).is_err()); - } - use super::*; - use crate::btree::BTreeIndexMeta; - use crate::catalog::Identifier; - use crate::io::FileIOBuilder; - use crate::spec::stats::BinaryTableStats; - use crate::spec::{ - ArrayType, BinaryRowBuilder, BinaryType, DoubleType, FloatType, GlobalIndexSearchMode, - IndexManifest, IntType, ManifestEntry, Predicate, PredicateBuilder, PredicateOperator, - Schema, TableSchema, TimeType, VarBinaryType, VarCharType, - }; - use crate::table::global_index_scanner::{evaluate_global_index, GlobalIndexEvaluation}; - use crate::table::{merge_row_ranges, SnapshotManager, TableCommit, TableWrite}; - use arrow_array::builder::{Int32Builder, ListBuilder, Time32MillisecondBuilder}; - use arrow_array::{ArrayRef, Float32Array, Float64Array, Int32Array, Int64Array, StringArray}; - use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; - use chrono::{DateTime, Utc}; - use std::sync::Arc; - - fn data_file(name: &str, first_row_id: Option, row_count: i64) -> DataFileMeta { - DataFileMeta { - file_name: name.to_string(), - file_size: 128, - row_count, - min_key: vec![], - max_key: vec![], - key_stats: BinaryTableStats::new(vec![], vec![], vec![]), - value_stats: BinaryTableStats::new(vec![], vec![], vec![]), - min_sequence_number: 0, - max_sequence_number: 0, - schema_id: 0, - level: 0, - extra_files: vec![], - creation_time: Some( - "2024-09-06T07:45:55.039+00:00" - .parse::>() - .unwrap(), - ), - delete_row_count: None, - embedded_index: None, - first_row_id, - write_cols: None, - external_path: None, - file_source: None, - value_stats_cols: None, - column_max_sequence_numbers: None, - } - } - - fn partial_file(name: &str, first_row_id: Option, row_count: i64) -> DataFileMeta { - let mut file = data_file(name, first_row_id, row_count); - file.write_cols = Some(vec!["name".to_string()]); - file - } - - fn manifest_entry(file: DataFileMeta) -> ManifestEntry { - manifest_entry_with_bucket(file, 0, 1) - } - - fn manifest_entry_with_bucket( - file: DataFileMeta, - bucket: i32, - total_buckets: i32, - ) -> ManifestEntry { - ManifestEntry::new(FileKind::Add, vec![], bucket, total_buckets, file, 2) - } - - fn table_options(records_per_range: &str) -> HashMap { - HashMap::from([ - ("row-tracking.enabled".to_string(), "true".to_string()), - ("data-evolution.enabled".to_string(), "true".to_string()), - ("global-index.enabled".to_string(), "true".to_string()), - ( - "sorted-index.records-per-range".to_string(), - records_per_range.to_string(), - ), - ]) - } - - fn test_table(options: HashMap) -> Table { - test_table_with_path("memory:/test_btree_global_index_builder", options) - } - - fn test_table_with_path(table_path: &str, options: HashMap) -> Table { - let schema = Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column("name", DataType::VarChar(VarCharType::string_type())) - .options(options) - .build() - .unwrap(); - Table::new( - FileIOBuilder::new("memory").build().unwrap(), - Identifier::new("default", "test_table"), - table_path.to_string(), - TableSchema::new(0, &schema), - None, - ) - } - - fn multivalue_table(table_path: &str) -> Table { - let schema = Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column( - "items", - DataType::Array(ArrayType::new(DataType::Int(IntType::new()))), - ) - .options(table_options("10")) - .build() - .unwrap(); - Table::new( - FileIOBuilder::new("memory").build().unwrap(), - Identifier::new("default", "test_multivalue_table"), - table_path.to_string(), - TableSchema::new(0, &schema), - None, - ) - } - - fn plan( - entries: Vec, - records_per_range: i64, - ) -> Result> { - let table = test_table(table_options(&records_per_range.to_string())); - let core = CoreOptions::new(table.schema().options()); - plan_sorted_index_shards( - table.location(), - table.schema().partition_keys(), - table.schema().fields(), - &core, - 1, - entries, - records_per_range, - &[], - ) - } - - #[test] - fn test_planner_splits_single_file_across_ranges() { - let shards = plan(vec![manifest_entry(data_file("a", Some(0), 25))], 10).unwrap(); - - assert_eq!( - shards - .iter() - .map(|s| (s.row_range_start, s.row_range_end)) - .collect::>(), - vec![(0, 9), (10, 19), (20, 24)] - ); - } - - #[test] - fn test_planner_merges_contiguous_normal_files() { - let shards = plan( - vec![ - manifest_entry(data_file("a", Some(0), 5)), - manifest_entry(data_file("b", Some(5), 5)), - ], - 20, - ) - .unwrap(); - - assert_eq!(shards.len(), 1); - assert_eq!((shards[0].row_range_start, shards[0].row_range_end), (0, 9)); - } - - #[test] - fn test_planner_splits_row_id_gap_into_separate_shards() { - let shards = plan( - vec![ - manifest_entry(data_file("a", Some(0), 5)), - manifest_entry(data_file("b", Some(10), 5)), - ], - 20, - ) - .unwrap(); - - assert_eq!( - shards - .iter() - .map(|s| (s.row_range_start, s.row_range_end)) - .collect::>(), - vec![(0, 4), (10, 14)] - ); - } - - #[test] - fn test_planner_rejects_missing_first_row_id() { - let err = plan(vec![manifest_entry(data_file("a", None, 5))], 10) - .expect_err("missing first_row_id should fail"); - assert!( - matches!(err, Error::DataInvalid { message, .. } if message.contains("missing first_row_id")) - ); - } - - #[test] - fn test_planner_keeps_buckets_separate() { - let shards = plan( - vec![ - manifest_entry_with_bucket(data_file("a", Some(0), 5), 0, 2), - manifest_entry_with_bucket(data_file("b", Some(5), 5), 1, 2), - ], - 20, - ) - .unwrap(); - - assert_eq!( - shards - .iter() - .map(|s| ( - s.source_bucket, - s.total_buckets, - s.row_range_start, - s.row_range_end - )) - .collect::>(), - vec![(0, 2, 0, 4), (1, 2, 5, 9)] - ); - } - - #[test] - fn test_planner_keeps_partial_file_in_read_group_without_expanding_coverage() { - let shards = plan( - vec![ - manifest_entry(data_file("base", Some(0), 5)), - manifest_entry(partial_file("partial", Some(0), 5)), - ], - 20, - ) - .unwrap(); - - assert_eq!(shards.len(), 1); - assert_eq!((shards[0].row_range_start, shards[0].row_range_end), (0, 4)); - assert_eq!(shards[0].files.len(), 2); - } - - #[test] - fn test_build_read_splits_groups_only_overlapping_partial_files() { - let shards = plan( - vec![ - manifest_entry(data_file("a", Some(0), 5)), - manifest_entry(data_file("b", Some(5), 5)), - manifest_entry(partial_file("partial", Some(0), 5)), - ], - 20, - ) - .unwrap(); - assert_eq!(shards.len(), 1); - - let splits = build_read_splits_for_shard(&shards[0]).unwrap(); - - assert_eq!(splits.len(), 2); - assert_eq!( - splits[0] - .data_files() - .iter() - .map(|file| file.file_name.as_str()) - .collect::>(), - vec!["a", "partial"] - ); - assert_eq!(splits[0].row_ranges(), Some(&[RowRange::new(0, 4)][..])); - assert!(!splits[0].raw_convertible()); - - assert_eq!( - splits[1] - .data_files() - .iter() - .map(|file| file.file_name.as_str()) - .collect::>(), - vec!["b"] - ); - assert_eq!(splits[1].row_ranges(), Some(&[RowRange::new(5, 9)][..])); - assert!(splits[1].raw_convertible()); - } - - #[test] - fn test_validate_btree_field_rejects_complex_type() { - let field = DataField::new( - 0, - "items".to_string(), - DataType::Array(crate::spec::ArrayType::new(DataType::Int(IntType::new()))), - ); - let err = validate_btree_field(&field).expect_err("array should be rejected"); - assert!(matches!(err, Error::Unsupported { message } if message.contains("scalar"))); - } - - #[test] - fn test_validate_btree_field_rejects_binary_types() { - for data_type in [ - DataType::Binary(BinaryType::new(4).unwrap()), - DataType::VarBinary(VarBinaryType::try_new(true, 4).unwrap()), - ] { - let field = DataField::new(0, "bytes".to_string(), data_type); - let err = validate_btree_field(&field).expect_err("binary should be rejected"); - assert!(matches!(err, Error::Unsupported { message } if message.contains("scalar"))); - } - } - - fn index_batch(values: Vec>, row_ids: Vec>) -> RecordBatch { - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", ArrowDataType::Int32, true), - ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), - ])); - RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from(values)) as ArrayRef, - Arc::new(Int64Array::from(row_ids)) as ArrayRef, - ], - ) - .unwrap() - } - - #[test] - fn test_extract_index_rows_serializes_keys_and_local_row_ids() { - let batch = index_batch( - vec![Some(10), None, Some(30)], - vec![Some(5), Some(6), Some(7)], - ); - let rows = extract_index_rows_from_batches( - &[batch], - "id", - &DataType::Int(IntType::new()), - 5, - 3, - serialize_datum, - ) - .unwrap(); - - assert_eq!( - rows, - vec![ - (Some(10i32.to_le_bytes().to_vec()), 0), - (None, 1), - (Some(30i32.to_le_bytes().to_vec()), 2), - ] - ); - } - - #[test] - fn test_extract_multivalue_rows_skips_null_arrays_and_elements() { - let element = Arc::new(ArrowField::new("element", ArrowDataType::Int32, true)); - let mut items = ListBuilder::new(Int32Builder::new()).with_field(element.clone()); - items.values().append_value(10); - items.values().append_null(); - items.values().append_value(10); - items.append(true); - items.append(true); // empty array - items.append(false); // null array - items.values().append_value(30); - items.append(true); - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("items", ArrowDataType::List(element), true), - ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, false), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(items.finish()) as ArrayRef, - Arc::new(Int64Array::from_iter_values(5..9)) as ArrayRef, - ], - ) - .unwrap(); - - let rows = extract_multivalue_index_rows_from_batches( - &[batch], - "items", - &DataType::Int(IntType::new()), - 5, - 4, - serialize_bitmap_datum, - ) - .unwrap(); - - assert_eq!( - rows, - vec![ - (Some(10i32.to_le_bytes().to_vec()), 0), - (Some(10i32.to_le_bytes().to_vec()), 0), - (Some(30i32.to_le_bytes().to_vec()), 3), - ] - ); - } - - #[test] - fn test_extract_multivalue_time_rows_matches_java_int_serializer() { - let element = Arc::new(ArrowField::new( - "element", - ArrowDataType::Time32(arrow_schema::TimeUnit::Millisecond), - true, - )); - let mut items = - ListBuilder::new(Time32MillisecondBuilder::new()).with_field(element.clone()); - items.values().append_value(12_345); - items.values().append_null(); - items.append(true); - items.values().append_value(86_399_999); - items.append(true); - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("items", ArrowDataType::List(element), true), - ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, false), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(items.finish()) as ArrayRef, - Arc::new(Int64Array::from_iter_values(7..9)) as ArrayRef, - ], - ) - .unwrap(); - let time_type = DataType::Time(TimeType::new(3).unwrap()); - - let rows = extract_multivalue_index_rows_from_batches( - &[batch], - "items", - &time_type, - 7, - 2, - serialize_bitmap_datum, - ) - .unwrap(); - - assert_eq!( - rows, - vec![ - (Some(12_345i32.to_le_bytes().to_vec()), 0), - (Some(86_399_999i32.to_le_bytes().to_vec()), 1), - ] - ); - } - - #[test] - fn test_index_key_codec_scopes_java_nan_semantics_to_bitmap() { - fn assert_codec( - data_type: DataType, - negative_nan: Datum, - raw_nan_key: Vec, - canonical_nan_key: Vec, - zero: Datum, - ) { - let (btree_cmp, btree_serialize) = - make_index_key_codec(BTREE_GLOBAL_INDEX_TYPE, &data_type); - let btree_nan_key = btree_serialize(&negative_nan, &data_type); - let zero_key = btree_serialize(&zero, &data_type); - assert_eq!(btree_nan_key, raw_nan_key); - assert!(btree_cmp(&btree_nan_key, &zero_key).is_lt()); - - let (bitmap_cmp, bitmap_serialize) = - make_index_key_codec(BITMAP_GLOBAL_INDEX_TYPE, &data_type); - let bitmap_nan_key = bitmap_serialize(&negative_nan, &data_type); - assert_eq!(bitmap_nan_key, canonical_nan_key); - assert!(bitmap_cmp(&bitmap_nan_key, &zero_key).is_gt()); - } - - assert_codec( - DataType::Float(FloatType::new()), - Datum::Float(f32::from_bits(0xffc0_0001)), - 0xffc0_0001u32.to_le_bytes().to_vec(), - 0x7fc0_0000u32.to_le_bytes().to_vec(), - Datum::Float(0.0), - ); - assert_codec( - DataType::Double(DoubleType::new()), - Datum::Double(f64::from_bits(0xfff8_0000_0000_0001)), - 0xfff8_0000_0000_0001u64.to_le_bytes().to_vec(), - 0x7ff8_0000_0000_0000u64.to_le_bytes().to_vec(), - Datum::Double(0.0), - ); - } - - #[test] - fn test_extract_index_rows_rejects_row_id_gap() { - let batch = index_batch(vec![Some(10), Some(30)], vec![Some(5), Some(7)]); - let err = extract_index_rows_from_batches( - &[batch], - "id", - &DataType::Int(IntType::new()), - 5, - 2, - serialize_datum, - ) - .expect_err("row-id gap should fail"); - - assert!( - matches!(err, Error::DataInvalid { message, .. } if message.contains("expected _ROW_ID")) - ); - } - - #[test] - fn test_sort_index_rows_orders_nulls_then_keys() { - let mut rows = vec![ - (Some(3i32.to_le_bytes().to_vec()), 0), - (None, 1), - (Some(1i32.to_le_bytes().to_vec()), 2), - (Some(1i32.to_le_bytes().to_vec()), 3), - ]; - let cmp = make_key_comparator(&DataType::Int(IntType::new())); - - sort_index_rows(&mut rows, &cmp); - - assert_eq!( - rows, - vec![ - (None, 1), - (Some(1i32.to_le_bytes().to_vec()), 2), - (Some(1i32.to_le_bytes().to_vec()), 3), - (Some(3i32.to_le_bytes().to_vec()), 0), - ] - ); - } - - #[test] - fn test_extract_index_rows_accepts_string_column() { - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("name", ArrowDataType::Utf8, true), - ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(StringArray::from(vec![Some("alice"), None])) as ArrayRef, - Arc::new(Int64Array::from(vec![Some(10), Some(11)])) as ArrayRef, - ], - ) - .unwrap(); - - let rows = extract_index_rows_from_batches( - &[batch], - "name", - &DataType::VarChar(VarCharType::string_type()), - 10, - 2, - serialize_datum, - ) - .unwrap(); - - assert_eq!(rows, vec![(Some(b"alice".to_vec()), 0), (None, 1)]); - } - - fn data_batch(ids: Vec, names: Vec<&str>) -> RecordBatch { - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", ArrowDataType::Int32, false), - ArrowField::new("name", ArrowDataType::Utf8, true), - ])); - RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from(ids)) as ArrayRef, - Arc::new(StringArray::from(names)) as ArrayRef, - ], - ) - .unwrap() - } - - fn nullable_name_batch(ids: Vec, names: Vec>) -> RecordBatch { - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", ArrowDataType::Int32, false), - ArrowField::new("name", ArrowDataType::Utf8, true), - ])); - RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from(ids)) as ArrayRef, - Arc::new(StringArray::from(names)) as ArrayRef, - ], - ) - .unwrap() - } - - fn multivalue_batch() -> RecordBatch { - let element = Arc::new(ArrowField::new("element", ArrowDataType::Int32, true)); - let mut items = ListBuilder::new(Int32Builder::new()).with_field(element.clone()); - - items.values().append_value(10); - items.values().append_null(); - items.values().append_value(10); - items.append(true); - - items.append(true); // empty array - items.append(false); // null array - - items.values().append_value(10); - items.values().append_value(30); - items.append(true); - - items.values().append_value(30); - items.values().append_value(40); - items.append(true); - - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", ArrowDataType::Int32, false), - ArrowField::new("items", ArrowDataType::List(element), true), - ])); - RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from_iter_values(1..=5)) as ArrayRef, - Arc::new(items.finish()) as ArrayRef, - ], - ) - .unwrap() - } - - async fn setup_dirs(table: &Table) { - table - .file_io() - .mkdirs(&format!("{}/snapshot/", table.location())) - .await - .unwrap(); - table - .file_io() - .mkdirs(&format!("{}/manifest/", table.location())) - .await - .unwrap(); - } - - async fn scan_ids(table: &Table, predicate: Predicate) -> Vec { - let mut builder = table.new_read_builder(); - builder.with_filter(predicate); - let plan = builder.new_scan().plan().await.unwrap(); - let read = builder.new_read().unwrap(); - let batches = read - .to_arrow(plan.splits()) - .unwrap() - .try_collect::>() - .await - .unwrap(); - let mut ids = batches - .iter() - .flat_map(|batch| { - batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .values() - .iter() - .copied() - }) - .collect::>(); - ids.sort_unstable(); - ids - } - - #[tokio::test] - async fn test_execute_writes_btree_index_manifest_and_file() { - let table_path = "memory:/test_btree_global_index_builder_e2e"; - let table = test_table_with_path(table_path, table_options("10")); - setup_dirs(&table).await; - - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "alice"])) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - let shard_count = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .execute() - .await - .unwrap(); - assert_eq!(shard_count, 1); - - let snapshot_manager = - SnapshotManager::new(table.file_io().clone(), table.location().to_string()); - let snapshot = snapshot_manager - .get_latest_snapshot() - .await - .unwrap() - .unwrap(); - let index_manifest = snapshot.index_manifest().expect("index manifest"); - let index_entries = IndexManifest::read( - table.file_io(), - &format!("{table_path}/manifest/{index_manifest}"), - ) - .await - .unwrap(); - assert_eq!(index_entries.len(), 1); - - let index_file = &index_entries[0].index_file; - assert_eq!(index_file.index_type, BTREE_GLOBAL_INDEX_TYPE); - assert!(index_file.file_name.starts_with("btree-global-index-")); - assert_eq!(index_file.row_count, 3); - assert!(index_file.file_size > 0); - - let global_meta = index_file - .global_index_meta - .as_ref() - .expect("global index meta"); - assert_eq!(global_meta.row_range_start, 0); - assert_eq!(global_meta.row_range_end, 2); - assert_eq!(global_meta.index_field_id, 1); - let btree_meta = - crate::btree::BTreeIndexMeta::deserialize(global_meta.index_meta.as_ref().unwrap()) - .unwrap(); - assert_eq!(btree_meta.first_key, Some(b"alice".to_vec())); - assert_eq!(btree_meta.last_key, Some(b"bob".to_vec())); - assert!(!btree_meta.has_nulls); - - let predicate = PredicateBuilder::new(table.schema().fields()) - .equal("name", crate::spec::Datum::String("alice".to_string())) - .unwrap(); - let row_ranges = evaluate_global_index(GlobalIndexEvaluation { - file_io: table.file_io(), - table_path: table.location(), - index_entries: &index_entries, - predicates: &[predicate], - schema_fields: table.schema().fields(), - search_mode: GlobalIndexSearchMode::Fast, - global_index_thread_num: 32, - btree_fallback_scan_max_size: i64::MAX, - bitmap_fallback_scan_max_size: i64::MAX, - fm_read_options: crate::fm_index::FMReadOptions::default(), - next_row_id: snapshot.next_row_id(), - data_ranges: &[], - }) - .await - .unwrap() - .unwrap(); - assert_eq!(row_ranges, vec![RowRange::new(0, 0), RowRange::new(2, 2)]); - - // Reopen the same table without an explicit global-index override and - // verify that the regular scan path still uses the committed index. - let mut options = table.schema().options().clone(); - assert_eq!( - options.remove("global-index.enabled"), - Some("true".to_string()) - ); - let scan_table = Table::new( - table.file_io().clone(), - table.identifier().clone(), - table.location().to_string(), - table.schema().copy_with_replaced_options(options), - None, - ); - let predicate = PredicateBuilder::new(scan_table.schema().fields()) - .equal("name", crate::spec::Datum::String("alice".to_string())) - .unwrap(); - let mut read_builder = scan_table.new_read_builder(); - read_builder.with_filter(predicate); - let plan = read_builder.new_scan().plan().await.unwrap(); - - assert_eq!(plan.splits().len(), 1); - assert_eq!( - plan.splits()[0].row_ranges(), - Some(&[RowRange::new(0, 0), RowRange::new(2, 2)][..]) - ); - } - - #[tokio::test] - async fn test_execute_writes_and_queries_fm_index() { - let table_path = "memory:/test_fm_global_index_builder_e2e"; - let mut options = table_options("10"); - options.insert("fm-index.sa-sample-rate".to_string(), "1".to_string()); - options.insert("fm-index.locate-cost-ratio".to_string(), "1".to_string()); - options.insert("fm-index.compression".to_string(), "none".to_string()); - let table = test_table_with_path(table_path, options); - setup_dirs(&table).await; - - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&nullable_name_batch( - vec![1, 2, 3, 4], - vec![Some("banana"), Some("bandana"), None, Some("")], - )) - .await - .unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(table_write.prepare_commit().await.unwrap()) - .await - .unwrap(); - - assert_eq!( - table - .new_sorted_global_index_build_builder() - .with_index_column("name") - .with_index_type(FM_GLOBAL_INDEX_TYPE) - .execute() - .await - .unwrap(), - 1 - ); - - let snapshot = table - .snapshot_manager() - .get_latest_snapshot() - .await - .unwrap() - .unwrap(); - let index_entries = IndexManifest::read( - table.file_io(), - &format!( - "{table_path}/manifest/{}", - snapshot.index_manifest().expect("index manifest") - ), - ) - .await - .unwrap(); - assert_eq!(index_entries.len(), 1); - let index_file = &index_entries[0].index_file; - assert_eq!(index_file.index_type, FM_GLOBAL_INDEX_TYPE); - assert!(index_file.file_name.starts_with("fm-global-index-")); - assert_eq!(index_file.row_count, 4); - crate::fm_index::validate_manifest_meta( - index_file - .global_index_meta - .as_ref() - .unwrap() - .index_meta - .as_ref() - .unwrap(), - ) - .unwrap(); - - let predicates = PredicateBuilder::new(table.schema().fields()); - assert_eq!( - scan_ids( - &table, - predicates - .contains("name", Datum::String("ana".to_string())) - .unwrap(), - ) - .await, - vec![1, 2] - ); - assert_eq!( - scan_ids(&table, predicates.is_null("name").unwrap()).await, - vec![3] - ); - assert_eq!( - scan_ids( - &table, - predicates - .contains("name", Datum::String(String::new())) - .unwrap(), - ) - .await, - vec![1, 2, 4] - ); - - // A dense match that exceeds the FM locate budget must decline the - // index and let the normal row filter scan the source, never produce - // a false empty result. - let mut fallback_options = table.schema().options().clone(); - fallback_options.insert( - "fm-index.locate-cost-ratio".to_string(), - "0.000001".to_string(), - ); - let fallback_table = Table::new( - table.file_io().clone(), - table.identifier().clone(), - table.location().to_string(), - table.schema().copy_with_replaced_options(fallback_options), - None, - ); - let predicate = PredicateBuilder::new(fallback_table.schema().fields()) - .contains("name", Datum::String("a".to_string())) - .unwrap(); - assert_eq!(scan_ids(&fallback_table, predicate).await, vec![1, 2]); - } - - #[tokio::test] - async fn test_failed_fm_build_removes_partial_index_file() { - let table_path = "memory:/test_failed_fm_build_cleanup"; - let table = test_table_with_path(table_path, table_options("3")); - setup_dirs(&table).await; - - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch( - vec![1, 2, 3, 4, 5, 6], - vec!["one", "two", "six", "red", "blue", "value-too-long"], - )) - .await - .unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(table_write.prepare_commit().await.unwrap()) - .await - .unwrap(); - - let mut options = HashMap::new(); - options.insert("fm-index.partition-size".to_string(), "8".to_string()); - options.insert("fm-index.partition-row-count".to_string(), "1".to_string()); - options.insert("fm-index.compression".to_string(), "none".to_string()); - let error = table - .new_sorted_global_index_build_builder() - .with_index_column("name") - .with_index_type(FM_GLOBAL_INDEX_TYPE) - .with_options(options) - .execute() - .await - .expect_err("the oversized FM value must fail the build"); - assert!(matches!(error, Error::DataInvalid { .. })); - - let files = table - .file_io() - .list_status(&format!("{table_path}/index")) - .await - .unwrap(); - assert!( - files - .iter() - .all(|file| !file.path.contains("fm-global-index-")), - "failed FM build left a partial index file: {files:?}" - ); - } - - #[tokio::test] - async fn test_global_index_prunes_during_manifest_read() { - for (search_mode, expected_manifest_pruned, expected_entries_read) in - [("fast", 1, 1), ("full", 1, 1)] - { - let table_path = format!("memory:/test_global_index_manifest_pruning_{search_mode}"); - let mut options = table_options("2"); - options.insert( - "global-index.search-mode".to_string(), - search_mode.to_string(), - ); - let table = test_table_with_path(&table_path, options); - setup_dirs(&table).await; - - for (user, ids, names) in [ - ("writer-1", vec![1, 2], vec!["alice", "bob"]), - ("writer-2", vec![3, 4], vec!["carol", "dave"]), - ] { - let mut table_write = TableWrite::new(&table, user.to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch(ids, names)) - .await - .unwrap(); - TableCommit::new(table.clone(), user.to_string()) - .commit(table_write.prepare_commit().await.unwrap()) - .await - .unwrap(); - } - - table - .new_btree_global_index_build_builder() - .with_index_column("name") - .execute() - .await - .unwrap(); - - let predicate = PredicateBuilder::new(table.schema().fields()) - .equal("name", crate::spec::Datum::String("alice".to_string())) - .unwrap(); - let mut read_builder = table.new_read_builder(); - read_builder.with_filter(predicate); - let (plan, trace) = read_builder.new_scan().plan_with_trace().await.unwrap(); - - assert_eq!( - plan.splits() - .iter() - .flat_map(|split| split.data_files()) - .count(), - 1 - ); - assert_eq!( - trace.manifest_files_pruned_by_row_ranges, - expected_manifest_pruned - ); - assert_eq!(trace.manifest_entries_read, expected_entries_read); - assert_eq!(trace.manifest_entries_pruned_by_row_ranges, 0); - assert_eq!( - trace.manifest_entries_after_manifest_filters, - expected_entries_read - ); - } - } - - #[tokio::test] - async fn test_scalar_full_search_includes_unindexed_rows() { - let mut options = table_options("2"); - options.insert("scalar-index.search-mode".to_string(), "full".to_string()); - let table = test_table_with_path( - "memory:/test_scalar_full_search_includes_unindexed_rows", - options, - ); - setup_dirs(&table).await; - - let mut first_write = TableWrite::new(&table, "writer-1".to_string()).unwrap(); - first_write - .write_arrow_batch(&data_batch(vec![1, 2], vec!["alice", "bob"])) - .await - .unwrap(); - TableCommit::new(table.clone(), "writer-1".to_string()) - .commit(first_write.prepare_commit().await.unwrap()) - .await - .unwrap(); - table - .new_btree_global_index_build_builder() - .with_index_column("name") - .execute() - .await - .unwrap(); - - let mut second_write = TableWrite::new(&table, "writer-2".to_string()).unwrap(); - second_write - .write_arrow_batch(&data_batch(vec![3, 4], vec!["alice", "dave"])) - .await - .unwrap(); - TableCommit::new(table.clone(), "writer-2".to_string()) - .commit(second_write.prepare_commit().await.unwrap()) - .await - .unwrap(); - - let predicate = PredicateBuilder::new(table.schema().fields()) - .equal("name", crate::spec::Datum::String("alice".to_string())) - .unwrap(); - let mut read_builder = table.new_read_builder(); - read_builder.with_filter(predicate); - let plan = read_builder.new_scan().plan().await.unwrap(); - let planned_ranges = merge_row_ranges( - plan.splits() - .iter() - .flat_map(|split| split.row_ranges().unwrap_or_default()) - .cloned() - .collect(), - ); - - assert_eq!( - planned_ranges, - vec![RowRange::new(0, 0), RowRange::new(2, 3)] - ); - assert_eq!( - scan_ids( - &table, - PredicateBuilder::new(table.schema().fields()) - .equal("name", crate::spec::Datum::String("alice".to_string())) - .unwrap(), - ) - .await, - vec![1, 3] - ); - } - - #[tokio::test] - async fn test_empty_global_index_ranges_skip_legacy_manifests() { - for search_mode in ["fast", "full"] { - let table_path = format!("memory:/test_empty_global_index_ranges_{search_mode}"); - let mut options = table_options("10"); - options.insert( - "global-index.search-mode".to_string(), - search_mode.to_string(), - ); - let table = test_table_with_path(&table_path, options); - setup_dirs(&table).await; - - let mut table_write = TableWrite::new(&table, "writer".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch(vec![1, 2], vec!["alice", "bob"])) - .await - .unwrap(); - TableCommit::new(table.clone(), "writer".to_string()) - .commit(table_write.prepare_commit().await.unwrap()) - .await - .unwrap(); - table - .new_btree_global_index_build_builder() - .with_index_column("name") - .execute() - .await - .unwrap(); - - let mut legacy_file = data_file("legacy.parquet", None, 2); - legacy_file.level = 1; - legacy_file.file_source = Some(1); // FileSource.COMPACT - TableCommit::new(table.clone(), "legacy-writer".to_string()) - .commit(vec![CommitMessage::new( - BinaryRowBuilder::new(0).build_serialized(), - 0, - vec![legacy_file], - )]) - .await - .unwrap(); - - let snapshot = table - .snapshot_manager() - .get_latest_snapshot() - .await - .unwrap() - .unwrap(); - assert_eq!(snapshot.next_row_id(), Some(2)); - - let predicate = PredicateBuilder::new(table.schema().fields()) - .equal( - "name", - crate::spec::Datum::String("not-present".to_string()), - ) - .unwrap(); - let mut read_builder = table.new_read_builder(); - read_builder.with_filter(predicate); - - let plan = read_builder.new_scan().plan().await.unwrap(); - assert!(plan.splits().is_empty()); - - let (traced_plan, trace) = read_builder.new_scan().plan_with_trace().await.unwrap(); - let delta_plan = read_builder - .new_scan() - .plan_snapshot_delta(&snapshot) - .await - .unwrap(); - - assert!(traced_plan.splits().is_empty()); - assert_eq!(trace.manifest_entries_read, 0); - assert_eq!(trace.final_splits, 0); - assert_eq!(trace.final_files, 0); - assert!(delta_plan.splits().is_empty()); - } - } - - #[tokio::test] - async fn test_detail_mode_defers_manifest_pruning_for_unindexed_ranges() { - let table_path = "memory:/test_detail_manifest_pruning"; - let mut options = table_options("2"); - options.insert("global-index.search-mode".to_string(), "detail".to_string()); - let table = test_table_with_path(table_path, options); - setup_dirs(&table).await; - - let mut table_write = TableWrite::new(&table, "writer-1".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch(vec![1, 2], vec!["alice", "bob"])) - .await - .unwrap(); - TableCommit::new(table.clone(), "writer-1".to_string()) - .commit(table_write.prepare_commit().await.unwrap()) - .await - .unwrap(); - table - .new_btree_global_index_build_builder() - .with_index_column("name") - .execute() - .await - .unwrap(); - - let mut table_write = TableWrite::new(&table, "writer-2".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch(vec![3, 4], vec!["alice", "dave"])) - .await - .unwrap(); - TableCommit::new(table.clone(), "writer-2".to_string()) - .commit(table_write.prepare_commit().await.unwrap()) - .await - .unwrap(); - - let predicate = PredicateBuilder::new(table.schema().fields()) - .equal("name", crate::spec::Datum::String("alice".to_string())) - .unwrap(); - let mut read_builder = table.new_read_builder(); - read_builder.with_filter(predicate); - let (plan, trace) = read_builder.new_scan().plan_with_trace().await.unwrap(); - let planned_ranges = merge_row_ranges( - plan.splits() - .iter() - .flat_map(|split| split.row_ranges().unwrap_or_default()) - .cloned() - .collect(), - ); - - assert_eq!( - plan.splits() - .iter() - .flat_map(|split| split.data_files()) - .count(), - 2 - ); - assert_eq!( - planned_ranges, - vec![RowRange::new(0, 0), RowRange::new(2, 3)] - ); - assert_eq!(trace.manifest_files_pruned_by_row_ranges, 0); - assert_eq!(trace.manifest_entries_read, 2); - } - - #[tokio::test] - async fn test_execute_writes_bitmap_index_manifest_and_java_file() { - let table_path = "memory:/test_bitmap_global_index_builder_e2e"; - let table = test_table_with_path(table_path, table_options("10")); - setup_dirs(&table).await; - - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "alice"])) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - let shard_count = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .with_index_type(BITMAP_GLOBAL_INDEX_TYPE) - .execute() - .await - .unwrap(); - assert_eq!(shard_count, 1); - - let snapshot_manager = - SnapshotManager::new(table.file_io().clone(), table.location().to_string()); - let snapshot = snapshot_manager - .get_latest_snapshot() - .await - .unwrap() - .unwrap(); - let index_manifest = snapshot.index_manifest().expect("index manifest"); - let index_entries = IndexManifest::read( - table.file_io(), - &format!("{table_path}/manifest/{index_manifest}"), - ) - .await - .unwrap(); - assert_eq!(index_entries.len(), 1); - - let index_file = &index_entries[0].index_file; - assert_eq!(index_file.index_type, BITMAP_GLOBAL_INDEX_TYPE); - assert!(index_file.file_name.starts_with("bitmap-global-index-")); - assert_eq!(index_file.row_count, 3); - assert!(index_file.file_size > 0); - - let global_meta = index_file - .global_index_meta - .as_ref() - .expect("global index meta"); - let bitmap_meta = - crate::btree::BTreeIndexMeta::deserialize(global_meta.index_meta.as_ref().unwrap()) - .unwrap(); - assert_eq!(bitmap_meta.first_key, Some(b"alice".to_vec())); - assert_eq!(bitmap_meta.last_key, Some(b"bob".to_vec())); - assert!(!bitmap_meta.has_nulls); - - let index_path = format!("{table_path}/index/{}", index_file.file_name); - let input = table.file_io().new_input(&index_path).unwrap(); - let file_size = input.metadata().await.unwrap().size; - let reader = input.reader().await.unwrap(); - let bitmap_reader = - crate::table::bitmap_global_index_reader::BitmapGlobalIndexReader::open( - Box::new(reader), - file_size, - ) - .await - .unwrap(); - let bitmap = bitmap_reader - .query( - crate::spec::PredicateOperator::Eq, - &[crate::spec::Datum::String("alice".to_string())], - table.schema().fields()[1].data_type(), - ) - .await - .unwrap(); - assert_eq!(bitmap.iter().collect::>(), vec![0, 2]); - - let predicate = PredicateBuilder::new(table.schema().fields()) - .equal("name", crate::spec::Datum::String("alice".to_string())) - .unwrap(); - let row_ranges = evaluate_global_index(GlobalIndexEvaluation { - file_io: table.file_io(), - table_path: table.location(), - index_entries: &index_entries, - predicates: &[predicate], - schema_fields: table.schema().fields(), - search_mode: GlobalIndexSearchMode::Fast, - global_index_thread_num: 32, - btree_fallback_scan_max_size: i64::MAX, - bitmap_fallback_scan_max_size: i64::MAX, - fm_read_options: crate::fm_index::FMReadOptions::default(), - next_row_id: snapshot.next_row_id(), - data_ranges: &[], - }) - .await - .unwrap() - .unwrap(); - assert_eq!(row_ranges, vec![RowRange::new(0, 0), RowRange::new(2, 2)]); - } - - #[tokio::test] - async fn test_execute_multivalue_index_and_array_queries_end_to_end() { - let table_path = "memory:/test_multivalue_global_index_builder_e2e"; - let table = multivalue_table(table_path); - setup_dirs(&table).await; - - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&multivalue_batch()) - .await - .unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(table_write.prepare_commit().await.unwrap()) - .await - .unwrap(); - - let shard_count = table - .new_btree_global_index_build_builder() - .with_index_column("items") - .with_index_type(MULTIVALUE_GLOBAL_INDEX_TYPE) - .execute() - .await - .unwrap(); - assert_eq!(shard_count, 1); - - let snapshot = SnapshotManager::new(table.file_io().clone(), table.location().to_string()) - .get_latest_snapshot() - .await - .unwrap() - .unwrap(); - let index_manifest = snapshot.index_manifest().expect("index manifest"); - let index_entries = IndexManifest::read( - table.file_io(), - &format!("{table_path}/manifest/{index_manifest}"), - ) - .await - .unwrap(); - assert_eq!(index_entries.len(), 1); - - let index_file = &index_entries[0].index_file; - assert_eq!(index_file.index_type, MULTIVALUE_GLOBAL_INDEX_TYPE); - assert_eq!(index_file.row_count, 5, "source rows, not postings"); - let global_meta = index_file.global_index_meta.as_ref().unwrap(); - let serialized_meta = global_meta.index_meta.as_deref().unwrap(); - let sorted_meta = BTreeIndexMeta::deserialize(serialized_meta).unwrap(); - assert_eq!(serialized_meta, sorted_meta.serialize()); - let element_type = DataType::Int(IntType::new()); - - let index_path = format!("{table_path}/index/{}", index_file.file_name); - let input = table.file_io().new_input(&index_path).unwrap(); - let file_size = input.metadata().await.unwrap().size; - let reader = input.reader().await.unwrap(); - let bitmap_reader = - crate::table::bitmap_global_index_reader::BitmapGlobalIndexReader::open( - Box::new(reader), - file_size, - ) - .await - .unwrap(); - let contains = bitmap_reader - .query( - PredicateOperator::ArrayContains, - &[Datum::Int(10)], - &element_type, - ) - .await - .unwrap(); - assert_eq!(contains.iter().collect::>(), vec![0, 3]); - let overlap = bitmap_reader - .query( - PredicateOperator::ArraysOverlap, - &[Datum::Int(40), Datum::Int(10), Datum::Int(10)], - &element_type, - ) - .await - .unwrap(); - assert_eq!(overlap.iter().collect::>(), vec![0, 3, 4]); - let contains_all = bitmap_reader - .query( - PredicateOperator::ArrayContainsAll, - &[Datum::Int(10), Datum::Int(30), Datum::Int(10)], - &element_type, - ) - .await - .unwrap(); - assert_eq!(contains_all.iter().collect::>(), vec![3]); - - let fields = table.schema().fields(); - let contains_all_predicate = PredicateBuilder::new(fields) - .array_contains_all( - "items", - vec![Datum::Int(10), Datum::Int(30), Datum::Int(10)], - ) - .unwrap(); - let ranges = evaluate_global_index(GlobalIndexEvaluation { - file_io: table.file_io(), - table_path: table.location(), - index_entries: &index_entries, - predicates: std::slice::from_ref(&contains_all_predicate), - schema_fields: fields, - search_mode: GlobalIndexSearchMode::Fast, - global_index_thread_num: 32, - btree_fallback_scan_max_size: i64::MAX, - bitmap_fallback_scan_max_size: i64::MAX, - fm_read_options: crate::fm_index::FMReadOptions::default(), - next_row_id: snapshot.next_row_id(), - data_ranges: &[], - }) - .await - .unwrap() - .unwrap(); - assert_eq!(ranges, vec![RowRange::new(3, 3)]); - - assert_eq!(scan_ids(&table, contains_all_predicate).await, vec![4]); - let empty_contains_all = PredicateBuilder::new(fields) - .array_contains_all("items", vec![]) - .unwrap(); - assert_eq!( - scan_ids(&table, empty_contains_all).await, - vec![1, 2, 4, 5], - "empty contains-all matches every non-null array and must fall back" - ); - } - - #[tokio::test] - async fn test_bitmap_floating_candidates_preserve_residual_results() { - let table_path = "memory:/test_bitmap_floating_residual_candidates"; - let schema = Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column("f", DataType::Float(FloatType::new())) - .column("d", DataType::Double(DoubleType::new())) - .options(table_options("100")) - .build() - .unwrap(); - let table = Table::new( - FileIOBuilder::new("memory").build().unwrap(), - Identifier::new("default", "test_bitmap_floating_residual_candidates"), - table_path.to_string(), - TableSchema::new(0, &schema), - None, - ); - setup_dirs(&table).await; - - let float_negative_nan = f32::from_bits(0xffc0_0001); - let double_negative_nan = f64::from_bits(0xfff8_0000_0000_0001); - let arrow_schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", ArrowDataType::Int32, false), - ArrowField::new("f", ArrowDataType::Float32, true), - ArrowField::new("d", ArrowDataType::Float64, true), - ])); - let batch = RecordBatch::try_new( - arrow_schema, - vec![ - Arc::new(Int32Array::from_iter_values(0..9)) as ArrayRef, - Arc::new(Float32Array::from(vec![ - Some(float_negative_nan), - Some(f32::from_bits(0xffff_1234)), - Some(f32::NAN), - Some(f32::from_bits(0x7fc0_0010)), - Some(-1.0), - Some(-0.0), - Some(0.0), - Some(1.0), - None, - ])) as ArrayRef, - Arc::new(Float64Array::from(vec![ - Some(double_negative_nan), - Some(f64::from_bits(0xffff_1234_5678_9abc)), - Some(f64::NAN), - Some(f64::from_bits(0x7ff8_0000_0000_0010)), - Some(-1.0), - Some(-0.0), - Some(0.0), - Some(1.0), - None, - ])) as ArrayRef, - ], - ) - .unwrap(); - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write.write_arrow_batch(&batch).await.unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - for column in ["f", "d"] { - let shard_count = table - .new_btree_global_index_build_builder() - .with_index_column(column) - .with_index_type(BITMAP_GLOBAL_INDEX_TYPE) - .execute() - .await - .unwrap(); - assert_eq!(shard_count, 1); - } - - let mut disabled_options = table.schema().options().clone(); - disabled_options.insert("global-index.enabled".to_string(), "false".to_string()); - let table_without_index = Table::new( - table.file_io().clone(), - table.identifier().clone(), - table.location().to_string(), - table.schema().copy_with_replaced_options(disabled_options), - None, - ); - - let predicates = PredicateBuilder::new(table.schema().fields()); - let cases = [ - ( - "Float < 0", - predicates.less_than("f", Datum::Float(0.0)).unwrap(), - vec![0, 1, 4, 5], - ), - ( - "Double < 0", - predicates.less_than("d", Datum::Double(0.0)).unwrap(), - vec![0, 1, 4, 5], - ), - ( - "Float = canonical NaN", - predicates.equal("f", Datum::Float(f32::NAN)).unwrap(), - vec![2], - ), - ( - "Double = canonical NaN", - predicates.equal("d", Datum::Double(f64::NAN)).unwrap(), - vec![2], - ), - ( - "Float = negative NaN", - predicates - .equal("f", Datum::Float(float_negative_nan)) - .unwrap(), - vec![0], - ), - ( - "Double = negative NaN", - predicates - .equal("d", Datum::Double(double_negative_nan)) - .unwrap(), - vec![0], - ), - ( - "Float IN NaNs", - predicates - .is_in( - "f", - vec![Datum::Float(float_negative_nan), Datum::Float(f32::NAN)], - ) - .unwrap(), - vec![0, 2], - ), - ( - "Double IN NaNs", - predicates - .is_in( - "d", - vec![Datum::Double(double_negative_nan), Datum::Double(f64::NAN)], - ) - .unwrap(), - vec![0, 2], - ), - ( - "Float != canonical NaN", - predicates.not_equal("f", Datum::Float(f32::NAN)).unwrap(), - vec![0, 1, 3, 4, 5, 6, 7], - ), - ( - "Double != canonical NaN", - predicates.not_equal("d", Datum::Double(f64::NAN)).unwrap(), - vec![0, 1, 3, 4, 5, 6, 7], - ), - ( - "Float NOT IN", - predicates - .is_not_in("f", vec![Datum::Float(f32::NAN), Datum::Float(0.0)]) - .unwrap(), - vec![0, 1, 3, 4, 5, 7], - ), - ( - "Double NOT IN", - predicates - .is_not_in("d", vec![Datum::Double(f64::NAN), Datum::Double(0.0)]) - .unwrap(), - vec![0, 1, 3, 4, 5, 7], - ), - ( - "Float combined range", - Predicate::and(vec![ - predicates - .greater_or_equal("f", Datum::Float(float_negative_nan)) - .unwrap(), - predicates.less_or_equal("f", Datum::Float(0.0)).unwrap(), - ]), - vec![0, 4, 5, 6], - ), - ( - "Double combined range", - Predicate::and(vec![ - predicates - .greater_or_equal("d", Datum::Double(double_negative_nan)) - .unwrap(), - predicates.less_or_equal("d", Datum::Double(0.0)).unwrap(), - ]), - vec![0, 4, 5, 6], - ), - ]; - - for (name, predicate, expected) in cases { - let without_index = scan_ids(&table_without_index, predicate.clone()).await; - assert_eq!(without_index, expected, "{name}: residual baseline"); - let with_index = scan_ids(&table, predicate).await; - assert_eq!( - with_index, without_index, - "{name}: global index changed rows" - ); - } - } - - /// Bitmap is built through the same sorted builder; a second build with no - /// new data must be a no-op keyed on the bitmap coverage — not error, and - /// not be confused by any btree coverage of the same field. - #[tokio::test] - async fn bitmap_second_build_without_new_data_is_noop() { - let table_path = "memory:/test_bitmap_global_index_second_build_noop"; - let table = test_table_with_path(table_path, table_options("10")); - setup_dirs(&table).await; - - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "carol"])) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - let first_built = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .with_index_type(BITMAP_GLOBAL_INDEX_TYPE) - .execute() - .await - .unwrap(); - assert!( - first_built > 0, - "first bitmap build must index initial rows" - ); - - let files_after_first = latest_bitmap_index_files(&table).await; - assert!(!files_after_first.is_empty()); - - let built = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .with_index_type(BITMAP_GLOBAL_INDEX_TYPE) - .execute() - .await - .unwrap(); - assert_eq!( - built, 0, - "fully-indexed bitmap table must build nothing on re-run" - ); - - let names_first = files_after_first - .iter() - .map(|f| f.file_name.clone()) - .collect::>(); - let names_second = latest_bitmap_index_files(&table) - .await - .iter() - .map(|f| f.file_name.clone()) - .collect::>(); - assert_eq!(names_first, names_second, "re-run must not change entries"); - } - - /// A btree index over the SAME field must NOT count as bitmap coverage: a - /// bitmap build after a btree build over identical rows must still produce a - /// bitmap index (regression guard for the index_type-keyed gap computation — - /// the merge-residual bug hard-coded btree here, which would have skipped - /// these rows for a bitmap build). - #[tokio::test] - async fn bitmap_build_after_btree_on_same_field_still_indexes() { - let table_path = "memory:/test_bitmap_after_btree_same_field"; - let table = test_table_with_path(table_path, table_options("10")); - setup_dirs(&table).await; - - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "carol"])) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - let btree_built = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .execute() - .await - .unwrap(); - assert!(btree_built > 0); - - let bitmap_built = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .with_index_type(BITMAP_GLOBAL_INDEX_TYPE) - .execute() - .await - .unwrap(); - assert!( - bitmap_built > 0, - "bitmap build must index rows even when a btree index already covers the same field" - ); - - let bitmap_files = latest_bitmap_index_files(&table).await; - assert!( - !bitmap_files.is_empty(), - "a bitmap index file must be written" - ); - let coverage = data_row_id_coverage(&table).await; - let bitmap_start = bitmap_files - .iter() - .filter_map(|f| f.global_index_meta.as_ref()) - .map(|m| m.row_range_start) - .min() - .unwrap(); - assert_eq!( - bitmap_start, - coverage[0].from(), - "bitmap coverage must span from the first data row, not skip btree-covered rows" - ); - } - - /// Bitmap incremental: build, append, build again → only the appended range - /// gets a new bitmap file; the first bitmap file is retained (append-only). - #[tokio::test] - async fn bitmap_incremental_build_indexes_only_new_rows() { - let table_path = "memory:/test_bitmap_global_index_incremental"; - let table = test_table_with_path(table_path, table_options("10")); - setup_dirs(&table).await; - - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "carol"])) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - let first_built = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .with_index_type(BITMAP_GLOBAL_INDEX_TYPE) - .execute() - .await - .unwrap(); - assert!(first_built > 0); - let first_names = latest_bitmap_index_files(&table) - .await - .iter() - .map(|f| f.file_name.clone()) - .collect::>(); - let n: i64 = 3; - - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch(vec![4, 5, 6], vec!["dave", "erin", "frank"])) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - let second_built = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .with_index_type(BITMAP_GLOBAL_INDEX_TYPE) - .execute() - .await - .unwrap(); - assert!(second_built > 0, "appended rows must be indexed"); - - let all_files = latest_bitmap_index_files(&table).await; - let all_names = all_files - .iter() - .map(|f| f.file_name.clone()) - .collect::>(); - assert!( - first_names.iter().all(|name| all_names.contains(name)), - "build #1 bitmap files must be retained untouched" - ); - let new_files = all_files - .iter() - .filter(|f| !first_names.contains(&f.file_name)) - .collect::>(); - assert!(!new_files.is_empty(), "build #2 must add new bitmap files"); - for file in new_files { - let meta = file - .global_index_meta - .as_ref() - .expect("global index meta on new bitmap file"); - assert!( - meta.row_range_start >= n, - "new bitmap file range must start at or after {}, got [{}, {}]", - n, - meta.row_range_start, - meta.row_range_end - ); - } - } - - async fn latest_bitmap_index_files(table: &Table) -> Vec { - let snapshot_manager = - SnapshotManager::new(table.file_io().clone(), table.location().to_string()); - let snapshot = snapshot_manager - .get_latest_snapshot() - .await - .unwrap() - .unwrap(); - let Some(index_manifest_name) = snapshot.index_manifest() else { - return Vec::new(); - }; - IndexManifest::read( - table.file_io(), - &snapshot_manager.manifest_path(index_manifest_name), - ) - .await - .unwrap() - .into_iter() - .filter(|entry| { - entry.kind == FileKind::Add && entry.index_file.index_type == BITMAP_GLOBAL_INDEX_TYPE - }) - .map(|entry| entry.index_file) - .collect() - } - - async fn latest_btree_index_files(table: &Table) -> Vec { - let snapshot_manager = - SnapshotManager::new(table.file_io().clone(), table.location().to_string()); - let snapshot = snapshot_manager - .get_latest_snapshot() - .await - .unwrap() - .unwrap(); - let Some(index_manifest_name) = snapshot.index_manifest() else { - return Vec::new(); - }; - IndexManifest::read( - table.file_io(), - &snapshot_manager.manifest_path(index_manifest_name), - ) - .await - .unwrap() - .into_iter() - .filter(|entry| { - entry.kind == FileKind::Add && entry.index_file.index_type == BTREE_GLOBAL_INDEX_TYPE - }) - .map(|entry| entry.index_file) - .collect() - } - - /// Row-id coverage of the committed data files, read back from the data - /// manifest (never hard-coded) and merged into contiguous ranges. Mirrors - /// how `execute` gathers `manifest_entries` so tests observe the exact - /// row-ids the writer assigned. - async fn data_row_id_coverage(table: &Table) -> Vec { - let snapshot_manager = - SnapshotManager::new(table.file_io().clone(), table.location().to_string()); - let snapshot = snapshot_manager - .get_latest_snapshot() - .await - .unwrap() - .unwrap(); - let entries = table - .new_read_builder() - .new_scan() - .with_scan_all_files() - .plan_manifest_entries(&snapshot) - .await - .unwrap(); - let ranges = entries - .iter() - .filter(|entry| *entry.kind() == FileKind::Add) - .filter_map(|entry| { - entry - .file() - .row_id_range() - .map(|(start, end)| RowRange::new(start, end)) - }) - .collect::>(); - merge_row_ranges(ranges) - } - - /// Second build with no new data must be a clean no-op (returns 0), not an - /// overlap error. This is the core bug fix: today the second call errors. - #[tokio::test] - async fn second_build_without_new_data_is_noop() { - let table_path = "memory:/test_btree_global_index_second_build_noop"; - let table = test_table_with_path(table_path, table_options("10")); - setup_dirs(&table).await; - - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "carol"])) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - let first_built = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .execute() - .await - .unwrap(); - assert!(first_built > 0, "first build must index the initial rows"); - - let files_after_first = latest_btree_index_files(&table).await; - assert!(!files_after_first.is_empty()); - - let built = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .execute() - .await - .unwrap(); - assert_eq!(built, 0, "fully-indexed table must build nothing on re-run"); - - let files_after_second = latest_btree_index_files(&table).await; - let names_first = files_after_first - .iter() - .map(|f| f.file_name.clone()) - .collect::>(); - let names_second = files_after_second - .iter() - .map(|f| f.file_name.clone()) - .collect::>(); - assert_eq!( - names_first, names_second, - "re-run must not add or remove index manifest entries" - ); - } - - /// Build, append new rows, build again -> only the appended row range is - /// indexed; the first build's index files are retained untouched (append-only). - #[tokio::test] - async fn incremental_build_indexes_only_new_rows() { - let table_path = "memory:/test_btree_global_index_incremental"; - let table = test_table_with_path(table_path, table_options("10")); - setup_dirs(&table).await; - - // Build #1 over rows [0..3). - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "carol"])) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - let first_built = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .execute() - .await - .unwrap(); - assert!(first_built > 0); - - let first_files = latest_btree_index_files(&table).await; - let first_names = first_files - .iter() - .map(|f| f.file_name.clone()) - .collect::>(); - let n: i64 = 3; - - // Append a second batch (new row-ids [3..6)). - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch(vec![4, 5, 6], vec!["dave", "erin", "frank"])) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - let second_built = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .execute() - .await - .unwrap(); - assert!(second_built > 0, "appended rows must be indexed"); - - let all_files = latest_btree_index_files(&table).await; - let all_names = all_files - .iter() - .map(|f| f.file_name.clone()) - .collect::>(); - - // Every build-#1 file is still present (append-only, no rewrite/delete). - assert!( - first_names.iter().all(|name| all_names.contains(name)), - "build #1 index files must be retained untouched" - ); - - // Every build-#2 file covers only the appended range [N, ..]. - let new_files = all_files - .iter() - .filter(|f| !first_names.contains(&f.file_name)) - .collect::>(); - assert!(!new_files.is_empty(), "build #2 must add new index files"); - for file in new_files { - let meta = file - .global_index_meta - .as_ref() - .expect("global index meta on new btree file"); - assert!( - meta.row_range_start >= n, - "new index file range must start at or after {}, got [{}, {}]", - n, - meta.row_range_start, - meta.row_range_end - ); - } - } - - /// Regression: first build (no existing index) must equal the pre-change - /// full build -- subtraction with empty `indexed` = full coverage. - #[tokio::test] - async fn first_build_indexes_full_coverage() { - let table_path = "memory:/test_btree_global_index_first_full_coverage"; - let table = test_table_with_path(table_path, table_options("10")); - setup_dirs(&table).await; - - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "carol"])) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - let built = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .execute() - .await - .unwrap(); - assert_eq!( - built, 1, - "first build must index the full coverage in one shard" - ); - - let files = latest_btree_index_files(&table).await; - assert_eq!(files.len(), 1); - let meta = files[0] - .global_index_meta - .as_ref() - .expect("global index meta"); - assert_eq!(meta.row_range_start, 0); - assert_eq!(meta.row_range_end, 2); - } - - /// Grid boundary (spec edge 4): with `records-per-range = 4`, an appended - /// gap that spans several grid cells must be split so each new index file's - /// range stays inside one cell, the ranges are contiguous, and together - /// they exactly cover the gap. Row-ids are read back from the manifests, - /// never hard-coded. - #[tokio::test] - async fn incremental_build_splits_gap_across_records_per_range_grid() { - const RPR: i64 = 4; - let table_path = "memory:/test_btree_global_index_grid_boundary"; - let table = test_table_with_path(table_path, table_options("4")); - setup_dirs(&table).await; - - // Build #1 over an initial batch (row-ids the writer assigns). - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "carol"])) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - let first_built = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .execute() - .await - .unwrap(); - assert!(first_built > 0, "first build must index the initial rows"); - - // Row range already covered by build #1 (read back, not hard-coded). - let first_index_files = latest_btree_index_files(&table).await; - let indexed_before = merge_row_ranges( - first_index_files - .iter() - .filter_map(|f| f.global_index_meta.as_ref()) - .map(|m| RowRange::new(m.row_range_start, m.row_range_end)) - .collect(), - ); - assert_eq!( - indexed_before.len(), - 1, - "build #1 should cover one contiguous range" - ); - let gap_start = indexed_before[0].to() + 1; - let before_names = first_index_files - .iter() - .map(|f| f.file_name.clone()) - .collect::>(); - - // Append rows so the new gap crosses records_per_range (=4) boundaries. - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch( - vec![4, 5, 6, 7, 8, 9, 10], - vec!["d", "e", "f", "g", "h", "i", "j"], - )) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - // Total data coverage read back from the data manifest. - let coverage = data_row_id_coverage(&table).await; - assert_eq!( - coverage.len(), - 1, - "appended data must be contiguous with build #1" - ); - let gap_end = coverage[0].to(); - assert!( - gap_end - gap_start + 1 > RPR, - "gap [{gap_start}, {gap_end}] must span more than one records_per_range cell" - ); - - let second_built = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .execute() - .await - .unwrap(); - assert!(second_built > 0, "appended rows must be indexed"); - - // Only the newly written index files (build #1 files are retained). - let mut new_metas = latest_btree_index_files(&table) - .await - .into_iter() - .filter(|f| !before_names.contains(&f.file_name)) - .filter_map(|f| f.global_index_meta) - .map(|m| (m.row_range_start, m.row_range_end)) - .collect::>(); - new_metas.sort(); - assert!(!new_metas.is_empty(), "build #2 must add new index files"); - - // (a) Each range lies within a single grid cell: no multiple of RPR is - // strictly interior, i.e. start and end share the same cell index. - for (start, end) in &new_metas { - assert!(end >= start, "range must be non-empty: [{start}, {end}]"); - assert_eq!( - start / RPR, - end / RPR, - "range [{start}, {end}] straddles a records_per_range boundary" - ); - } - // (b) Contiguous with no gaps or overlaps. - for pair in new_metas.windows(2) { - assert_eq!( - pair[1].0, - pair[0].1 + 1, - "ranges must be contiguous: {:?} then {:?}", - pair[0], - pair[1] - ); - } - // (c) Together they exactly cover the appended gap [gap_start, gap_end]. - assert_eq!( - new_metas.first().unwrap().0, - gap_start, - "coverage must start at the gap start" - ); - assert_eq!( - new_metas.last().unwrap().1, - gap_end, - "coverage must end at the gap end" - ); - } - - /// Hole splitting (spec edge 5) at build level: a mid-coverage indexed range - /// (constructed directly, as the drop-builder tests build `GlobalIndexMeta` - /// entries) must carve the data coverage into two build segments, one on - /// each side, and the hole itself must not be re-indexed. - #[tokio::test] - async fn incremental_build_splits_gap_around_mid_coverage_indexed_hole() { - let table_path = "memory:/test_btree_global_index_mid_hole"; - // records-per-range large so the grid never splits: the only split is - // the hole itself. - let table = test_table_with_path(table_path, table_options("100")); - setup_dirs(&table).await; - - // Real data spanning row-ids [0, 9]. - let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&data_batch( - (1..=10).collect(), - vec!["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"], - )) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - - let coverage = data_row_id_coverage(&table).await; - assert_eq!(coverage.len(), 1, "data must be one contiguous range"); - assert_eq!(coverage[0].from(), 0); - let last_row = coverage[0].to(); - assert!(last_row >= 9, "need at least 10 rows for a mid hole"); - - // Inject a mid-coverage indexed range [hole_start, hole_end] for the - // `name` field directly into the index manifest. - let name_field_id = find_index_field(&table, "name").unwrap().id(); - let hole_start = 4; - let hole_end = 6; - let synthetic = IndexFileMeta { - index_type: BTREE_GLOBAL_INDEX_TYPE.to_string(), - file_name: "btree-synthetic-hole.index".to_string(), - file_size: 1, - row_count: (hole_end - hole_start + 1), - deletion_vectors_ranges: None, - global_index_meta: Some(GlobalIndexMeta { - row_range_start: hole_start, - row_range_end: hole_end, - index_field_id: name_field_id, - extra_field_ids: None, - source_meta: None, - index_meta: None, - }), - }; - let mut message = CommitMessage::new(BinaryRow::new(0).to_serialized_bytes(), 0, vec![]); - message.new_index_files = vec![synthetic]; - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(vec![message]) - .await - .unwrap(); - - let before_names = latest_btree_index_files(&table) - .await - .into_iter() - .map(|f| f.file_name) - .collect::>(); - - // Build: gap = coverage minus the hole = [0, hole_start-1] and - // [hole_end+1, last_row]; two shards since the grid does not split here. - let built = table - .new_btree_global_index_build_builder() - .with_index_column("name") - .execute() - .await - .unwrap(); - assert_eq!( - built, 2, - "mid-coverage hole must split the gap into two shards" - ); - - let mut new_metas = latest_btree_index_files(&table) - .await - .into_iter() - .filter(|f| !before_names.contains(&f.file_name)) - .filter_map(|f| f.global_index_meta) - .map(|m| (m.row_range_start, m.row_range_end)) - .collect::>(); - new_metas.sort(); - - assert_eq!( - new_metas, - vec![(0, hole_start - 1), (hole_end + 1, last_row)], - "new shards must fill the coverage on both sides of the indexed hole" - ); - for (start, end) in &new_metas { - assert!( - *end < hole_start || *start > hole_end, - "new shard [{start}, {end}] must not overlap indexed hole [{hole_start}, {hole_end}]" - ); - } - } -} +mod tests; diff --git a/crates/paimon/src/table/sorted_global_index_build_builder/extraction.rs b/crates/paimon/src/table/sorted_global_index_build_builder/extraction.rs new file mode 100644 index 000000000..f50799c93 --- /dev/null +++ b/crates/paimon/src/table/sorted_global_index_build_builder/extraction.rs @@ -0,0 +1,343 @@ +// 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. + +//! Projected table reads and conversion into serialized index rows. + +use super::planning::SortedGlobalIndexShard; +use super::validation::checked_row_count; +use super::{SerializeKeyFn, SortedIndexKeyRow}; +use crate::spec::{ + extract_datum_from_array, extract_datum_from_arrow, DataField, DataType, ROW_ID_FIELD_NAME, +}; +use crate::table::global_index_types::MULTIVALUE_GLOBAL_INDEX_TYPE; +use crate::table::stats_filter::group_by_overlapping_row_id; +use crate::table::{DataSplit, DataSplitBuilder, RowRange, Table}; +use crate::{Error, Result}; +use arrow_array::{Array, FixedSizeListArray, Int64Array, LargeListArray, ListArray, RecordBatch}; +use futures::TryStreamExt; +use std::cmp::Ordering; + +pub(super) async fn extract_index_rows( + table: &Table, + shard: &SortedGlobalIndexShard, + index_column: &str, + index_field: &DataField, + index_type: &str, + serialize_key: SerializeKeyFn, +) -> Result> { + let splits = build_read_splits_for_shard(shard)?; + + let mut read_builder = table.new_read_builder(); + read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; + let read = read_builder.new_read()?; + let batches = read.to_arrow(&splits)?.try_collect::>().await?; + let expected_row_count = checked_row_count(shard.row_range_start, shard.row_range_end)?; + if index_type == MULTIVALUE_GLOBAL_INDEX_TYPE { + let DataType::Array(array_type) = index_field.data_type() else { + unreachable!("multivalue field was validated before extraction") + }; + extract_multivalue_index_rows_from_batches( + &batches, + index_column, + array_type.element_type(), + shard.row_range_start, + expected_row_count, + serialize_key, + ) + } else { + extract_index_rows_from_batches( + &batches, + index_column, + index_field.data_type(), + shard.row_range_start, + expected_row_count, + serialize_key, + ) + } +} + +pub(super) fn build_read_splits_for_shard( + shard: &SortedGlobalIndexShard, +) -> Result> { + let shard_range = RowRange::new(shard.row_range_start, shard.row_range_end); + group_by_overlapping_row_id(shard.files.clone()) + .into_iter() + .filter_map(|files| { + let ranges = files + .iter() + .filter_map(|file| { + file.row_id_range() + .and_then(|(start, end)| shard_range.intersect_inclusive(start, end)) + }) + .collect::>(); + let ranges = crate::table::merge_row_ranges(ranges); + if ranges.is_empty() { + return None; + } + let raw_convertible = files.len() == 1; + Some( + DataSplitBuilder::new() + .with_snapshot(shard.snapshot_id) + .with_partition(shard.partition.clone()) + .with_bucket(shard.source_bucket) + .with_bucket_path(shard.bucket_path.clone()) + .with_total_buckets(shard.total_buckets) + .with_data_files(files) + .with_row_ranges(ranges) + .with_raw_convertible(raw_convertible) + .build(), + ) + }) + .collect() +} + +pub(super) fn extract_index_rows_from_batches( + batches: &[RecordBatch], + index_column: &str, + data_type: &DataType, + row_range_start: i64, + expected_row_count: i64, + serialize_key: SerializeKeyFn, +) -> Result> { + let row_count = batches.iter().map(RecordBatch::num_rows).sum::(); + let mut rows = Vec::with_capacity(row_count); + let mut expected_row_id = row_range_start; + for batch in batches { + let value_index = + batch + .schema() + .index_of(index_column) + .map_err(|e| Error::DataInvalid { + message: format!("Index column '{index_column}' not found in read batch: {e}"), + source: None, + })?; + let row_id_index = + batch + .schema() + .index_of(ROW_ID_FIELD_NAME) + .map_err(|e| Error::DataInvalid { + message: format!("_ROW_ID column not found in read batch: {e}"), + source: None, + })?; + let row_ids = batch + .column(row_id_index) + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::DataInvalid { + message: "Sorted global index build requires non-null Int64 _ROW_ID".to_string(), + source: None, + })?; + + for row in 0..batch.num_rows() { + if row_ids.is_null(row) { + return Err(Error::DataInvalid { + message: "Sorted global index build found null _ROW_ID".to_string(), + source: None, + }); + } + let row_id = row_ids.value(row); + if row_id != expected_row_id { + return Err(Error::DataInvalid { + message: format!( + "Sorted global index build expected _ROW_ID {}, got {}", + expected_row_id, row_id + ), + source: None, + }); + } + expected_row_id += 1; + + let key = extract_datum_from_arrow(batch, row, value_index, data_type)? + .map(|datum| serialize_key(&datum, data_type)); + rows.push((key, row_id - row_range_start)); + } + } + let actual_row_count = expected_row_id - row_range_start; + if actual_row_count != expected_row_count { + return Err(Error::DataInvalid { + message: format!( + "Sorted global index build expected {} rows, got {}", + expected_row_count, actual_row_count + ), + source: None, + }); + } + Ok(rows) +} + +pub(super) fn extract_multivalue_index_rows_from_batches( + batches: &[RecordBatch], + index_column: &str, + element_type: &DataType, + row_range_start: i64, + expected_row_count: i64, + serialize_key: SerializeKeyFn, +) -> Result> { + let mut rows = Vec::new(); + let mut expected_row_id = row_range_start; + for batch in batches { + let value_index = + batch + .schema() + .index_of(index_column) + .map_err(|e| Error::DataInvalid { + message: format!("Index column '{index_column}' not found in read batch: {e}"), + source: None, + })?; + let row_id_index = + batch + .schema() + .index_of(ROW_ID_FIELD_NAME) + .map_err(|e| Error::DataInvalid { + message: format!("_ROW_ID column not found in read batch: {e}"), + source: None, + })?; + let row_ids = batch + .column(row_id_index) + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::DataInvalid { + message: "Multivalue global index build requires non-null Int64 _ROW_ID" + .to_string(), + source: None, + })?; + + #[derive(Clone, Copy)] + enum ArrayLayout<'a> { + List(&'a ListArray), + LargeList(&'a LargeListArray), + Fixed(&'a FixedSizeListArray), + } + let column = batch.column(value_index); + let layout = if let Some(array) = column.as_any().downcast_ref::() { + ArrayLayout::List(array) + } else if let Some(array) = column.as_any().downcast_ref::() { + ArrayLayout::LargeList(array) + } else if let Some(array) = column.as_any().downcast_ref::() { + ArrayLayout::Fixed(array) + } else { + return Err(Error::DataInvalid { + message: format!( + "Multivalue global index extraction requires an Arrow list column, got {:?}", + column.data_type() + ), + source: None, + }); + }; + let values = match layout { + ArrayLayout::List(array) => array.values(), + ArrayLayout::LargeList(array) => array.values(), + ArrayLayout::Fixed(array) => array.values(), + }; + + for row in 0..batch.num_rows() { + if row_ids.is_null(row) { + return Err(Error::DataInvalid { + message: "Multivalue global index build found null _ROW_ID".to_string(), + source: None, + }); + } + let row_id = row_ids.value(row); + if row_id != expected_row_id { + return Err(Error::DataInvalid { + message: format!( + "Multivalue global index build expected _ROW_ID {}, got {}", + expected_row_id, row_id + ), + source: None, + }); + } + expected_row_id += 1; + + let is_null = match layout { + ArrayLayout::List(array) => array.is_null(row), + ArrayLayout::LargeList(array) => array.is_null(row), + ArrayLayout::Fixed(array) => array.is_null(row), + }; + if is_null { + continue; + } + let (start, end) = match layout { + ArrayLayout::List(array) => { + let offsets = array.value_offsets(); + ( + usize::try_from(offsets[row]), + usize::try_from(offsets[row + 1]), + ) + } + ArrayLayout::LargeList(array) => { + let offsets = array.value_offsets(); + ( + usize::try_from(offsets[row]), + usize::try_from(offsets[row + 1]), + ) + } + ArrayLayout::Fixed(array) => { + let start = usize::try_from(array.value_offset(row)); + let end = usize::try_from(array.value_offset(row) + array.value_length()); + (start, end) + } + }; + let (start, end) = match (start, end) { + (Ok(start), Ok(end)) => (start, end), + _ => { + return Err(Error::DataInvalid { + message: "Multivalue global index found a negative array offset" + .to_string(), + source: None, + }) + } + }; + for element_index in start..end { + if let Some(datum) = + extract_datum_from_array(values, element_index, value_index, element_type)? + { + rows.push(( + Some(serialize_key(&datum, element_type)), + row_id - row_range_start, + )); + } + } + } + } + + let actual_row_count = expected_row_id - row_range_start; + if actual_row_count != expected_row_count { + return Err(Error::DataInvalid { + message: format!( + "Multivalue global index build expected {} source rows, got {}", + expected_row_count, actual_row_count + ), + source: None, + }); + } + Ok(rows) +} + +pub(super) fn sort_index_rows( + rows: &mut [SortedIndexKeyRow], + cmp: &dyn Fn(&[u8], &[u8]) -> Ordering, +) { + rows.sort_by(|left, right| match (&left.0, &right.0) { + (None, None) => left.1.cmp(&right.1), + (None, Some(_)) => Ordering::Less, + (Some(_), None) => Ordering::Greater, + (Some(left_key), Some(right_key)) => { + cmp(left_key, right_key).then_with(|| left.1.cmp(&right.1)) + } + }); +} diff --git a/crates/paimon/src/table/sorted_global_index_build_builder/planning.rs b/crates/paimon/src/table/sorted_global_index_build_builder/planning.rs new file mode 100644 index 000000000..7c52fd806 --- /dev/null +++ b/crates/paimon/src/table/sorted_global_index_build_builder/planning.rs @@ -0,0 +1,245 @@ +// 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. + +//! Planning of uncovered, row-range-aligned index shards. + +use super::validation::ranges_overlap; +use crate::spec::{ + bucket_dir_name, BinaryRow, CoreOptions, DataField, DataFileMeta, FileKind, ManifestEntry, + PartitionComputer, +}; +use crate::table::source::{exclude_row_ranges, is_data_evolution_normal_file}; +use crate::table::RowRange; +use crate::{Error, Result}; +use std::collections::HashMap; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct SortedGlobalIndexShard { + pub partition: BinaryRow, + pub partition_bytes: Vec, + pub files: Vec, + pub row_range_start: i64, + pub row_range_end: i64, + pub(super) snapshot_id: i64, + pub(super) source_bucket: i32, + pub(super) total_buckets: i32, + pub(super) bucket_path: String, +} + +#[allow(clippy::too_many_arguments)] +pub(super) fn plan_sorted_index_shards( + table_location: &str, + partition_keys: &[String], + schema_fields: &[DataField], + core_options: &CoreOptions, + snapshot_id: i64, + entries: Vec, + records_per_range: i64, + indexed: &[RowRange], +) -> Result> { + if records_per_range <= 0 { + return Err(Error::DataInvalid { + message: format!( + "Option 'sorted-index.records-per-range' must be greater than 0, got: {records_per_range}" + ), + source: None, + }); + } + + let mut by_partition_bucket: HashMap<(Vec, i32, i32), Vec> = HashMap::new(); + for entry in entries { + if *entry.kind() != FileKind::Add { + continue; + } + if entry.file().first_row_id.is_none() { + return Err(Error::DataInvalid { + message: format!( + "Data file '{}' is missing first_row_id; cannot build a complete sorted global index", + entry.file().file_name + ), + source: None, + }); + } + let (partition, bucket, total_buckets, file) = entry.into_parts(); + by_partition_bucket + .entry((partition, bucket, total_buckets)) + .or_default() + .push(file); + } + + let mut result = Vec::new(); + for ((partition_bytes, source_bucket, total_buckets), files) in by_partition_bucket { + let partition = if partition_keys.is_empty() { + BinaryRow::new(0) + } else { + BinaryRow::from_serialized_bytes(&partition_bytes)? + }; + let bucket_path = bucket_path( + table_location, + partition_keys, + schema_fields, + core_options, + &partition, + source_bucket, + )?; + let normal_groups = group_normal_file_ranges(files)?; + for group in normal_groups { + let (coverage_start, coverage_end) = normal_coverage_range(&group.files)?; + let build_segments = + exclude_row_ranges(&[RowRange::new(coverage_start, coverage_end)], indexed); + for seg in build_segments { + let seg_start = seg.from(); + let seg_end = seg.to(); + let start_range = seg_start / records_per_range; + let end_range = seg_end / records_per_range; + for range_id in start_range..=end_range { + let range_start = range_id * records_per_range; + let range_end = range_start + records_per_range - 1; + let row_range_start = seg_start.max(range_start); + let row_range_end = seg_end.min(range_end); + result.push(SortedGlobalIndexShard { + partition: partition.clone(), + partition_bytes: partition_bytes.clone(), + files: group.files.clone(), + row_range_start, + row_range_end, + snapshot_id, + source_bucket, + total_buckets, + bucket_path: bucket_path.clone(), + }); + } + } + } + } + result.sort_by(|a, b| { + a.partition + .to_serialized_bytes() + .cmp(&b.partition.to_serialized_bytes()) + .then(a.source_bucket.cmp(&b.source_bucket)) + .then(a.row_range_start.cmp(&b.row_range_start)) + }); + Ok(result) +} + +#[derive(Debug)] +struct PlannedFileGroup { + files: Vec, +} + +fn group_normal_file_ranges(files: Vec) -> Result> { + if files.is_empty() { + return Ok(Vec::new()); + } + for file in &files { + file.row_id_range().ok_or_else(|| Error::DataInvalid { + message: format!( + "Data file '{}' is missing first_row_id; cannot build a complete sorted global index", + file.file_name + ), + source: None, + })?; + } + + let mut normal_ranges = files + .iter() + .filter(|file| is_data_evolution_normal_file(file)) + .filter_map(DataFileMeta::row_id_range) + .collect::>(); + normal_ranges.sort_by_key(|(start, _)| *start); + + let mut coverage_ranges: Vec<(i64, i64)> = Vec::new(); + for (file_start, file_end) in normal_ranges { + match coverage_ranges.last_mut() { + Some((_, end)) if file_start <= *end + 1 => { + *end = (*end).max(file_end); + } + _ => coverage_ranges.push((file_start, file_end)), + } + } + + coverage_ranges + .into_iter() + .map(|(start, end)| { + let mut group_files = files + .iter() + .filter(|file| { + file.row_id_range().is_some_and(|(file_start, file_end)| { + ranges_overlap(start, end, file_start, file_end) + }) + }) + .cloned() + .collect::>(); + group_files.sort_by_key(|file| { + ( + file.first_row_id.unwrap_or(i64::MAX), + !is_data_evolution_normal_file(file), + file.file_name.clone(), + ) + }); + Ok(PlannedFileGroup { files: group_files }) + }) + .collect() +} + +fn normal_coverage_range(files: &[DataFileMeta]) -> Result<(i64, i64)> { + let mut start = None; + let mut end = None; + for file in files + .iter() + .filter(|file| is_data_evolution_normal_file(file)) + { + let (file_start, file_end) = file.row_id_range().ok_or_else(|| Error::DataInvalid { + message: format!( + "Data file '{}' is missing first_row_id; cannot build a complete sorted global index", + file.file_name + ), + source: None, + })?; + start = Some(start.map_or(file_start, |value: i64| value.min(file_start))); + end = Some(end.map_or(file_end, |value: i64| value.max(file_end))); + } + start.zip(end).ok_or_else(|| Error::DataInvalid { + message: "Sorted global index shard has no normal data files".to_string(), + source: None, + }) +} + +fn bucket_path( + table_location: &str, + partition_keys: &[String], + schema_fields: &[DataField], + core_options: &CoreOptions, + partition: &BinaryRow, + bucket: i32, +) -> Result { + let base = table_location.trim_end_matches('/'); + if partition_keys.is_empty() { + return Ok(format!("{base}/{}", bucket_dir_name(bucket))); + } + let computer = PartitionComputer::new( + partition_keys, + schema_fields, + core_options.partition_default_name(), + core_options.legacy_partition_name(), + )?; + Ok(format!( + "{base}/{}{}", + computer.generate_partition_path(partition)?, + bucket_dir_name(bucket) + )) +} diff --git a/crates/paimon/src/table/sorted_global_index_build_builder/tests.rs b/crates/paimon/src/table/sorted_global_index_build_builder/tests.rs new file mode 100644 index 000000000..1c8c438f1 --- /dev/null +++ b/crates/paimon/src/table/sorted_global_index_build_builder/tests.rs @@ -0,0 +1,2332 @@ +// 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::btree::BTreeIndexMeta; +use crate::catalog::Identifier; +use crate::io::FileIOBuilder; +use crate::spec::stats::BinaryTableStats; +use crate::spec::{ + ArrayType, BinaryRow, BinaryRowBuilder, BinaryType, DataField, DataFileMeta, DoubleType, + FileKind, FloatType, GlobalIndexMeta, GlobalIndexSearchMode, IndexFileMeta, IndexManifest, + IntType, ManifestEntry, Predicate, PredicateBuilder, PredicateOperator, Schema, TableSchema, + TimeType, VarBinaryType, VarCharType, ROW_ID_FIELD_NAME, +}; +use crate::table::global_index_scanner::{evaluate_global_index, GlobalIndexEvaluation}; +use crate::table::global_index_types::MULTIVALUE_GLOBAL_INDEX_TYPE; +use crate::table::{merge_row_ranges, SnapshotManager, TableCommit, TableWrite}; +use arrow_array::builder::{Int32Builder, ListBuilder, Time32MillisecondBuilder}; +use arrow_array::{ + Array, ArrayRef, Float32Array, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray, +}; +use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; +use chrono::{DateTime, Utc}; +use futures::TryStreamExt; +use std::sync::Arc; + +/// A row range wider than `i32::MAX` yields the full count instead of being +/// rejected, and an inverted or overflowing range is still an error. +#[test] +fn checked_row_count_spans_beyond_i32() { + let start = 0; + let end = i64::from(i32::MAX) + 10; + assert_eq!( + super::checked_row_count(start, end).unwrap(), + i64::from(i32::MAX) + 11 + ); + assert!(super::checked_row_count(5, 4).is_err()); + assert!(super::checked_row_count(i64::MIN, i64::MAX).is_err()); +} + +fn data_file(name: &str, first_row_id: Option, row_count: i64) -> DataFileMeta { + DataFileMeta { + file_name: name.to_string(), + file_size: 128, + row_count, + min_key: vec![], + max_key: vec![], + key_stats: BinaryTableStats::new(vec![], vec![], vec![]), + value_stats: BinaryTableStats::new(vec![], vec![], vec![]), + min_sequence_number: 0, + max_sequence_number: 0, + schema_id: 0, + level: 0, + extra_files: vec![], + creation_time: Some( + "2024-09-06T07:45:55.039+00:00" + .parse::>() + .unwrap(), + ), + delete_row_count: None, + embedded_index: None, + first_row_id, + write_cols: None, + external_path: None, + file_source: None, + value_stats_cols: None, + column_max_sequence_numbers: None, + } +} + +fn partial_file(name: &str, first_row_id: Option, row_count: i64) -> DataFileMeta { + let mut file = data_file(name, first_row_id, row_count); + file.write_cols = Some(vec!["name".to_string()]); + file +} + +fn manifest_entry(file: DataFileMeta) -> ManifestEntry { + manifest_entry_with_bucket(file, 0, 1) +} + +fn manifest_entry_with_bucket( + file: DataFileMeta, + bucket: i32, + total_buckets: i32, +) -> ManifestEntry { + ManifestEntry::new(FileKind::Add, vec![], bucket, total_buckets, file, 2) +} + +fn table_options(records_per_range: &str) -> HashMap { + HashMap::from([ + ("row-tracking.enabled".to_string(), "true".to_string()), + ("data-evolution.enabled".to_string(), "true".to_string()), + ("global-index.enabled".to_string(), "true".to_string()), + ( + "sorted-index.records-per-range".to_string(), + records_per_range.to_string(), + ), + ]) +} + +fn test_table(options: HashMap) -> Table { + test_table_with_path("memory:/test_btree_global_index_builder", options) +} + +fn test_table_with_path(table_path: &str, options: HashMap) -> Table { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("name", DataType::VarChar(VarCharType::string_type())) + .options(options) + .build() + .unwrap(); + Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "test_table"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ) +} + +fn multivalue_table(table_path: &str) -> Table { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + "items", + DataType::Array(ArrayType::new(DataType::Int(IntType::new()))), + ) + .options(table_options("10")) + .build() + .unwrap(); + Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "test_multivalue_table"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ) +} + +fn plan( + entries: Vec, + records_per_range: i64, +) -> Result> { + let table = test_table(table_options(&records_per_range.to_string())); + let core = CoreOptions::new(table.schema().options()); + plan_sorted_index_shards( + table.location(), + table.schema().partition_keys(), + table.schema().fields(), + &core, + 1, + entries, + records_per_range, + &[], + ) +} + +#[test] +fn test_planner_splits_single_file_across_ranges() { + let shards = plan(vec![manifest_entry(data_file("a", Some(0), 25))], 10).unwrap(); + + assert_eq!( + shards + .iter() + .map(|s| (s.row_range_start, s.row_range_end)) + .collect::>(), + vec![(0, 9), (10, 19), (20, 24)] + ); +} + +#[test] +fn test_planner_merges_contiguous_normal_files() { + let shards = plan( + vec![ + manifest_entry(data_file("a", Some(0), 5)), + manifest_entry(data_file("b", Some(5), 5)), + ], + 20, + ) + .unwrap(); + + assert_eq!(shards.len(), 1); + assert_eq!((shards[0].row_range_start, shards[0].row_range_end), (0, 9)); +} + +#[test] +fn test_planner_splits_row_id_gap_into_separate_shards() { + let shards = plan( + vec![ + manifest_entry(data_file("a", Some(0), 5)), + manifest_entry(data_file("b", Some(10), 5)), + ], + 20, + ) + .unwrap(); + + assert_eq!( + shards + .iter() + .map(|s| (s.row_range_start, s.row_range_end)) + .collect::>(), + vec![(0, 4), (10, 14)] + ); +} + +#[test] +fn test_planner_rejects_missing_first_row_id() { + let err = plan(vec![manifest_entry(data_file("a", None, 5))], 10) + .expect_err("missing first_row_id should fail"); + assert!( + matches!(err, Error::DataInvalid { message, .. } if message.contains("missing first_row_id")) + ); +} + +#[test] +fn test_planner_keeps_buckets_separate() { + let shards = plan( + vec![ + manifest_entry_with_bucket(data_file("a", Some(0), 5), 0, 2), + manifest_entry_with_bucket(data_file("b", Some(5), 5), 1, 2), + ], + 20, + ) + .unwrap(); + + assert_eq!( + shards + .iter() + .map(|s| ( + s.source_bucket, + s.total_buckets, + s.row_range_start, + s.row_range_end + )) + .collect::>(), + vec![(0, 2, 0, 4), (1, 2, 5, 9)] + ); +} + +#[test] +fn test_planner_keeps_partial_file_in_read_group_without_expanding_coverage() { + let shards = plan( + vec![ + manifest_entry(data_file("base", Some(0), 5)), + manifest_entry(partial_file("partial", Some(0), 5)), + ], + 20, + ) + .unwrap(); + + assert_eq!(shards.len(), 1); + assert_eq!((shards[0].row_range_start, shards[0].row_range_end), (0, 4)); + assert_eq!(shards[0].files.len(), 2); +} + +#[test] +fn test_build_read_splits_groups_only_overlapping_partial_files() { + let shards = plan( + vec![ + manifest_entry(data_file("a", Some(0), 5)), + manifest_entry(data_file("b", Some(5), 5)), + manifest_entry(partial_file("partial", Some(0), 5)), + ], + 20, + ) + .unwrap(); + assert_eq!(shards.len(), 1); + + let splits = build_read_splits_for_shard(&shards[0]).unwrap(); + + assert_eq!(splits.len(), 2); + assert_eq!( + splits[0] + .data_files() + .iter() + .map(|file| file.file_name.as_str()) + .collect::>(), + vec!["a", "partial"] + ); + assert_eq!(splits[0].row_ranges(), Some(&[RowRange::new(0, 4)][..])); + assert!(!splits[0].raw_convertible()); + + assert_eq!( + splits[1] + .data_files() + .iter() + .map(|file| file.file_name.as_str()) + .collect::>(), + vec!["b"] + ); + assert_eq!(splits[1].row_ranges(), Some(&[RowRange::new(5, 9)][..])); + assert!(splits[1].raw_convertible()); +} + +#[test] +fn test_validate_btree_field_rejects_complex_type() { + let field = DataField::new( + 0, + "items".to_string(), + DataType::Array(crate::spec::ArrayType::new(DataType::Int(IntType::new()))), + ); + let err = validate_btree_field(&field).expect_err("array should be rejected"); + assert!(matches!(err, Error::Unsupported { message } if message.contains("scalar"))); +} + +#[test] +fn test_validate_btree_field_rejects_binary_types() { + for data_type in [ + DataType::Binary(BinaryType::new(4).unwrap()), + DataType::VarBinary(VarBinaryType::try_new(true, 4).unwrap()), + ] { + let field = DataField::new(0, "bytes".to_string(), data_type); + let err = validate_btree_field(&field).expect_err("binary should be rejected"); + assert!(matches!(err, Error::Unsupported { message } if message.contains("scalar"))); + } +} + +fn index_batch(values: Vec>, row_ids: Vec>) -> RecordBatch { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, true), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(values)) as ArrayRef, + Arc::new(Int64Array::from(row_ids)) as ArrayRef, + ], + ) + .unwrap() +} + +#[test] +fn test_extract_index_rows_serializes_keys_and_local_row_ids() { + let batch = index_batch( + vec![Some(10), None, Some(30)], + vec![Some(5), Some(6), Some(7)], + ); + let rows = extract_index_rows_from_batches( + &[batch], + "id", + &DataType::Int(IntType::new()), + 5, + 3, + serialize_datum, + ) + .unwrap(); + + assert_eq!( + rows, + vec![ + (Some(10i32.to_le_bytes().to_vec()), 0), + (None, 1), + (Some(30i32.to_le_bytes().to_vec()), 2), + ] + ); +} + +#[test] +fn test_extract_multivalue_rows_skips_null_arrays_and_elements() { + let element = Arc::new(ArrowField::new("element", ArrowDataType::Int32, true)); + let mut items = ListBuilder::new(Int32Builder::new()).with_field(element.clone()); + items.values().append_value(10); + items.values().append_null(); + items.values().append_value(10); + items.append(true); + items.append(true); // empty array + items.append(false); // null array + items.values().append_value(30); + items.append(true); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("items", ArrowDataType::List(element), true), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(items.finish()) as ArrayRef, + Arc::new(Int64Array::from_iter_values(5..9)) as ArrayRef, + ], + ) + .unwrap(); + + let rows = extract_multivalue_index_rows_from_batches( + &[batch], + "items", + &DataType::Int(IntType::new()), + 5, + 4, + serialize_bitmap_datum, + ) + .unwrap(); + + assert_eq!( + rows, + vec![ + (Some(10i32.to_le_bytes().to_vec()), 0), + (Some(10i32.to_le_bytes().to_vec()), 0), + (Some(30i32.to_le_bytes().to_vec()), 3), + ] + ); +} + +#[test] +fn test_extract_multivalue_time_rows_matches_java_int_serializer() { + let element = Arc::new(ArrowField::new( + "element", + ArrowDataType::Time32(arrow_schema::TimeUnit::Millisecond), + true, + )); + let mut items = ListBuilder::new(Time32MillisecondBuilder::new()).with_field(element.clone()); + items.values().append_value(12_345); + items.values().append_null(); + items.append(true); + items.values().append_value(86_399_999); + items.append(true); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("items", ArrowDataType::List(element), true), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(items.finish()) as ArrayRef, + Arc::new(Int64Array::from_iter_values(7..9)) as ArrayRef, + ], + ) + .unwrap(); + let time_type = DataType::Time(TimeType::new(3).unwrap()); + + let rows = extract_multivalue_index_rows_from_batches( + &[batch], + "items", + &time_type, + 7, + 2, + serialize_bitmap_datum, + ) + .unwrap(); + + assert_eq!( + rows, + vec![ + (Some(12_345i32.to_le_bytes().to_vec()), 0), + (Some(86_399_999i32.to_le_bytes().to_vec()), 1), + ] + ); +} + +#[test] +fn test_index_key_codec_scopes_java_nan_semantics_to_bitmap() { + fn assert_codec( + data_type: DataType, + negative_nan: Datum, + raw_nan_key: Vec, + canonical_nan_key: Vec, + zero: Datum, + ) { + let (btree_cmp, btree_serialize) = + make_index_key_codec(BTREE_GLOBAL_INDEX_TYPE, &data_type); + let btree_nan_key = btree_serialize(&negative_nan, &data_type); + let zero_key = btree_serialize(&zero, &data_type); + assert_eq!(btree_nan_key, raw_nan_key); + assert!(btree_cmp(&btree_nan_key, &zero_key).is_lt()); + + let (bitmap_cmp, bitmap_serialize) = + make_index_key_codec(BITMAP_GLOBAL_INDEX_TYPE, &data_type); + let bitmap_nan_key = bitmap_serialize(&negative_nan, &data_type); + assert_eq!(bitmap_nan_key, canonical_nan_key); + assert!(bitmap_cmp(&bitmap_nan_key, &zero_key).is_gt()); + } + + assert_codec( + DataType::Float(FloatType::new()), + Datum::Float(f32::from_bits(0xffc0_0001)), + 0xffc0_0001u32.to_le_bytes().to_vec(), + 0x7fc0_0000u32.to_le_bytes().to_vec(), + Datum::Float(0.0), + ); + assert_codec( + DataType::Double(DoubleType::new()), + Datum::Double(f64::from_bits(0xfff8_0000_0000_0001)), + 0xfff8_0000_0000_0001u64.to_le_bytes().to_vec(), + 0x7ff8_0000_0000_0000u64.to_le_bytes().to_vec(), + Datum::Double(0.0), + ); +} + +#[test] +fn test_extract_index_rows_rejects_row_id_gap() { + let batch = index_batch(vec![Some(10), Some(30)], vec![Some(5), Some(7)]); + let err = extract_index_rows_from_batches( + &[batch], + "id", + &DataType::Int(IntType::new()), + 5, + 2, + serialize_datum, + ) + .expect_err("row-id gap should fail"); + + assert!( + matches!(err, Error::DataInvalid { message, .. } if message.contains("expected _ROW_ID")) + ); +} + +#[test] +fn test_sort_index_rows_orders_nulls_then_keys() { + let mut rows = vec![ + (Some(3i32.to_le_bytes().to_vec()), 0), + (None, 1), + (Some(1i32.to_le_bytes().to_vec()), 2), + (Some(1i32.to_le_bytes().to_vec()), 3), + ]; + let cmp = make_key_comparator(&DataType::Int(IntType::new())); + + sort_index_rows(&mut rows, &cmp); + + assert_eq!( + rows, + vec![ + (None, 1), + (Some(1i32.to_le_bytes().to_vec()), 2), + (Some(1i32.to_le_bytes().to_vec()), 3), + (Some(3i32.to_le_bytes().to_vec()), 0), + ] + ); +} + +#[test] +fn test_extract_index_rows_accepts_string_column() { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("name", ArrowDataType::Utf8, true), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(vec![Some("alice"), None])) as ArrayRef, + Arc::new(Int64Array::from(vec![Some(10), Some(11)])) as ArrayRef, + ], + ) + .unwrap(); + + let rows = extract_index_rows_from_batches( + &[batch], + "name", + &DataType::VarChar(VarCharType::string_type()), + 10, + 2, + serialize_datum, + ) + .unwrap(); + + assert_eq!(rows, vec![(Some(b"alice".to_vec()), 0), (None, 1)]); +} + +fn data_batch(ids: Vec, names: Vec<&str>) -> RecordBatch { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("name", ArrowDataType::Utf8, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(ids)) as ArrayRef, + Arc::new(StringArray::from(names)) as ArrayRef, + ], + ) + .unwrap() +} + +fn nullable_name_batch(ids: Vec, names: Vec>) -> RecordBatch { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("name", ArrowDataType::Utf8, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(ids)) as ArrayRef, + Arc::new(StringArray::from(names)) as ArrayRef, + ], + ) + .unwrap() +} + +fn multivalue_batch() -> RecordBatch { + let element = Arc::new(ArrowField::new("element", ArrowDataType::Int32, true)); + let mut items = ListBuilder::new(Int32Builder::new()).with_field(element.clone()); + + items.values().append_value(10); + items.values().append_null(); + items.values().append_value(10); + items.append(true); + + items.append(true); // empty array + items.append(false); // null array + + items.values().append_value(10); + items.values().append_value(30); + items.append(true); + + items.values().append_value(30); + items.values().append_value(40); + items.append(true); + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("items", ArrowDataType::List(element), true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from_iter_values(1..=5)) as ArrayRef, + Arc::new(items.finish()) as ArrayRef, + ], + ) + .unwrap() +} + +async fn setup_dirs(table: &Table) { + table + .file_io() + .mkdirs(&format!("{}/snapshot/", table.location())) + .await + .unwrap(); + table + .file_io() + .mkdirs(&format!("{}/manifest/", table.location())) + .await + .unwrap(); +} + +async fn scan_ids(table: &Table, predicate: Predicate) -> Vec { + let mut builder = table.new_read_builder(); + builder.with_filter(predicate); + let plan = builder.new_scan().plan().await.unwrap(); + let read = builder.new_read().unwrap(); + let batches = read + .to_arrow(plan.splits()) + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut ids = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + ids.sort_unstable(); + ids +} + +#[tokio::test] +async fn test_execute_writes_btree_index_manifest_and_file() { + let table_path = "memory:/test_btree_global_index_builder_e2e"; + let table = test_table_with_path(table_path, table_options("10")); + setup_dirs(&table).await; + + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "alice"])) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let shard_count = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .execute() + .await + .unwrap(); + assert_eq!(shard_count, 1); + + let snapshot_manager = + SnapshotManager::new(table.file_io().clone(), table.location().to_string()); + let snapshot = snapshot_manager + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let index_manifest = snapshot.index_manifest().expect("index manifest"); + let index_entries = IndexManifest::read( + table.file_io(), + &format!("{table_path}/manifest/{index_manifest}"), + ) + .await + .unwrap(); + assert_eq!(index_entries.len(), 1); + + let index_file = &index_entries[0].index_file; + assert_eq!(index_file.index_type, BTREE_GLOBAL_INDEX_TYPE); + assert!(index_file.file_name.starts_with("btree-global-index-")); + assert_eq!(index_file.row_count, 3); + assert!(index_file.file_size > 0); + + let global_meta = index_file + .global_index_meta + .as_ref() + .expect("global index meta"); + assert_eq!(global_meta.row_range_start, 0); + assert_eq!(global_meta.row_range_end, 2); + assert_eq!(global_meta.index_field_id, 1); + let btree_meta = + crate::btree::BTreeIndexMeta::deserialize(global_meta.index_meta.as_ref().unwrap()) + .unwrap(); + assert_eq!(btree_meta.first_key, Some(b"alice".to_vec())); + assert_eq!(btree_meta.last_key, Some(b"bob".to_vec())); + assert!(!btree_meta.has_nulls); + + let predicate = PredicateBuilder::new(table.schema().fields()) + .equal("name", crate::spec::Datum::String("alice".to_string())) + .unwrap(); + let row_ranges = evaluate_global_index(GlobalIndexEvaluation { + file_io: table.file_io(), + table_path: table.location(), + index_entries: &index_entries, + predicates: &[predicate], + schema_fields: table.schema().fields(), + search_mode: GlobalIndexSearchMode::Fast, + global_index_thread_num: 32, + btree_fallback_scan_max_size: i64::MAX, + bitmap_fallback_scan_max_size: i64::MAX, + fm_read_options: crate::fm_index::FMReadOptions::default(), + next_row_id: snapshot.next_row_id(), + data_ranges: &[], + }) + .await + .unwrap() + .unwrap(); + assert_eq!(row_ranges, vec![RowRange::new(0, 0), RowRange::new(2, 2)]); + + // Reopen the same table without an explicit global-index override and + // verify that the regular scan path still uses the committed index. + let mut options = table.schema().options().clone(); + assert_eq!( + options.remove("global-index.enabled"), + Some("true".to_string()) + ); + let scan_table = Table::new( + table.file_io().clone(), + table.identifier().clone(), + table.location().to_string(), + table.schema().copy_with_replaced_options(options), + None, + ); + let predicate = PredicateBuilder::new(scan_table.schema().fields()) + .equal("name", crate::spec::Datum::String("alice".to_string())) + .unwrap(); + let mut read_builder = scan_table.new_read_builder(); + read_builder.with_filter(predicate); + let plan = read_builder.new_scan().plan().await.unwrap(); + + assert_eq!(plan.splits().len(), 1); + assert_eq!( + plan.splits()[0].row_ranges(), + Some(&[RowRange::new(0, 0), RowRange::new(2, 2)][..]) + ); +} + +#[tokio::test] +async fn test_execute_writes_and_queries_fm_index() { + let table_path = "memory:/test_fm_global_index_builder_e2e"; + let mut options = table_options("10"); + options.insert("fm-index.sa-sample-rate".to_string(), "1".to_string()); + options.insert("fm-index.locate-cost-ratio".to_string(), "1".to_string()); + options.insert("fm-index.compression".to_string(), "none".to_string()); + let table = test_table_with_path(table_path, options); + setup_dirs(&table).await; + + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&nullable_name_batch( + vec![1, 2, 3, 4], + vec![Some("banana"), Some("bandana"), None, Some("")], + )) + .await + .unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(table_write.prepare_commit().await.unwrap()) + .await + .unwrap(); + + assert_eq!( + table + .new_sorted_global_index_build_builder() + .with_index_column("name") + .with_index_type(FM_GLOBAL_INDEX_TYPE) + .execute() + .await + .unwrap(), + 1 + ); + + let snapshot = table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let index_entries = IndexManifest::read( + table.file_io(), + &format!( + "{table_path}/manifest/{}", + snapshot.index_manifest().expect("index manifest") + ), + ) + .await + .unwrap(); + assert_eq!(index_entries.len(), 1); + let index_file = &index_entries[0].index_file; + assert_eq!(index_file.index_type, FM_GLOBAL_INDEX_TYPE); + assert!(index_file.file_name.starts_with("fm-global-index-")); + assert_eq!(index_file.row_count, 4); + crate::fm_index::validate_manifest_meta( + index_file + .global_index_meta + .as_ref() + .unwrap() + .index_meta + .as_ref() + .unwrap(), + ) + .unwrap(); + + let predicates = PredicateBuilder::new(table.schema().fields()); + assert_eq!( + scan_ids( + &table, + predicates + .contains("name", Datum::String("ana".to_string())) + .unwrap(), + ) + .await, + vec![1, 2] + ); + assert_eq!( + scan_ids(&table, predicates.is_null("name").unwrap()).await, + vec![3] + ); + assert_eq!( + scan_ids( + &table, + predicates + .contains("name", Datum::String(String::new())) + .unwrap(), + ) + .await, + vec![1, 2, 4] + ); + + // A dense match that exceeds the FM locate budget must decline the + // index and let the normal row filter scan the source, never produce + // a false empty result. + let mut fallback_options = table.schema().options().clone(); + fallback_options.insert( + "fm-index.locate-cost-ratio".to_string(), + "0.000001".to_string(), + ); + let fallback_table = Table::new( + table.file_io().clone(), + table.identifier().clone(), + table.location().to_string(), + table.schema().copy_with_replaced_options(fallback_options), + None, + ); + let predicate = PredicateBuilder::new(fallback_table.schema().fields()) + .contains("name", Datum::String("a".to_string())) + .unwrap(); + assert_eq!(scan_ids(&fallback_table, predicate).await, vec![1, 2]); +} + +#[tokio::test] +async fn test_failed_fm_build_removes_partial_index_file() { + let table_path = "memory:/test_failed_fm_build_cleanup"; + let table = test_table_with_path(table_path, table_options("3")); + setup_dirs(&table).await; + + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch( + vec![1, 2, 3, 4, 5, 6], + vec!["one", "two", "six", "red", "blue", "value-too-long"], + )) + .await + .unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(table_write.prepare_commit().await.unwrap()) + .await + .unwrap(); + + let mut options = HashMap::new(); + options.insert("fm-index.partition-size".to_string(), "8".to_string()); + options.insert("fm-index.partition-row-count".to_string(), "1".to_string()); + options.insert("fm-index.compression".to_string(), "none".to_string()); + let error = table + .new_sorted_global_index_build_builder() + .with_index_column("name") + .with_index_type(FM_GLOBAL_INDEX_TYPE) + .with_options(options) + .execute() + .await + .expect_err("the oversized FM value must fail the build"); + assert!(matches!(error, Error::DataInvalid { .. })); + + let files = table + .file_io() + .list_status(&format!("{table_path}/index")) + .await + .unwrap(); + assert!( + files + .iter() + .all(|file| !file.path.contains("fm-global-index-")), + "failed FM build left a partial index file: {files:?}" + ); +} + +#[tokio::test] +async fn test_global_index_prunes_during_manifest_read() { + for (search_mode, expected_manifest_pruned, expected_entries_read) in + [("fast", 1, 1), ("full", 1, 1)] + { + let table_path = format!("memory:/test_global_index_manifest_pruning_{search_mode}"); + let mut options = table_options("2"); + options.insert( + "global-index.search-mode".to_string(), + search_mode.to_string(), + ); + let table = test_table_with_path(&table_path, options); + setup_dirs(&table).await; + + for (user, ids, names) in [ + ("writer-1", vec![1, 2], vec!["alice", "bob"]), + ("writer-2", vec![3, 4], vec!["carol", "dave"]), + ] { + let mut table_write = TableWrite::new(&table, user.to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch(ids, names)) + .await + .unwrap(); + TableCommit::new(table.clone(), user.to_string()) + .commit(table_write.prepare_commit().await.unwrap()) + .await + .unwrap(); + } + + table + .new_btree_global_index_build_builder() + .with_index_column("name") + .execute() + .await + .unwrap(); + + let predicate = PredicateBuilder::new(table.schema().fields()) + .equal("name", crate::spec::Datum::String("alice".to_string())) + .unwrap(); + let mut read_builder = table.new_read_builder(); + read_builder.with_filter(predicate); + let (plan, trace) = read_builder.new_scan().plan_with_trace().await.unwrap(); + + assert_eq!( + plan.splits() + .iter() + .flat_map(|split| split.data_files()) + .count(), + 1 + ); + assert_eq!( + trace.manifest_files_pruned_by_row_ranges, + expected_manifest_pruned + ); + assert_eq!(trace.manifest_entries_read, expected_entries_read); + assert_eq!(trace.manifest_entries_pruned_by_row_ranges, 0); + assert_eq!( + trace.manifest_entries_after_manifest_filters, + expected_entries_read + ); + } +} + +#[tokio::test] +async fn test_scalar_full_search_includes_unindexed_rows() { + let mut options = table_options("2"); + options.insert("scalar-index.search-mode".to_string(), "full".to_string()); + let table = test_table_with_path( + "memory:/test_scalar_full_search_includes_unindexed_rows", + options, + ); + setup_dirs(&table).await; + + let mut first_write = TableWrite::new(&table, "writer-1".to_string()).unwrap(); + first_write + .write_arrow_batch(&data_batch(vec![1, 2], vec!["alice", "bob"])) + .await + .unwrap(); + TableCommit::new(table.clone(), "writer-1".to_string()) + .commit(first_write.prepare_commit().await.unwrap()) + .await + .unwrap(); + table + .new_btree_global_index_build_builder() + .with_index_column("name") + .execute() + .await + .unwrap(); + + let mut second_write = TableWrite::new(&table, "writer-2".to_string()).unwrap(); + second_write + .write_arrow_batch(&data_batch(vec![3, 4], vec!["alice", "dave"])) + .await + .unwrap(); + TableCommit::new(table.clone(), "writer-2".to_string()) + .commit(second_write.prepare_commit().await.unwrap()) + .await + .unwrap(); + + let predicate = PredicateBuilder::new(table.schema().fields()) + .equal("name", crate::spec::Datum::String("alice".to_string())) + .unwrap(); + let mut read_builder = table.new_read_builder(); + read_builder.with_filter(predicate); + let plan = read_builder.new_scan().plan().await.unwrap(); + let planned_ranges = merge_row_ranges( + plan.splits() + .iter() + .flat_map(|split| split.row_ranges().unwrap_or_default()) + .cloned() + .collect(), + ); + + assert_eq!( + planned_ranges, + vec![RowRange::new(0, 0), RowRange::new(2, 3)] + ); + assert_eq!( + scan_ids( + &table, + PredicateBuilder::new(table.schema().fields()) + .equal("name", crate::spec::Datum::String("alice".to_string())) + .unwrap(), + ) + .await, + vec![1, 3] + ); +} + +#[tokio::test] +async fn test_empty_global_index_ranges_skip_legacy_manifests() { + for search_mode in ["fast", "full"] { + let table_path = format!("memory:/test_empty_global_index_ranges_{search_mode}"); + let mut options = table_options("10"); + options.insert( + "global-index.search-mode".to_string(), + search_mode.to_string(), + ); + let table = test_table_with_path(&table_path, options); + setup_dirs(&table).await; + + let mut table_write = TableWrite::new(&table, "writer".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch(vec![1, 2], vec!["alice", "bob"])) + .await + .unwrap(); + TableCommit::new(table.clone(), "writer".to_string()) + .commit(table_write.prepare_commit().await.unwrap()) + .await + .unwrap(); + table + .new_btree_global_index_build_builder() + .with_index_column("name") + .execute() + .await + .unwrap(); + + let mut legacy_file = data_file("legacy.parquet", None, 2); + legacy_file.level = 1; + legacy_file.file_source = Some(1); // FileSource.COMPACT + TableCommit::new(table.clone(), "legacy-writer".to_string()) + .commit(vec![CommitMessage::new( + BinaryRowBuilder::new(0).build_serialized(), + 0, + vec![legacy_file], + )]) + .await + .unwrap(); + + let snapshot = table + .snapshot_manager() + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + assert_eq!(snapshot.next_row_id(), Some(2)); + + let predicate = PredicateBuilder::new(table.schema().fields()) + .equal( + "name", + crate::spec::Datum::String("not-present".to_string()), + ) + .unwrap(); + let mut read_builder = table.new_read_builder(); + read_builder.with_filter(predicate); + + let plan = read_builder.new_scan().plan().await.unwrap(); + assert!(plan.splits().is_empty()); + + let (traced_plan, trace) = read_builder.new_scan().plan_with_trace().await.unwrap(); + let delta_plan = read_builder + .new_scan() + .plan_snapshot_delta(&snapshot) + .await + .unwrap(); + + assert!(traced_plan.splits().is_empty()); + assert_eq!(trace.manifest_entries_read, 0); + assert_eq!(trace.final_splits, 0); + assert_eq!(trace.final_files, 0); + assert!(delta_plan.splits().is_empty()); + } +} + +#[tokio::test] +async fn test_detail_mode_defers_manifest_pruning_for_unindexed_ranges() { + let table_path = "memory:/test_detail_manifest_pruning"; + let mut options = table_options("2"); + options.insert("global-index.search-mode".to_string(), "detail".to_string()); + let table = test_table_with_path(table_path, options); + setup_dirs(&table).await; + + let mut table_write = TableWrite::new(&table, "writer-1".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch(vec![1, 2], vec!["alice", "bob"])) + .await + .unwrap(); + TableCommit::new(table.clone(), "writer-1".to_string()) + .commit(table_write.prepare_commit().await.unwrap()) + .await + .unwrap(); + table + .new_btree_global_index_build_builder() + .with_index_column("name") + .execute() + .await + .unwrap(); + + let mut table_write = TableWrite::new(&table, "writer-2".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch(vec![3, 4], vec!["alice", "dave"])) + .await + .unwrap(); + TableCommit::new(table.clone(), "writer-2".to_string()) + .commit(table_write.prepare_commit().await.unwrap()) + .await + .unwrap(); + + let predicate = PredicateBuilder::new(table.schema().fields()) + .equal("name", crate::spec::Datum::String("alice".to_string())) + .unwrap(); + let mut read_builder = table.new_read_builder(); + read_builder.with_filter(predicate); + let (plan, trace) = read_builder.new_scan().plan_with_trace().await.unwrap(); + let planned_ranges = merge_row_ranges( + plan.splits() + .iter() + .flat_map(|split| split.row_ranges().unwrap_or_default()) + .cloned() + .collect(), + ); + + assert_eq!( + plan.splits() + .iter() + .flat_map(|split| split.data_files()) + .count(), + 2 + ); + assert_eq!( + planned_ranges, + vec![RowRange::new(0, 0), RowRange::new(2, 3)] + ); + assert_eq!(trace.manifest_files_pruned_by_row_ranges, 0); + assert_eq!(trace.manifest_entries_read, 2); +} + +#[tokio::test] +async fn test_execute_writes_bitmap_index_manifest_and_java_file() { + let table_path = "memory:/test_bitmap_global_index_builder_e2e"; + let table = test_table_with_path(table_path, table_options("10")); + setup_dirs(&table).await; + + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "alice"])) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let shard_count = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .with_index_type(BITMAP_GLOBAL_INDEX_TYPE) + .execute() + .await + .unwrap(); + assert_eq!(shard_count, 1); + + let snapshot_manager = + SnapshotManager::new(table.file_io().clone(), table.location().to_string()); + let snapshot = snapshot_manager + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let index_manifest = snapshot.index_manifest().expect("index manifest"); + let index_entries = IndexManifest::read( + table.file_io(), + &format!("{table_path}/manifest/{index_manifest}"), + ) + .await + .unwrap(); + assert_eq!(index_entries.len(), 1); + + let index_file = &index_entries[0].index_file; + assert_eq!(index_file.index_type, BITMAP_GLOBAL_INDEX_TYPE); + assert!(index_file.file_name.starts_with("bitmap-global-index-")); + assert_eq!(index_file.row_count, 3); + assert!(index_file.file_size > 0); + + let global_meta = index_file + .global_index_meta + .as_ref() + .expect("global index meta"); + let bitmap_meta = + crate::btree::BTreeIndexMeta::deserialize(global_meta.index_meta.as_ref().unwrap()) + .unwrap(); + assert_eq!(bitmap_meta.first_key, Some(b"alice".to_vec())); + assert_eq!(bitmap_meta.last_key, Some(b"bob".to_vec())); + assert!(!bitmap_meta.has_nulls); + + let index_path = format!("{table_path}/index/{}", index_file.file_name); + let input = table.file_io().new_input(&index_path).unwrap(); + let file_size = input.metadata().await.unwrap().size; + let reader = input.reader().await.unwrap(); + let bitmap_reader = crate::table::bitmap_global_index_reader::BitmapGlobalIndexReader::open( + Box::new(reader), + file_size, + ) + .await + .unwrap(); + let bitmap = bitmap_reader + .query( + crate::spec::PredicateOperator::Eq, + &[crate::spec::Datum::String("alice".to_string())], + table.schema().fields()[1].data_type(), + ) + .await + .unwrap(); + assert_eq!(bitmap.iter().collect::>(), vec![0, 2]); + + let predicate = PredicateBuilder::new(table.schema().fields()) + .equal("name", crate::spec::Datum::String("alice".to_string())) + .unwrap(); + let row_ranges = evaluate_global_index(GlobalIndexEvaluation { + file_io: table.file_io(), + table_path: table.location(), + index_entries: &index_entries, + predicates: &[predicate], + schema_fields: table.schema().fields(), + search_mode: GlobalIndexSearchMode::Fast, + global_index_thread_num: 32, + btree_fallback_scan_max_size: i64::MAX, + bitmap_fallback_scan_max_size: i64::MAX, + fm_read_options: crate::fm_index::FMReadOptions::default(), + next_row_id: snapshot.next_row_id(), + data_ranges: &[], + }) + .await + .unwrap() + .unwrap(); + assert_eq!(row_ranges, vec![RowRange::new(0, 0), RowRange::new(2, 2)]); +} + +#[tokio::test] +async fn test_execute_multivalue_index_and_array_queries_end_to_end() { + let table_path = "memory:/test_multivalue_global_index_builder_e2e"; + let table = multivalue_table(table_path); + setup_dirs(&table).await; + + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&multivalue_batch()) + .await + .unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(table_write.prepare_commit().await.unwrap()) + .await + .unwrap(); + + let shard_count = table + .new_btree_global_index_build_builder() + .with_index_column("items") + .with_index_type(MULTIVALUE_GLOBAL_INDEX_TYPE) + .execute() + .await + .unwrap(); + assert_eq!(shard_count, 1); + + let snapshot = SnapshotManager::new(table.file_io().clone(), table.location().to_string()) + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let index_manifest = snapshot.index_manifest().expect("index manifest"); + let index_entries = IndexManifest::read( + table.file_io(), + &format!("{table_path}/manifest/{index_manifest}"), + ) + .await + .unwrap(); + assert_eq!(index_entries.len(), 1); + + let index_file = &index_entries[0].index_file; + assert_eq!(index_file.index_type, MULTIVALUE_GLOBAL_INDEX_TYPE); + assert_eq!(index_file.row_count, 5, "source rows, not postings"); + let global_meta = index_file.global_index_meta.as_ref().unwrap(); + let serialized_meta = global_meta.index_meta.as_deref().unwrap(); + let sorted_meta = BTreeIndexMeta::deserialize(serialized_meta).unwrap(); + assert_eq!(serialized_meta, sorted_meta.serialize()); + let element_type = DataType::Int(IntType::new()); + + let index_path = format!("{table_path}/index/{}", index_file.file_name); + let input = table.file_io().new_input(&index_path).unwrap(); + let file_size = input.metadata().await.unwrap().size; + let reader = input.reader().await.unwrap(); + let bitmap_reader = crate::table::bitmap_global_index_reader::BitmapGlobalIndexReader::open( + Box::new(reader), + file_size, + ) + .await + .unwrap(); + let contains = bitmap_reader + .query( + PredicateOperator::ArrayContains, + &[Datum::Int(10)], + &element_type, + ) + .await + .unwrap(); + assert_eq!(contains.iter().collect::>(), vec![0, 3]); + let overlap = bitmap_reader + .query( + PredicateOperator::ArraysOverlap, + &[Datum::Int(40), Datum::Int(10), Datum::Int(10)], + &element_type, + ) + .await + .unwrap(); + assert_eq!(overlap.iter().collect::>(), vec![0, 3, 4]); + let contains_all = bitmap_reader + .query( + PredicateOperator::ArrayContainsAll, + &[Datum::Int(10), Datum::Int(30), Datum::Int(10)], + &element_type, + ) + .await + .unwrap(); + assert_eq!(contains_all.iter().collect::>(), vec![3]); + + let fields = table.schema().fields(); + let contains_all_predicate = PredicateBuilder::new(fields) + .array_contains_all( + "items", + vec![Datum::Int(10), Datum::Int(30), Datum::Int(10)], + ) + .unwrap(); + let ranges = evaluate_global_index(GlobalIndexEvaluation { + file_io: table.file_io(), + table_path: table.location(), + index_entries: &index_entries, + predicates: std::slice::from_ref(&contains_all_predicate), + schema_fields: fields, + search_mode: GlobalIndexSearchMode::Fast, + global_index_thread_num: 32, + btree_fallback_scan_max_size: i64::MAX, + bitmap_fallback_scan_max_size: i64::MAX, + fm_read_options: crate::fm_index::FMReadOptions::default(), + next_row_id: snapshot.next_row_id(), + data_ranges: &[], + }) + .await + .unwrap() + .unwrap(); + assert_eq!(ranges, vec![RowRange::new(3, 3)]); + + assert_eq!(scan_ids(&table, contains_all_predicate).await, vec![4]); + let empty_contains_all = PredicateBuilder::new(fields) + .array_contains_all("items", vec![]) + .unwrap(); + assert_eq!( + scan_ids(&table, empty_contains_all).await, + vec![1, 2, 4, 5], + "empty contains-all matches every non-null array and must fall back" + ); +} + +#[tokio::test] +async fn test_bitmap_floating_candidates_preserve_residual_results() { + let table_path = "memory:/test_bitmap_floating_residual_candidates"; + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("f", DataType::Float(FloatType::new())) + .column("d", DataType::Double(DoubleType::new())) + .options(table_options("100")) + .build() + .unwrap(); + let table = Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "test_bitmap_floating_residual_candidates"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + setup_dirs(&table).await; + + let float_negative_nan = f32::from_bits(0xffc0_0001); + let double_negative_nan = f64::from_bits(0xfff8_0000_0000_0001); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("f", ArrowDataType::Float32, true), + ArrowField::new("d", ArrowDataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![ + Arc::new(Int32Array::from_iter_values(0..9)) as ArrayRef, + Arc::new(Float32Array::from(vec![ + Some(float_negative_nan), + Some(f32::from_bits(0xffff_1234)), + Some(f32::NAN), + Some(f32::from_bits(0x7fc0_0010)), + Some(-1.0), + Some(-0.0), + Some(0.0), + Some(1.0), + None, + ])) as ArrayRef, + Arc::new(Float64Array::from(vec![ + Some(double_negative_nan), + Some(f64::from_bits(0xffff_1234_5678_9abc)), + Some(f64::NAN), + Some(f64::from_bits(0x7ff8_0000_0000_0010)), + Some(-1.0), + Some(-0.0), + Some(0.0), + Some(1.0), + None, + ])) as ArrayRef, + ], + ) + .unwrap(); + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write.write_arrow_batch(&batch).await.unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + for column in ["f", "d"] { + let shard_count = table + .new_btree_global_index_build_builder() + .with_index_column(column) + .with_index_type(BITMAP_GLOBAL_INDEX_TYPE) + .execute() + .await + .unwrap(); + assert_eq!(shard_count, 1); + } + + let mut disabled_options = table.schema().options().clone(); + disabled_options.insert("global-index.enabled".to_string(), "false".to_string()); + let table_without_index = Table::new( + table.file_io().clone(), + table.identifier().clone(), + table.location().to_string(), + table.schema().copy_with_replaced_options(disabled_options), + None, + ); + + let predicates = PredicateBuilder::new(table.schema().fields()); + let cases = [ + ( + "Float < 0", + predicates.less_than("f", Datum::Float(0.0)).unwrap(), + vec![0, 1, 4, 5], + ), + ( + "Double < 0", + predicates.less_than("d", Datum::Double(0.0)).unwrap(), + vec![0, 1, 4, 5], + ), + ( + "Float = canonical NaN", + predicates.equal("f", Datum::Float(f32::NAN)).unwrap(), + vec![2], + ), + ( + "Double = canonical NaN", + predicates.equal("d", Datum::Double(f64::NAN)).unwrap(), + vec![2], + ), + ( + "Float = negative NaN", + predicates + .equal("f", Datum::Float(float_negative_nan)) + .unwrap(), + vec![0], + ), + ( + "Double = negative NaN", + predicates + .equal("d", Datum::Double(double_negative_nan)) + .unwrap(), + vec![0], + ), + ( + "Float IN NaNs", + predicates + .is_in( + "f", + vec![Datum::Float(float_negative_nan), Datum::Float(f32::NAN)], + ) + .unwrap(), + vec![0, 2], + ), + ( + "Double IN NaNs", + predicates + .is_in( + "d", + vec![Datum::Double(double_negative_nan), Datum::Double(f64::NAN)], + ) + .unwrap(), + vec![0, 2], + ), + ( + "Float != canonical NaN", + predicates.not_equal("f", Datum::Float(f32::NAN)).unwrap(), + vec![0, 1, 3, 4, 5, 6, 7], + ), + ( + "Double != canonical NaN", + predicates.not_equal("d", Datum::Double(f64::NAN)).unwrap(), + vec![0, 1, 3, 4, 5, 6, 7], + ), + ( + "Float NOT IN", + predicates + .is_not_in("f", vec![Datum::Float(f32::NAN), Datum::Float(0.0)]) + .unwrap(), + vec![0, 1, 3, 4, 5, 7], + ), + ( + "Double NOT IN", + predicates + .is_not_in("d", vec![Datum::Double(f64::NAN), Datum::Double(0.0)]) + .unwrap(), + vec![0, 1, 3, 4, 5, 7], + ), + ( + "Float combined range", + Predicate::and(vec![ + predicates + .greater_or_equal("f", Datum::Float(float_negative_nan)) + .unwrap(), + predicates.less_or_equal("f", Datum::Float(0.0)).unwrap(), + ]), + vec![0, 4, 5, 6], + ), + ( + "Double combined range", + Predicate::and(vec![ + predicates + .greater_or_equal("d", Datum::Double(double_negative_nan)) + .unwrap(), + predicates.less_or_equal("d", Datum::Double(0.0)).unwrap(), + ]), + vec![0, 4, 5, 6], + ), + ]; + + for (name, predicate, expected) in cases { + let without_index = scan_ids(&table_without_index, predicate.clone()).await; + assert_eq!(without_index, expected, "{name}: residual baseline"); + let with_index = scan_ids(&table, predicate).await; + assert_eq!( + with_index, without_index, + "{name}: global index changed rows" + ); + } +} + +/// Bitmap is built through the same sorted builder; a second build with no +/// new data must be a no-op keyed on the bitmap coverage — not error, and +/// not be confused by any btree coverage of the same field. +#[tokio::test] +async fn bitmap_second_build_without_new_data_is_noop() { + let table_path = "memory:/test_bitmap_global_index_second_build_noop"; + let table = test_table_with_path(table_path, table_options("10")); + setup_dirs(&table).await; + + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "carol"])) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let first_built = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .with_index_type(BITMAP_GLOBAL_INDEX_TYPE) + .execute() + .await + .unwrap(); + assert!( + first_built > 0, + "first bitmap build must index initial rows" + ); + + let files_after_first = latest_bitmap_index_files(&table).await; + assert!(!files_after_first.is_empty()); + + let built = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .with_index_type(BITMAP_GLOBAL_INDEX_TYPE) + .execute() + .await + .unwrap(); + assert_eq!( + built, 0, + "fully-indexed bitmap table must build nothing on re-run" + ); + + let names_first = files_after_first + .iter() + .map(|f| f.file_name.clone()) + .collect::>(); + let names_second = latest_bitmap_index_files(&table) + .await + .iter() + .map(|f| f.file_name.clone()) + .collect::>(); + assert_eq!(names_first, names_second, "re-run must not change entries"); +} + +/// A btree index over the SAME field must NOT count as bitmap coverage: a +/// bitmap build after a btree build over identical rows must still produce a +/// bitmap index (regression guard for the index_type-keyed gap computation — +/// the merge-residual bug hard-coded btree here, which would have skipped +/// these rows for a bitmap build). +#[tokio::test] +async fn bitmap_build_after_btree_on_same_field_still_indexes() { + let table_path = "memory:/test_bitmap_after_btree_same_field"; + let table = test_table_with_path(table_path, table_options("10")); + setup_dirs(&table).await; + + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "carol"])) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let btree_built = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .execute() + .await + .unwrap(); + assert!(btree_built > 0); + + let bitmap_built = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .with_index_type(BITMAP_GLOBAL_INDEX_TYPE) + .execute() + .await + .unwrap(); + assert!( + bitmap_built > 0, + "bitmap build must index rows even when a btree index already covers the same field" + ); + + let bitmap_files = latest_bitmap_index_files(&table).await; + assert!( + !bitmap_files.is_empty(), + "a bitmap index file must be written" + ); + let coverage = data_row_id_coverage(&table).await; + let bitmap_start = bitmap_files + .iter() + .filter_map(|f| f.global_index_meta.as_ref()) + .map(|m| m.row_range_start) + .min() + .unwrap(); + assert_eq!( + bitmap_start, + coverage[0].from(), + "bitmap coverage must span from the first data row, not skip btree-covered rows" + ); +} + +/// Bitmap incremental: build, append, build again → only the appended range +/// gets a new bitmap file; the first bitmap file is retained (append-only). +#[tokio::test] +async fn bitmap_incremental_build_indexes_only_new_rows() { + let table_path = "memory:/test_bitmap_global_index_incremental"; + let table = test_table_with_path(table_path, table_options("10")); + setup_dirs(&table).await; + + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "carol"])) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let first_built = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .with_index_type(BITMAP_GLOBAL_INDEX_TYPE) + .execute() + .await + .unwrap(); + assert!(first_built > 0); + let first_names = latest_bitmap_index_files(&table) + .await + .iter() + .map(|f| f.file_name.clone()) + .collect::>(); + let n: i64 = 3; + + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch(vec![4, 5, 6], vec!["dave", "erin", "frank"])) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let second_built = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .with_index_type(BITMAP_GLOBAL_INDEX_TYPE) + .execute() + .await + .unwrap(); + assert!(second_built > 0, "appended rows must be indexed"); + + let all_files = latest_bitmap_index_files(&table).await; + let all_names = all_files + .iter() + .map(|f| f.file_name.clone()) + .collect::>(); + assert!( + first_names.iter().all(|name| all_names.contains(name)), + "build #1 bitmap files must be retained untouched" + ); + let new_files = all_files + .iter() + .filter(|f| !first_names.contains(&f.file_name)) + .collect::>(); + assert!(!new_files.is_empty(), "build #2 must add new bitmap files"); + for file in new_files { + let meta = file + .global_index_meta + .as_ref() + .expect("global index meta on new bitmap file"); + assert!( + meta.row_range_start >= n, + "new bitmap file range must start at or after {}, got [{}, {}]", + n, + meta.row_range_start, + meta.row_range_end + ); + } +} + +async fn latest_bitmap_index_files(table: &Table) -> Vec { + let snapshot_manager = + SnapshotManager::new(table.file_io().clone(), table.location().to_string()); + let snapshot = snapshot_manager + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let Some(index_manifest_name) = snapshot.index_manifest() else { + return Vec::new(); + }; + IndexManifest::read( + table.file_io(), + &snapshot_manager.manifest_path(index_manifest_name), + ) + .await + .unwrap() + .into_iter() + .filter(|entry| { + entry.kind == FileKind::Add && entry.index_file.index_type == BITMAP_GLOBAL_INDEX_TYPE + }) + .map(|entry| entry.index_file) + .collect() +} + +async fn latest_btree_index_files(table: &Table) -> Vec { + let snapshot_manager = + SnapshotManager::new(table.file_io().clone(), table.location().to_string()); + let snapshot = snapshot_manager + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let Some(index_manifest_name) = snapshot.index_manifest() else { + return Vec::new(); + }; + IndexManifest::read( + table.file_io(), + &snapshot_manager.manifest_path(index_manifest_name), + ) + .await + .unwrap() + .into_iter() + .filter(|entry| { + entry.kind == FileKind::Add && entry.index_file.index_type == BTREE_GLOBAL_INDEX_TYPE + }) + .map(|entry| entry.index_file) + .collect() +} + +/// Row-id coverage of the committed data files, read back from the data +/// manifest (never hard-coded) and merged into contiguous ranges. Mirrors +/// how `execute` gathers `manifest_entries` so tests observe the exact +/// row-ids the writer assigned. +async fn data_row_id_coverage(table: &Table) -> Vec { + let snapshot_manager = + SnapshotManager::new(table.file_io().clone(), table.location().to_string()); + let snapshot = snapshot_manager + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let entries = table + .new_read_builder() + .new_scan() + .with_scan_all_files() + .plan_manifest_entries(&snapshot) + .await + .unwrap(); + let ranges = entries + .iter() + .filter(|entry| *entry.kind() == FileKind::Add) + .filter_map(|entry| { + entry + .file() + .row_id_range() + .map(|(start, end)| RowRange::new(start, end)) + }) + .collect::>(); + merge_row_ranges(ranges) +} + +/// Second build with no new data must be a clean no-op (returns 0), not an +/// overlap error. This is the core bug fix: today the second call errors. +#[tokio::test] +async fn second_build_without_new_data_is_noop() { + let table_path = "memory:/test_btree_global_index_second_build_noop"; + let table = test_table_with_path(table_path, table_options("10")); + setup_dirs(&table).await; + + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "carol"])) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let first_built = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .execute() + .await + .unwrap(); + assert!(first_built > 0, "first build must index the initial rows"); + + let files_after_first = latest_btree_index_files(&table).await; + assert!(!files_after_first.is_empty()); + + let built = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .execute() + .await + .unwrap(); + assert_eq!(built, 0, "fully-indexed table must build nothing on re-run"); + + let files_after_second = latest_btree_index_files(&table).await; + let names_first = files_after_first + .iter() + .map(|f| f.file_name.clone()) + .collect::>(); + let names_second = files_after_second + .iter() + .map(|f| f.file_name.clone()) + .collect::>(); + assert_eq!( + names_first, names_second, + "re-run must not add or remove index manifest entries" + ); +} + +/// Build, append new rows, build again -> only the appended row range is +/// indexed; the first build's index files are retained untouched (append-only). +#[tokio::test] +async fn incremental_build_indexes_only_new_rows() { + let table_path = "memory:/test_btree_global_index_incremental"; + let table = test_table_with_path(table_path, table_options("10")); + setup_dirs(&table).await; + + // Build #1 over rows [0..3). + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "carol"])) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let first_built = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .execute() + .await + .unwrap(); + assert!(first_built > 0); + + let first_files = latest_btree_index_files(&table).await; + let first_names = first_files + .iter() + .map(|f| f.file_name.clone()) + .collect::>(); + let n: i64 = 3; + + // Append a second batch (new row-ids [3..6)). + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch(vec![4, 5, 6], vec!["dave", "erin", "frank"])) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let second_built = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .execute() + .await + .unwrap(); + assert!(second_built > 0, "appended rows must be indexed"); + + let all_files = latest_btree_index_files(&table).await; + let all_names = all_files + .iter() + .map(|f| f.file_name.clone()) + .collect::>(); + + // Every build-#1 file is still present (append-only, no rewrite/delete). + assert!( + first_names.iter().all(|name| all_names.contains(name)), + "build #1 index files must be retained untouched" + ); + + // Every build-#2 file covers only the appended range [N, ..]. + let new_files = all_files + .iter() + .filter(|f| !first_names.contains(&f.file_name)) + .collect::>(); + assert!(!new_files.is_empty(), "build #2 must add new index files"); + for file in new_files { + let meta = file + .global_index_meta + .as_ref() + .expect("global index meta on new btree file"); + assert!( + meta.row_range_start >= n, + "new index file range must start at or after {}, got [{}, {}]", + n, + meta.row_range_start, + meta.row_range_end + ); + } +} + +/// Regression: first build (no existing index) must equal the pre-change +/// full build -- subtraction with empty `indexed` = full coverage. +#[tokio::test] +async fn first_build_indexes_full_coverage() { + let table_path = "memory:/test_btree_global_index_first_full_coverage"; + let table = test_table_with_path(table_path, table_options("10")); + setup_dirs(&table).await; + + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "carol"])) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let built = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .execute() + .await + .unwrap(); + assert_eq!( + built, 1, + "first build must index the full coverage in one shard" + ); + + let files = latest_btree_index_files(&table).await; + assert_eq!(files.len(), 1); + let meta = files[0] + .global_index_meta + .as_ref() + .expect("global index meta"); + assert_eq!(meta.row_range_start, 0); + assert_eq!(meta.row_range_end, 2); +} + +/// Grid boundary (spec edge 4): with `records-per-range = 4`, an appended +/// gap that spans several grid cells must be split so each new index file's +/// range stays inside one cell, the ranges are contiguous, and together +/// they exactly cover the gap. Row-ids are read back from the manifests, +/// never hard-coded. +#[tokio::test] +async fn incremental_build_splits_gap_across_records_per_range_grid() { + const RPR: i64 = 4; + let table_path = "memory:/test_btree_global_index_grid_boundary"; + let table = test_table_with_path(table_path, table_options("4")); + setup_dirs(&table).await; + + // Build #1 over an initial batch (row-ids the writer assigns). + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch(vec![1, 2, 3], vec!["alice", "bob", "carol"])) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let first_built = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .execute() + .await + .unwrap(); + assert!(first_built > 0, "first build must index the initial rows"); + + // Row range already covered by build #1 (read back, not hard-coded). + let first_index_files = latest_btree_index_files(&table).await; + let indexed_before = merge_row_ranges( + first_index_files + .iter() + .filter_map(|f| f.global_index_meta.as_ref()) + .map(|m| RowRange::new(m.row_range_start, m.row_range_end)) + .collect(), + ); + assert_eq!( + indexed_before.len(), + 1, + "build #1 should cover one contiguous range" + ); + let gap_start = indexed_before[0].to() + 1; + let before_names = first_index_files + .iter() + .map(|f| f.file_name.clone()) + .collect::>(); + + // Append rows so the new gap crosses records_per_range (=4) boundaries. + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch( + vec![4, 5, 6, 7, 8, 9, 10], + vec!["d", "e", "f", "g", "h", "i", "j"], + )) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + // Total data coverage read back from the data manifest. + let coverage = data_row_id_coverage(&table).await; + assert_eq!( + coverage.len(), + 1, + "appended data must be contiguous with build #1" + ); + let gap_end = coverage[0].to(); + assert!( + gap_end - gap_start + 1 > RPR, + "gap [{gap_start}, {gap_end}] must span more than one records_per_range cell" + ); + + let second_built = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .execute() + .await + .unwrap(); + assert!(second_built > 0, "appended rows must be indexed"); + + // Only the newly written index files (build #1 files are retained). + let mut new_metas = latest_btree_index_files(&table) + .await + .into_iter() + .filter(|f| !before_names.contains(&f.file_name)) + .filter_map(|f| f.global_index_meta) + .map(|m| (m.row_range_start, m.row_range_end)) + .collect::>(); + new_metas.sort(); + assert!(!new_metas.is_empty(), "build #2 must add new index files"); + + // (a) Each range lies within a single grid cell: no multiple of RPR is + // strictly interior, i.e. start and end share the same cell index. + for (start, end) in &new_metas { + assert!(end >= start, "range must be non-empty: [{start}, {end}]"); + assert_eq!( + start / RPR, + end / RPR, + "range [{start}, {end}] straddles a records_per_range boundary" + ); + } + // (b) Contiguous with no gaps or overlaps. + for pair in new_metas.windows(2) { + assert_eq!( + pair[1].0, + pair[0].1 + 1, + "ranges must be contiguous: {:?} then {:?}", + pair[0], + pair[1] + ); + } + // (c) Together they exactly cover the appended gap [gap_start, gap_end]. + assert_eq!( + new_metas.first().unwrap().0, + gap_start, + "coverage must start at the gap start" + ); + assert_eq!( + new_metas.last().unwrap().1, + gap_end, + "coverage must end at the gap end" + ); +} + +/// Hole splitting (spec edge 5) at build level: a mid-coverage indexed range +/// (constructed directly, as the drop-builder tests build `GlobalIndexMeta` +/// entries) must carve the data coverage into two build segments, one on +/// each side, and the hole itself must not be re-indexed. +#[tokio::test] +async fn incremental_build_splits_gap_around_mid_coverage_indexed_hole() { + let table_path = "memory:/test_btree_global_index_mid_hole"; + // records-per-range large so the grid never splits: the only split is + // the hole itself. + let table = test_table_with_path(table_path, table_options("100")); + setup_dirs(&table).await; + + // Real data spanning row-ids [0, 9]. + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&data_batch( + (1..=10).collect(), + vec!["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"], + )) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let coverage = data_row_id_coverage(&table).await; + assert_eq!(coverage.len(), 1, "data must be one contiguous range"); + assert_eq!(coverage[0].from(), 0); + let last_row = coverage[0].to(); + assert!(last_row >= 9, "need at least 10 rows for a mid hole"); + + // Inject a mid-coverage indexed range [hole_start, hole_end] for the + // `name` field directly into the index manifest. + let name_field_id = find_index_field(&table, "name").unwrap().id(); + let hole_start = 4; + let hole_end = 6; + let synthetic = IndexFileMeta { + index_type: BTREE_GLOBAL_INDEX_TYPE.to_string(), + file_name: "btree-synthetic-hole.index".to_string(), + file_size: 1, + row_count: (hole_end - hole_start + 1), + deletion_vectors_ranges: None, + global_index_meta: Some(GlobalIndexMeta { + row_range_start: hole_start, + row_range_end: hole_end, + index_field_id: name_field_id, + extra_field_ids: None, + source_meta: None, + index_meta: None, + }), + }; + let mut message = CommitMessage::new(BinaryRow::new(0).to_serialized_bytes(), 0, vec![]); + message.new_index_files = vec![synthetic]; + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(vec![message]) + .await + .unwrap(); + + let before_names = latest_btree_index_files(&table) + .await + .into_iter() + .map(|f| f.file_name) + .collect::>(); + + // Build: gap = coverage minus the hole = [0, hole_start-1] and + // [hole_end+1, last_row]; two shards since the grid does not split here. + let built = table + .new_btree_global_index_build_builder() + .with_index_column("name") + .execute() + .await + .unwrap(); + assert_eq!( + built, 2, + "mid-coverage hole must split the gap into two shards" + ); + + let mut new_metas = latest_btree_index_files(&table) + .await + .into_iter() + .filter(|f| !before_names.contains(&f.file_name)) + .filter_map(|f| f.global_index_meta) + .map(|m| (m.row_range_start, m.row_range_end)) + .collect::>(); + new_metas.sort(); + + assert_eq!( + new_metas, + vec![(0, hole_start - 1), (hole_end + 1, last_row)], + "new shards must fill the coverage on both sides of the indexed hole" + ); + for (start, end) in &new_metas { + assert!( + *end < hole_start || *start > hole_end, + "new shard [{start}, {end}] must not overlap indexed hole [{hole_start}, {hole_end}]" + ); + } +} diff --git a/crates/paimon/src/table/sorted_global_index_build_builder/validation.rs b/crates/paimon/src/table/sorted_global_index_build_builder/validation.rs new file mode 100644 index 000000000..59f9d001f --- /dev/null +++ b/crates/paimon/src/table/sorted_global_index_build_builder/validation.rs @@ -0,0 +1,179 @@ +// 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-option, field-type, and row-range validation. + +use crate::spec::{CoreOptions, DataField, DataType}; +use crate::table::global_index_types::{FM_GLOBAL_INDEX_TYPE, MULTIVALUE_GLOBAL_INDEX_TYPE}; +use crate::table::Table; +use crate::{Error, Result}; + +pub(super) fn validate_table_options(table: &Table, core_options: &CoreOptions) -> Result<()> { + if !core_options.row_tracking_enabled() { + return Err(Error::DataInvalid { + message: "Sorted global index build requires 'row-tracking.enabled' = 'true'" + .to_string(), + source: None, + }); + } + if !core_options.data_evolution_enabled() { + return Err(Error::DataInvalid { + message: "Sorted global index build requires 'data-evolution.enabled' = 'true'" + .to_string(), + source: None, + }); + } + if !core_options.global_index_enabled() { + return Err(Error::DataInvalid { + message: "Sorted global index build requires 'global-index.enabled' = 'true'" + .to_string(), + source: None, + }); + } + if !table.schema().primary_keys().is_empty() { + return Err(Error::Unsupported { + message: "Sorted global index build does not support primary-key tables".to_string(), + }); + } + if core_options.deletion_vectors_enabled() { + return Err(Error::Unsupported { + message: + "Sorted global index build does not support tables with deletion-vectors.enabled=true" + .to_string(), + }); + } + Ok(()) +} + +pub(super) fn find_index_field<'a>(table: &'a Table, column: &str) -> Result<&'a DataField> { + table + .schema() + .fields() + .iter() + .find(|field| field.name() == column) + .ok_or_else(|| Error::ColumnNotExist { + full_name: table.identifier().full_name(), + column: column.to_string(), + }) +} + +pub(super) fn validate_btree_field(field: &DataField) -> Result<()> { + if !is_btree_supported_data_type(field.data_type()) { + return Err(Error::Unsupported { + message: format!( + "Sorted global index only supports scalar columns, got {:?} for column '{}'", + field.data_type(), + field.name() + ), + }); + } + Ok(()) +} + +pub(super) fn index_key_type<'a>(index_type: &str, field: &'a DataField) -> Result<&'a DataType> { + if index_type == MULTIVALUE_GLOBAL_INDEX_TYPE { + let DataType::Array(array_type) = field.data_type() else { + return Err(Error::Unsupported { + message: format!( + "Multivalue global index requires an ARRAY column, got {:?} for column '{}'", + field.data_type(), + field.name() + ), + }); + }; + if !is_btree_supported_data_type(array_type.element_type()) { + return Err(Error::Unsupported { + message: format!( + "Multivalue global index does not support array element type {:?} for column '{}'", + array_type.element_type(), + field.name() + ), + }); + } + Ok(array_type.element_type()) + } else if index_type == FM_GLOBAL_INDEX_TYPE { + if !matches!(field.data_type(), DataType::Char(_) | DataType::VarChar(_)) { + return Err(Error::Unsupported { + message: format!( + "FM global index requires a character string column, got {:?} for column '{}'", + field.data_type(), + field.name() + ), + }); + } + Ok(field.data_type()) + } else { + validate_btree_field(field)?; + Ok(field.data_type()) + } +} + +pub(super) fn is_btree_supported_data_type(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Boolean(_) + | DataType::TinyInt(_) + | DataType::SmallInt(_) + | DataType::Int(_) + | DataType::BigInt(_) + | DataType::Decimal(_) + | DataType::Double(_) + | DataType::Float(_) + | DataType::Char(_) + | DataType::VarChar(_) + | DataType::Date(_) + | DataType::LocalZonedTimestamp(_) + | DataType::Time(_) + | DataType::Timestamp(_) + ) +} + +pub(super) fn checked_i64(value: u64, context: &str) -> Result { + i64::try_from(value).map_err(|_| Error::DataInvalid { + message: format!("{context}: {value}"), + source: None, + }) +} + +pub(super) fn checked_row_count(row_range_start: i64, row_range_end: i64) -> Result { + if row_range_end < row_range_start { + return Err(Error::DataInvalid { + message: format!( + "Invalid sorted global index row range [{row_range_start}, {row_range_end}]" + ), + source: None, + }); + } + row_range_end + .checked_sub(row_range_start) + .and_then(|span| span.checked_add(1)) + .ok_or_else(|| Error::DataInvalid { + message: format!( + "Row count overflows for row range [{row_range_start}, {row_range_end}]" + ), + source: None, + }) +} + +pub(super) fn ranges_overlap( + left_start: i64, + left_end: i64, + right_start: i64, + right_end: i64, +) -> bool { + left_start <= right_end && right_start <= left_end +} diff --git a/crates/paimon/src/table/sorted_global_index_build_builder/writer.rs b/crates/paimon/src/table/sorted_global_index_build_builder/writer.rs new file mode 100644 index 000000000..1c1ecb3e7 --- /dev/null +++ b/crates/paimon/src/table/sorted_global_index_build_builder/writer.rs @@ -0,0 +1,365 @@ +// 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. + +//! Type-specific index-file writing and failure cleanup. + +use super::extraction::{build_read_splits_for_shard, extract_index_rows, sort_index_rows}; +use super::planning::SortedGlobalIndexShard; +use super::validation::{checked_i64, checked_row_count, index_key_type}; +use super::{ + make_index_key_codec, GlobalIndexWriteOptions, SerializeKeyFn, SortedGlobalIndexBuildBuilder, + INDEX_DIR, +}; +use crate::btree::BTreeIndexWriter; +use crate::fm_index::{FMGlobalIndexWriter, FMWriteOptions}; +use crate::io::FileWrite; +use crate::spec::{ + extract_datum_from_arrow, DataField, GlobalIndexMeta, IndexFileMeta, ROW_ID_FIELD_NAME, +}; +use crate::table::bitmap_global_index_writer::{BitmapGlobalIndexWriter, BitmapWriteResult}; +use crate::table::global_index_types::{ + normalize_queryable_global_index_type, BITMAP_GLOBAL_INDEX_TYPE, BTREE_GLOBAL_INDEX_TYPE, + FM_GLOBAL_INDEX_TYPE, MULTIVALUE_GLOBAL_INDEX_TYPE, +}; +use crate::table::Table; +use crate::{Error, Result}; +use arrow_array::{Array, Int64Array}; +use futures::TryStreamExt; + +impl SortedGlobalIndexBuildBuilder<'_> { + pub(super) async fn build_index_file( + &self, + shard: &SortedGlobalIndexShard, + index_field: &DataField, + index_column: &str, + write_options: &GlobalIndexWriteOptions, + ) -> Result { + let index_type = normalize_queryable_global_index_type(&self.index_type).ok_or_else(|| { + Error::Unsupported { + message: format!( + "Scalar global index build only supports index_type => 'btree', 'bitmap', 'multivalue', or 'fm', got '{}'", + self.index_type + ), + } + })?; + let row_count = checked_row_count(shard.row_range_start, shard.row_range_end)?; + let key_type = index_key_type(index_type, index_field)?; + let codec_type = if matches!( + index_type, + MULTIVALUE_GLOBAL_INDEX_TYPE | FM_GLOBAL_INDEX_TYPE + ) { + BITMAP_GLOBAL_INDEX_TYPE + } else { + index_type + }; + let (cmp, serialize_key) = make_index_key_codec(codec_type, key_type); + let mut rows = if index_type == FM_GLOBAL_INDEX_TYPE { + Vec::new() + } else { + extract_index_rows( + self.table, + shard, + index_column, + index_field, + index_type, + serialize_key, + ) + .await? + }; + if !rows.is_empty() { + sort_index_rows(&mut rows, &cmp); + } + + self.table + .file_io() + .mkdirs(&format!( + "{}/{INDEX_DIR}/", + self.table.location().trim_end_matches('/') + )) + .await?; + let file_name = format!("{index_type}-global-index-{}.index", uuid::Uuid::new_v4()); + let index_path = format!( + "{}/{INDEX_DIR}/{}", + self.table.location().trim_end_matches('/'), + file_name + ); + let write_result: Result<(u64, Vec, i64)> = async { + let output = self.table.file_io().new_output(&index_path)?; + let writer = output.writer().await?; + let (written_row_count, index_meta) = match index_type { + BTREE_GLOBAL_INDEX_TYPE => { + let GlobalIndexWriteOptions::Sorted(write_options) = write_options else { + unreachable!("BTree uses sorted write options") + }; + let mut writer = BTreeIndexWriter::with_comparator_and_compression_level( + writer, + write_options.block_size, + write_options.compression_type, + write_options.compression_level, + cmp, + ); + for (key, local_row_id) in &rows { + writer + .write(key.as_deref(), *local_row_id) + .await + .map_err(|e| Error::DataInvalid { + message: format!( + "Failed to write BTree global index file '{file_name}'" + ), + source: Some(Box::new(e)), + })?; + } + let write_result = writer.finish().await.map_err(|e| Error::DataInvalid { + message: format!("Failed to finish BTree global index file '{file_name}'"), + source: Some(Box::new(e)), + })?; + (write_result.row_count, write_result.meta.serialize()) + } + BITMAP_GLOBAL_INDEX_TYPE => { + let GlobalIndexWriteOptions::Sorted(write_options) = write_options else { + unreachable!("bitmap uses sorted write options") + }; + let mut writer = BitmapGlobalIndexWriter::with_compression_level( + writer, + write_options.block_size, + write_options.compression_type, + write_options.compression_level, + cmp, + ); + for (key, local_row_id) in &rows { + writer.write(key.as_deref(), *local_row_id).map_err(|e| { + Error::DataInvalid { + message: format!( + "Failed to write bitmap global index file '{file_name}'" + ), + source: Some(Box::new(e)), + } + })?; + } + let BitmapWriteResult { row_count, meta } = + writer.finish().await.map_err(|e| Error::DataInvalid { + message: format!( + "Failed to finish bitmap global index file '{file_name}'" + ), + source: Some(Box::new(e)), + })?; + (row_count, meta.serialize()) + } + MULTIVALUE_GLOBAL_INDEX_TYPE => { + let GlobalIndexWriteOptions::Sorted(write_options) = write_options else { + unreachable!("multivalue uses sorted write options") + }; + let mut writer = BitmapGlobalIndexWriter::with_compression_level( + writer, + write_options.block_size, + write_options.compression_type, + write_options.compression_level, + cmp, + ); + for (key, local_row_id) in &rows { + let key = key + .as_deref() + .expect("multivalue extraction skips null keys"); + writer.write_posting(key, *local_row_id).map_err(|e| { + Error::DataInvalid { + message: format!( + "Failed to write multivalue global index file '{file_name}'" + ), + source: Some(Box::new(e)), + } + })?; + } + let BitmapWriteResult { row_count, meta } = writer + .finish_with_source_row_count(u64::try_from(row_count).unwrap()) + .await + .map_err(|e| Error::DataInvalid { + message: format!( + "Failed to finish multivalue global index file '{file_name}'" + ), + source: Some(Box::new(e)), + })?; + (row_count, meta.serialize()) + } + FM_GLOBAL_INDEX_TYPE => { + let GlobalIndexWriteOptions::FM(write_options) = write_options else { + unreachable!("FM uses FM write options") + }; + write_fm_index_streaming( + self.table, + shard, + index_column, + index_field, + serialize_key, + writer, + *write_options, + &file_name, + ) + .await? + } + _ => unreachable!("normalized queryable global index type"), + }; + + if written_row_count != u64::try_from(row_count).unwrap() { + return Err(Error::DataInvalid { + message: format!( + "Sorted global index expected {} rows, wrote {}", + row_count, written_row_count + ), + source: None, + }); + } + + let status = self.table.file_io().get_status(&index_path).await?; + let file_size = checked_i64( + status.size, + "Index file is too large for Rust IndexFileMeta", + )?; + Ok((written_row_count, index_meta, file_size)) + } + .await; + let (_, index_meta, file_size) = match write_result { + Ok(result) => result, + Err(error) => { + let _ = self.table.file_io().delete_file(&index_path).await; + return Err(error); + } + }; + Ok(IndexFileMeta { + index_type: index_type.to_string(), + file_name, + file_size, + row_count, + deletion_vectors_ranges: None, + global_index_meta: Some(GlobalIndexMeta { + row_range_start: shard.row_range_start, + row_range_end: shard.row_range_end, + index_field_id: index_field.id(), + extra_field_ids: None, + source_meta: None, + index_meta: Some(index_meta), + }), + }) + } +} + +#[allow(clippy::too_many_arguments)] +async fn write_fm_index_streaming( + table: &Table, + shard: &SortedGlobalIndexShard, + index_column: &str, + index_field: &DataField, + serialize_key: SerializeKeyFn, + output: Box, + write_options: FMWriteOptions, + file_name: &str, +) -> Result<(u64, Vec)> { + let mut writer = + FMGlobalIndexWriter::new(output, write_options).map_err(|error| Error::DataInvalid { + message: format!("Failed to create FM global index file '{file_name}'"), + source: Some(Box::new(error)), + })?; + let splits = build_read_splits_for_shard(shard)?; + let mut read_builder = table.new_read_builder(); + read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; + let read = read_builder.new_read()?; + let mut batches = read.to_arrow(&splits)?; + let expected_row_count = checked_row_count(shard.row_range_start, shard.row_range_end)?; + let mut expected_row_id = shard.row_range_start; + + while let Some(batch) = batches.try_next().await? { + let value_index = + batch + .schema() + .index_of(index_column) + .map_err(|error| Error::DataInvalid { + message: format!( + "Index column '{index_column}' not found in FM read batch: {error}" + ), + source: None, + })?; + let row_id_index = batch + .schema() + .index_of(ROW_ID_FIELD_NAME) + .map_err(|error| Error::DataInvalid { + message: format!("_ROW_ID column not found in FM read batch: {error}"), + source: None, + })?; + let row_ids = batch + .column(row_id_index) + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::DataInvalid { + message: "FM global index build requires non-null Int64 _ROW_ID".to_string(), + source: None, + })?; + + for row in 0..batch.num_rows() { + if row_ids.is_null(row) { + return Err(Error::DataInvalid { + message: "FM global index build found null _ROW_ID".to_string(), + source: None, + }); + } + let row_id = row_ids.value(row); + if row_id != expected_row_id { + return Err(Error::DataInvalid { + message: format!( + "FM global index build expected _ROW_ID {expected_row_id}, got {row_id}" + ), + source: None, + }); + } + expected_row_id = expected_row_id + .checked_add(1) + .ok_or_else(|| Error::DataInvalid { + message: "FM global index row ID overflow".to_string(), + source: None, + })?; + let local_row_id = + u64::try_from(row_id - shard.row_range_start).map_err(|_| Error::DataInvalid { + message: format!( + "FM global index file '{file_name}' has a negative local row ID" + ), + source: None, + })?; + let key = extract_datum_from_arrow(&batch, row, value_index, index_field.data_type())? + .map(|datum| serialize_key(&datum, index_field.data_type())); + writer + .write(key.as_deref(), local_row_id) + .await + .map_err(|error| Error::DataInvalid { + message: format!("Failed to write FM global index file '{file_name}'"), + source: Some(Box::new(error)), + })?; + } + } + + let actual_row_count = expected_row_id - shard.row_range_start; + if actual_row_count != expected_row_count { + return Err(Error::DataInvalid { + message: format!( + "FM global index build expected {expected_row_count} rows, got {actual_row_count}" + ), + source: None, + }); + } + let result = writer.finish().await.map_err(|error| Error::DataInvalid { + message: format!("Failed to finish FM global index file '{file_name}'"), + source: Some(Box::new(error)), + })?; + Ok((result.row_count, result.index_meta)) +} diff --git a/crates/paimon/src/table/vindex_index_build_builder.rs b/crates/paimon/src/table/vindex_index_build_builder.rs index b0680fe73..9e06ff9a8 100644 --- a/crates/paimon/src/table/vindex_index_build_builder.rs +++ b/crates/paimon/src/table/vindex_index_build_builder.rs @@ -15,120 +15,22 @@ // specific language governing permissions and limitations // under the License. -use crate::spec::{ - bucket_dir_name, BinaryRow, CoreOptions, DataField, DataFileMeta, DataType, FileKind, - GlobalIndexMeta, IndexFileMeta, ROW_ID_FIELD_NAME, -}; -use crate::table::data_file_reader::DataFileReadTiming; -use crate::table::source::exclude_row_ranges; -use crate::table::table_read::configured_parquet_read_budget; -use crate::table::{ - CommitMessage, DataSplit, DataSplitBuilder, RowRange, SnapshotManager, Table, TableCommit, -}; +mod extraction; +mod planning; +mod timing; +mod validation; +mod writer; + +use planning::plan_vindex_shards; +use timing::vector_index_build_timing_enabled; +use validation::{checked_i32, find_index_field, validate_table_options, validate_vector_field}; + +use crate::spec::CoreOptions; +use crate::table::{CommitMessage, RowRange, SnapshotManager, Table, TableCommit}; use crate::vindex::{is_vindex_index_type, VindexVectorIndexOptions}; use crate::{Error, Result}; -use arrow_array::{Array, FixedSizeListArray, Float32Array, Int64Array, ListArray, RecordBatch}; -use arrow_buffer::MutableBuffer; -use futures::TryStreamExt; -use paimon_vindex_core::autotune::default_training_vector_count; -use paimon_vindex_core::index::{VectorIndexTrainer, VectorIndexWriter}; -use paimon_vindex_core::io::PosWriter; use std::collections::HashMap; -use std::io::{Read, Seek, SeekFrom}; -use std::sync::{Arc, OnceLock}; -use std::time::{Duration, Instant}; -use tokio::io::AsyncWriteExt; -use tokio_util::io::SyncIoBridge; - -const INDEX_DIR: &str = "index"; -const VECTOR_BUFFER_BYTES: usize = 8 * 1024 * 1024; -const VECTOR_INDEX_BUILD_TIMING_ENV: &str = "PAIMON_LOG_VECTOR_INDEX_BUILD_TIMING"; - -fn vector_index_build_timing_enabled() -> bool { - static ENABLED: OnceLock = OnceLock::new(); - *ENABLED.get_or_init(|| { - std::env::var_os(VECTOR_INDEX_BUILD_TIMING_ENV).is_some_and(|value| value == "1") - }) -} - -struct VectorIndexBuildTiming { - total_without_commit: Duration, - source_batch_wait: Duration, - oss_read: Duration, - parquet_decode: Duration, - file_schema_open: Duration, - first_batch_wait: Duration, - remaining_batch_wait: Duration, - parquet_row_group_count: u64, - parquet_projected_bytes_min: u64, - parquet_projected_bytes_max: u64, - parquet_projected_bytes_total: u64, - parquet_peak_inflight_row_groups: usize, - raw_temp_write: Duration, - train_finish: Duration, - raw_temp_reread: Duration, - index_add: Duration, - serialize_upload: Duration, - rows: usize, - training_rows_seen: usize, - training_rows_retained: usize, - batch_count: usize, - raw_temp_bytes: usize, - index_bytes: u64, - data_file_count: usize, - file_name: String, -} - -impl VectorIndexBuildTiming { - fn log(self, index_type: &str, commit: Duration) { - let total = self.total_without_commit.saturating_add(commit); - let accounted = self - .source_batch_wait - .saturating_add(self.raw_temp_write) - .saturating_add(self.train_finish) - .saturating_add(self.raw_temp_reread) - .saturating_add(self.index_add) - .saturating_add(self.serialize_upload) - .saturating_add(commit); - let unattributed = total.saturating_sub(accounted); - eprintln!( - "event=paimon_vector_index_build index_type={} file={} rows={} training_rows_seen={} training_rows_retained={} batch_count={} raw_temp_bytes={} index_bytes={} source_batch_wait_ms={:.3} oss_read_ms={:.3} parquet_decode_ms={:.3} file_schema_open_ms={:.3} first_batch_wait_ms={:.3} remaining_batch_wait_ms={:.3} parquet_row_group_count={} parquet_projected_bytes_min={} parquet_projected_bytes_max={} parquet_projected_bytes_total={} parquet_peak_inflight_row_groups={} raw_temp_write_ms={:.3} train_finish_ms={:.3} raw_temp_reread_ms={:.3} index_add_ms={:.3} serialize_upload_ms={:.3} commit_ms={:.3} sample_read_ms=0.000 full_scan_add_ms=0.000 pipeline_blocked_ms=0.000 producer_blocked_ms=0.000 consumer_add_ms=0.000 data_file_count={} data_file_read_concurrency=1 peak_ready_batches=0 total_ms={:.3} unattributed_ms={:.3}", - index_type, - self.file_name, - self.rows, - self.training_rows_seen, - self.training_rows_retained, - self.batch_count, - self.raw_temp_bytes, - self.index_bytes, - self.source_batch_wait.as_secs_f64() * 1000.0, - self.oss_read.as_secs_f64() * 1000.0, - self.parquet_decode.as_secs_f64() * 1000.0, - self.file_schema_open.as_secs_f64() * 1000.0, - self.first_batch_wait.as_secs_f64() * 1000.0, - self.remaining_batch_wait.as_secs_f64() * 1000.0, - self.parquet_row_group_count, - self.parquet_projected_bytes_min, - self.parquet_projected_bytes_max, - self.parquet_projected_bytes_total, - self.parquet_peak_inflight_row_groups, - self.raw_temp_write.as_secs_f64() * 1000.0, - self.train_finish.as_secs_f64() * 1000.0, - self.raw_temp_reread.as_secs_f64() * 1000.0, - self.index_add.as_secs_f64() * 1000.0, - self.serialize_upload.as_secs_f64() * 1000.0, - commit.as_secs_f64() * 1000.0, - self.data_file_count, - total.as_secs_f64() * 1000.0, - unattributed.as_secs_f64() * 1000.0, - ); - } -} - -struct BuiltIndexFile { - meta: IndexFileMeta, - timing: Option, -} +use std::time::Instant; pub struct VindexIndexBuildBuilder<'a> { table: &'a Table, @@ -304,1757 +206,7 @@ impl<'a> VindexIndexBuildBuilder<'a> { Ok(shard_count) } - - async fn build_index_file( - &self, - shard: &VindexIndexShard, - index_column: &str, - dimension: i32, - index_field_id: i32, - options: &VindexVectorIndexOptions, - index_meta: Vec, - ) -> Result { - let timing_enabled = vector_index_build_timing_enabled(); - let total_start = timing_enabled.then(Instant::now); - let mut source_batch_wait = Duration::ZERO; - let mut raw_temp_write = Duration::ZERO; - let read_timing = timing_enabled.then(|| Arc::new(DataFileReadTiming::default())); - let parquet_read_budget = if timing_enabled { - let budget = configured_parquet_read_budget(self.table)?; - budget.enable_diagnostics(); - Some(budget) - } else { - None - }; - let mut batch_count = 0usize; - let row_count = checked_row_count(shard.row_range_start, shard.row_range_end)?; - let row_count_usize = usize::try_from(row_count).map_err(|e| Error::DataInvalid { - message: format!("Invalid vindex row count: {row_count}"), - source: Some(Box::new(e)), - })?; - let dimension_usize = usize::try_from(dimension).map_err(|e| Error::DataInvalid { - message: format!("Invalid vindex dimension: {dimension}"), - source: Some(Box::new(e)), - })?; - if dimension_usize == 0 { - return Err(Error::DataInvalid { - message: "vindex vector dimension must be positive".to_string(), - source: None, - }); - } - let expected_bytes = checked_vector_bytes(row_count_usize, dimension_usize)?; - let training_vector_count = - checked_training_vector_count(row_count_usize, options.train_sample_ratio)?; - let training_buffer_rows = - (VECTOR_BUFFER_BYTES / checked_vector_bytes(1, dimension_usize)?).max(1); - let training_buffer_floats = training_buffer_rows - .checked_mul(dimension_usize) - .ok_or_else(|| Error::DataInvalid { - message: "vindex training buffer length overflows usize".to_string(), - source: None, - })?; - - let mut trainer = - VectorIndexTrainer::new(options.config.clone()).map_err(|e| Error::DataInvalid { - message: format!("Failed to initialize vindex trainer: {e}"), - source: Some(Box::new(e)), - })?; - let raw_file = tempfile::tempfile().map_err(|e| Error::UnexpectedError { - message: format!("Failed to create temporary vindex vector file: {e}"), - source: Some(Box::new(e)), - })?; - let mut raw_file = tokio::fs::File::from_std(raw_file); - let split = data_split_for_shard(shard)?; - let mut read_builder = self.table.new_read_builder(); - read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; - let read = read_builder.new_read()?; - let read = match read_timing.as_ref() { - Some(timing) => read.with_data_file_read_timing(Arc::clone(timing)), - None => read, - }; - let read = match parquet_read_budget.as_ref() { - Some(budget) => read.with_parquet_read_budget(Arc::clone(budget)), - None => read, - }; - let mut batches = read.to_arrow(&[split])?; - let mut expected_row_id = shard.row_range_start; - let mut rows_seen = 0usize; - let mut bytes_written = 0usize; - let mut next_training_sample = 0usize; - let mut training_buffer = Vec::with_capacity(training_buffer_floats); - - loop { - let source_start = timing_enabled.then(Instant::now); - let batch = batches.try_next().await?; - if let Some(source_start) = source_start { - source_batch_wait = source_batch_wait.saturating_add(source_start.elapsed()); - } - let Some(batch) = batch else { break }; - batch_count += 1; - let vectors = - validate_vector_batch(&batch, index_column, dimension_usize, &mut expected_row_id)?; - let batch_end = - rows_seen - .checked_add(vectors.row_count) - .ok_or_else(|| Error::DataInvalid { - message: "vindex streamed row count overflows usize".to_string(), - source: None, - })?; - - if training_vector_count == row_count_usize { - trainer - .add_training_vectors_mut(vectors.values, vectors.row_count) - .map_err(|e| Error::DataInvalid { - message: format!("Failed to add vindex training vectors: {e}"), - source: Some(Box::new(e)), - })?; - } else { - while next_training_sample < training_vector_count { - let sample_row = checked_training_sample_index( - next_training_sample, - row_count_usize, - training_vector_count, - )?; - if sample_row >= batch_end { - break; - } - let start = (sample_row - rows_seen) * dimension_usize; - training_buffer - .extend_from_slice(&vectors.values[start..start + dimension_usize]); - next_training_sample += 1; - if training_buffer.len() == training_buffer_floats { - trainer - .add_training_vectors_mut( - &training_buffer, - training_buffer.len() / dimension_usize, - ) - .map_err(|e| Error::DataInvalid { - message: format!("Failed to add vindex training vectors: {e}"), - source: Some(Box::new(e)), - })?; - training_buffer.clear(); - } - } - } - - let raw_write_start = timing_enabled.then(Instant::now); - raw_file - .write_all(vectors.bytes) - .await - .map_err(|e| Error::UnexpectedError { - message: format!("Failed to spill vindex vectors: {e}"), - source: Some(Box::new(e)), - })?; - if let Some(raw_write_start) = raw_write_start { - raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); - } - bytes_written = bytes_written - .checked_add(vectors.bytes.len()) - .ok_or_else(|| Error::DataInvalid { - message: "vindex spilled byte count overflows usize".to_string(), - source: None, - })?; - rows_seen = batch_end; - } - - if !training_buffer.is_empty() { - trainer - .add_training_vectors_mut(&training_buffer, training_buffer.len() / dimension_usize) - .map_err(|e| Error::DataInvalid { - message: format!("Failed to add vindex training vectors: {e}"), - source: Some(Box::new(e)), - })?; - } - if rows_seen != row_count_usize - || expected_row_id - != shard - .row_range_end - .checked_add(1) - .ok_or_else(|| Error::DataInvalid { - message: "vindex row range end overflows i64".to_string(), - source: None, - })? - || (training_vector_count != row_count_usize - && next_training_sample != training_vector_count) - || bytes_written != expected_bytes - { - return Err(Error::DataInvalid { - message: format!( - "vindex streamed data mismatch: rows={rows_seen}/{row_count_usize}, training={next_training_sample}/{training_vector_count}, bytes={bytes_written}/{expected_bytes}" - ), - source: None, - }); - } - let raw_write_start = timing_enabled.then(Instant::now); - raw_file.flush().await.map_err(|e| Error::UnexpectedError { - message: format!("Failed to flush temporary vindex vector file: {e}"), - source: Some(Box::new(e)), - })?; - if let Some(raw_write_start) = raw_write_start { - raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); - } - let raw_file_len = raw_file - .metadata() - .await - .map_err(|e| Error::UnexpectedError { - message: format!("Failed to inspect temporary vindex vector file: {e}"), - source: Some(Box::new(e)), - })? - .len(); - if raw_file_len != expected_bytes as u64 { - return Err(Error::DataInvalid { - message: format!( - "temporary vindex vector file size mismatch: {raw_file_len}/{expected_bytes}" - ), - source: None, - }); - } - let raw_file = raw_file.into_std().await; - // Diagnostics only: never fail the build for a timing log field. - let training_rows_retained = if timing_enabled { - default_training_vector_count(training_vector_count, options.config.nlist()) - .unwrap_or(0) - } else { - 0 - }; - - let (writer, train_finish, raw_temp_reread, index_add) = tokio::task::spawn_blocking( - move || -> std::io::Result<(VectorIndexWriter, Duration, Duration, Duration)> { - let train_start = timing_enabled.then(Instant::now); - let training = trainer.finish()?; - let train_finish = train_start.map_or(Duration::ZERO, |start| start.elapsed()); - let mut writer = VectorIndexWriter::new(training); - let mut raw_temp_reread = Duration::ZERO; - let mut index_add = Duration::ZERO; - let mut raw_file = raw_file; - let reread_start = timing_enabled.then(Instant::now); - raw_file.seek(SeekFrom::Start(0))?; - if let Some(start) = reread_start { - raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); - } - let batch_rows = training_buffer_rows.min(row_count_usize); - let batch_bytes = checked_std_vector_bytes(batch_rows, dimension_usize)?; - let mut buffer = MutableBuffer::new(batch_bytes); - let mut ids = Vec::with_capacity(batch_rows); - let mut rows_added = 0usize; - while rows_added < row_count_usize { - let rows = batch_rows.min(row_count_usize - rows_added); - buffer.resize(checked_std_vector_bytes(rows, dimension_usize)?, 0); - let reread_start = timing_enabled.then(Instant::now); - raw_file.read_exact(buffer.as_slice_mut())?; - if let Some(start) = reread_start { - raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); - } - ids.clear(); - for row in rows_added..rows_added + rows { - ids.push(i64::try_from(row).map_err(|_| { - std::io::Error::new( - std::io::ErrorKind::InvalidData, - "vindex row id does not fit i64", - ) - })?); - } - let add_start = timing_enabled.then(Instant::now); - writer.add_vectors(&ids, buffer.typed_data::(), rows)?; - if let Some(start) = add_start { - index_add = index_add.saturating_add(start.elapsed()); - } - rows_added += rows; - } - let mut trailing = [0u8; 1]; - let reread_start = timing_enabled.then(Instant::now); - if raw_file.read(&mut trailing)? != 0 { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "temporary vindex vector file contains trailing bytes", - )); - } - if let Some(start) = reread_start { - raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); - } - Ok((writer, train_finish, raw_temp_reread, index_add)) - }, - ) - .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 or add vectors to vindex index: {e}"), - source: Some(Box::new(e)), - })?; - - let serialize_upload_start = timing_enabled.then(Instant::now); - self.table - .file_io() - .mkdirs(&format!( - "{}/{INDEX_DIR}/", - self.table.location().trim_end_matches('/') - )) - .await?; - let file_name = format!( - "vector-{}-global-index-{}.index", - self.index_type, - uuid::Uuid::new_v4() - ); - let index_path = format!( - "{}/{INDEX_DIR}/{}", - self.table.location().trim_end_matches('/'), - file_name - ); - let write_result = async { - let async_writer = self - .table - .file_io() - .new_output(&index_path)? - .async_writer() - .await?; - let mut output = SyncIoBridge::new(async_writer); - tokio::task::spawn_blocking(move || -> std::io::Result<()> { - let mut writer = writer; - writer.write(&mut PosWriter::new(&mut output))?; - output.shutdown() - }) - .await - .map_err(|e| Error::UnexpectedError { - message: format!("vindex serialization task failed: {e}"), - source: None, - })? - .map_err(|e| Error::UnexpectedError { - message: format!("Failed to stream vindex index: {e}"), - source: Some(Box::new(e)), - })?; - self.table.file_io().get_status(&index_path).await - } - .await; - let status = match write_result { - Ok(status) => status, - Err(error) => { - let _ = self.table.file_io().delete_file(&index_path).await; - return Err(error); - } - }; - let serialize_upload = - serialize_upload_start.map_or(Duration::ZERO, |start| start.elapsed()); - let meta = IndexFileMeta { - index_type: self.index_type.clone(), - file_name: file_name.clone(), - file_size: checked_i64( - status.size, - "Index file is too large for Rust IndexFileMeta", - )?, - row_count, - deletion_vectors_ranges: None, - global_index_meta: Some(GlobalIndexMeta { - row_range_start: shard.row_range_start, - row_range_end: shard.row_range_end, - index_field_id, - extra_field_ids: None, - source_meta: None, - index_meta: Some(index_meta), - }), - }; - let (oss_read, parquet_decode) = read_timing - .as_ref() - .map_or((Duration::ZERO, Duration::ZERO), |timing| { - (timing.file_read(), timing.parquet_decode()) - }); - let (file_schema_open, first_batch_wait, remaining_batch_wait) = read_timing - .as_ref() - .map_or((Duration::ZERO, Duration::ZERO, Duration::ZERO), |timing| { - timing.file_waits() - }); - let parquet_diagnostics = parquet_read_budget - .as_ref() - .map_or_else(Default::default, |budget| budget.diagnostics()); - let timing = total_start.map(|start| VectorIndexBuildTiming { - total_without_commit: start.elapsed(), - source_batch_wait, - oss_read, - parquet_decode, - file_schema_open, - first_batch_wait, - remaining_batch_wait, - parquet_row_group_count: parquet_diagnostics.row_group_count, - parquet_projected_bytes_min: parquet_diagnostics.projected_bytes_min, - parquet_projected_bytes_max: parquet_diagnostics.projected_bytes_max, - parquet_projected_bytes_total: parquet_diagnostics.projected_bytes_total, - parquet_peak_inflight_row_groups: parquet_diagnostics.peak_inflight, - raw_temp_write, - train_finish, - raw_temp_reread, - index_add, - serialize_upload, - rows: row_count_usize, - training_rows_seen: training_vector_count, - training_rows_retained, - batch_count, - raw_temp_bytes: bytes_written, - index_bytes: status.size, - data_file_count: shard.files.len(), - file_name, - }); - Ok(BuiltIndexFile { meta, timing }) - } -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct VindexIndexShard { - pub partition: BinaryRow, - pub partition_bytes: Vec, - pub files: Vec, - pub row_range_start: i64, - pub row_range_end: i64, - snapshot_id: i64, - source_bucket: i32, - total_buckets: i32, - bucket_path: String, -} - -fn validate_table_options(table: &Table, core_options: &CoreOptions) -> Result<()> { - if !core_options.row_tracking_enabled() { - return Err(Error::DataInvalid { - message: "vindex index build requires 'row-tracking.enabled' = 'true'".to_string(), - source: None, - }); - } - if !core_options.data_evolution_enabled() { - return Err(Error::DataInvalid { - message: "vindex index build requires 'data-evolution.enabled' = 'true'".to_string(), - source: None, - }); - } - if !core_options.global_index_enabled() { - return Err(Error::DataInvalid { - message: "vindex index build requires 'global-index.enabled' = 'true'".to_string(), - source: None, - }); - } - if !table.schema().primary_keys().is_empty() { - return Err(Error::Unsupported { - message: "vindex index build does not support primary-key tables".to_string(), - }); - } - if core_options.deletion_vectors_enabled() { - return Err(Error::Unsupported { - message: - "vindex index build does not support tables with deletion-vectors.enabled=true" - .to_string(), - }); - } - Ok(()) -} - -fn find_index_field<'a>(table: &'a Table, column: &str) -> Result<&'a DataField> { - table - .schema() - .fields() - .iter() - .find(|field| field.name() == column) - .ok_or_else(|| Error::ColumnNotExist { - full_name: table.identifier().full_name(), - column: column.to_string(), - }) -} - -fn validate_vector_field(field: &DataField) -> Result<()> { - let is_array_float = matches!( - field.data_type(), - DataType::Array(array) if matches!(array.element_type(), DataType::Float(_)) - ); - let is_vector_float = matches!( - field.data_type(), - DataType::Vector(vector) if matches!(vector.element_type(), DataType::Float(_)) - ); - if !is_array_float && !is_vector_float { - return Err(Error::DataInvalid { - message: format!( - "vindex index requires ARRAY or VECTOR column, got {:?} for column '{}'", - field.data_type(), - field.name() - ), - source: None, - }); - } - Ok(()) -} - -#[allow(clippy::too_many_arguments)] -fn plan_vindex_shards( - table_location: &str, - partition_keys: &[String], - schema_fields: &[DataField], - core_options: &CoreOptions, - snapshot_id: i64, - entries: Vec, - rows_per_shard: i64, - indexed: &[RowRange], -) -> Result> { - if rows_per_shard <= 0 { - return Err(Error::DataInvalid { - message: format!( - "Option 'global-index.row-count-per-shard' must be greater than 0, got: {rows_per_shard}" - ), - source: None, - }); - } - - let mut by_partition_bucket: HashMap<(Vec, i32, i32), Vec> = HashMap::new(); - for entry in entries { - if *entry.kind() != FileKind::Add { - continue; - } - if entry.file().first_row_id.is_none() { - return Err(Error::DataInvalid { - message: format!( - "Data file '{}' is missing first_row_id; cannot build a complete vindex index", - entry.file().file_name - ), - source: None, - }); - } - let (partition, bucket, total_buckets, file) = entry.into_parts(); - by_partition_bucket - .entry((partition, bucket, total_buckets)) - .or_default() - .push(file); - } - - let mut result = Vec::new(); - for ((partition_bytes, source_bucket, total_buckets), files) in by_partition_bucket { - let partition = if partition_keys.is_empty() { - BinaryRow::new(0) - } else { - BinaryRow::from_serialized_bytes(&partition_bytes)? - }; - let bucket_path = bucket_path( - table_location, - partition_keys, - schema_fields, - core_options, - &partition, - source_bucket, - )?; - let mut files_by_shard: HashMap> = HashMap::new(); - for file in files { - let (file_start, file_end) = file.row_id_range().ok_or_else(|| Error::DataInvalid { - message: format!( - "Data file '{}' is missing first_row_id; cannot build a complete vindex index", - file.file_name - ), - source: None, - })?; - let start_shard = file_start / rows_per_shard; - let end_shard = file_end / rows_per_shard; - for shard_id in start_shard..=end_shard { - files_by_shard - .entry(shard_id * rows_per_shard) - .or_default() - .push(file.clone()); - } - } - - let mut shard_starts = files_by_shard.keys().copied().collect::>(); - shard_starts.sort_unstable(); - for shard_start in shard_starts { - let shard_end = shard_start + rows_per_shard - 1; - let mut shard_files = files_by_shard.remove(&shard_start).unwrap_or_default(); - shard_files.sort_by_key(|file| file.first_row_id); - let groups = group_contiguous_files(shard_files)?; - for group in groups { - let group_start = group - .first() - .and_then(|file| file.first_row_id) - .expect("planned groups are non-empty and row-id assigned"); - let group_end = group - .iter() - .map(|file| file.row_id_range().unwrap().1) - .max() - .unwrap(); - // Coverage of this group clamped to the current shard cell. Then - // subtract the already-indexed ranges so the build only covers - // the gap. Because grid-clamp and gap-subtraction are both range - // intersections, applying the gap here is equivalent to btree's - // "exclude then split" -- and each surviving segment stays inside - // one shard cell, preserving per-shard row-id contiguity (the - // reader errors on a row-id gap within a shard). - let coverage_start = group_start.max(shard_start); - let coverage_end = group_end.min(shard_end); - let build_segments = - exclude_row_ranges(&[RowRange::new(coverage_start, coverage_end)], indexed); - for seg in build_segments { - result.push(VindexIndexShard { - partition: partition.clone(), - partition_bytes: partition_bytes.clone(), - files: group.clone(), - row_range_start: seg.from(), - row_range_end: seg.to(), - snapshot_id, - source_bucket, - total_buckets, - bucket_path: bucket_path.clone(), - }); - } - } - } - } - result.sort_by(|a, b| { - a.partition - .to_serialized_bytes() - .cmp(&b.partition.to_serialized_bytes()) - .then(a.source_bucket.cmp(&b.source_bucket)) - .then(a.row_range_start.cmp(&b.row_range_start)) - }); - Ok(result) -} - -fn group_contiguous_files(mut files: Vec) -> Result>> { - if files.is_empty() { - return Ok(Vec::new()); - } - files.sort_by_key(|file| file.first_row_id); - let mut groups = Vec::new(); - let mut current = Vec::new(); - let mut current_end = None; - for file in files { - let (file_start, file_end) = file.row_id_range().ok_or_else(|| Error::DataInvalid { - message: format!( - "Data file '{}' is missing first_row_id; cannot build a complete vindex index", - file.file_name - ), - source: None, - })?; - match current_end { - None => { - current.push(file); - current_end = Some(file_end); - } - Some(end) if file_start <= end + 1 => { - current.push(file); - current_end = Some(end.max(file_end)); - } - Some(_) => { - groups.push(std::mem::take(&mut current)); - current.push(file); - current_end = Some(file_end); - } - } - } - if !current.is_empty() { - groups.push(current); - } - Ok(groups) -} - -fn bucket_path( - table_location: &str, - partition_keys: &[String], - schema_fields: &[DataField], - core_options: &CoreOptions, - partition: &BinaryRow, - bucket: i32, -) -> Result { - let base = table_location.trim_end_matches('/'); - if partition_keys.is_empty() { - return Ok(format!("{base}/{}", bucket_dir_name(bucket))); - } - let computer = crate::spec::PartitionComputer::new( - partition_keys, - schema_fields, - core_options.partition_default_name(), - core_options.legacy_partition_name(), - )?; - Ok(format!( - "{base}/{}{}", - computer.generate_partition_path(partition)?, - bucket_dir_name(bucket) - )) -} - -fn data_split_for_shard(shard: &VindexIndexShard) -> Result { - DataSplitBuilder::new() - .with_snapshot(shard.snapshot_id) - .with_partition(shard.partition.clone()) - .with_bucket(shard.source_bucket) - .with_bucket_path(shard.bucket_path.clone()) - .with_total_buckets(shard.total_buckets) - .with_data_files(shard.files.clone()) - .with_row_ranges(vec![RowRange::new( - shard.row_range_start, - shard.row_range_end, - )]) - .build() -} - -struct ValidatedVectorBatch<'a> { - values: &'a [f32], - bytes: &'a [u8], - row_count: usize, -} - -fn validate_vector_batch<'a>( - batch: &'a RecordBatch, - index_column: &str, - dimension: usize, - expected_row_id: &mut i64, -) -> Result> { - let vector_index = batch - .schema() - .index_of(index_column) - .map_err(|e| Error::DataInvalid { - message: format!("Vector column '{index_column}' not found in read batch: {e}"), - source: None, - })?; - let row_id_index = - batch - .schema() - .index_of(ROW_ID_FIELD_NAME) - .map_err(|e| Error::DataInvalid { - message: format!("_ROW_ID column not found in read batch: {e}"), - source: None, - })?; - let column = batch.column(vector_index); - let (values, start, end) = if let Some(array) = column.as_any().downcast_ref::() { - if array.null_count() != 0 { - return Err(Error::DataInvalid { - message: "vindex vector extraction found null vector row".to_string(), - source: None, - }); - } - let offsets = array.value_offsets(); - for offsets in offsets.windows(2) { - let actual = offsets[1] - offsets[0]; - if actual != dimension as i32 { - return Err(Error::DataInvalid { - message: format!( - "vindex vector dimension mismatch: expected {dimension}, got {actual}" - ), - source: None, - }); - } - } - let start = usize::try_from(offsets[0]).map_err(|e| Error::DataInvalid { - message: "vindex vector offset is negative".to_string(), - source: Some(Box::new(e)), - })?; - let end = usize::try_from(offsets[offsets.len() - 1]).map_err(|e| Error::DataInvalid { - message: "vindex vector offset is negative".to_string(), - source: Some(Box::new(e)), - })?; - (array.values(), start, end) - } else if let Some(array) = column.as_any().downcast_ref::() { - let actual = usize::try_from(array.value_length()).map_err(|e| Error::DataInvalid { - message: format!( - "Invalid vindex FixedSizeList dimension: {}", - array.value_length() - ), - source: Some(Box::new(e)), - })?; - if actual != dimension { - return Err(Error::DataInvalid { - message: format!( - "vindex vector dimension mismatch: expected {dimension}, got {actual}" - ), - source: None, - }); - } - if array.null_count() != 0 { - return Err(Error::DataInvalid { - message: "vindex vector extraction found null vector row".to_string(), - source: None, - }); - } - let end = batch - .num_rows() - .checked_mul(dimension) - .ok_or_else(|| Error::DataInvalid { - message: "vindex batch vector length overflows usize".to_string(), - source: None, - })?; - (array.values(), 0, end) - } else { - return Err(Error::DataInvalid { - message: - "vindex vector extraction requires Arrow List or FixedSizeList" - .to_string(), - source: None, - }); - }; - let values = values - .as_any() - .downcast_ref::() - .ok_or_else(|| Error::DataInvalid { - message: "vindex vector extraction requires Float32 vector elements".to_string(), - source: None, - })?; - if values.null_count() != 0 - && values - .nulls() - .is_some_and(|nulls| nulls.slice(start, end - start).null_count() != 0) - { - return Err(Error::DataInvalid { - message: "vindex vector extraction found null vector element".to_string(), - source: None, - }); - } - let row_ids = batch - .column(row_id_index) - .as_any() - .downcast_ref::() - .ok_or_else(|| Error::DataInvalid { - message: "vindex vector extraction requires non-null Int64 _ROW_ID".to_string(), - source: None, - })?; - if row_ids.null_count() != 0 { - return Err(Error::DataInvalid { - message: "vindex vector extraction found null _ROW_ID".to_string(), - source: None, - }); - } - for row_id in row_ids.values() { - if *row_id != *expected_row_id { - return Err(Error::DataInvalid { - message: format!( - "vindex vector extraction expected _ROW_ID {}, got {}", - expected_row_id, row_id - ), - source: None, - }); - } - *expected_row_id = expected_row_id - .checked_add(1) - .ok_or_else(|| Error::DataInvalid { - message: "vindex expected row id overflows i64".to_string(), - source: None, - })?; - } - - let byte_start = checked_vector_bytes(start, 1)?; - let byte_end = checked_vector_bytes(end, 1)?; - Ok(ValidatedVectorBatch { - values: &values.values()[start..end], - bytes: &values.values().inner().as_slice()[byte_start..byte_end], - row_count: batch.num_rows(), - }) -} - -fn checked_vector_bytes(row_count: usize, dimension: usize) -> Result { - row_count - .checked_mul(dimension) - .and_then(|values| values.checked_mul(std::mem::size_of::())) - .ok_or_else(|| Error::DataInvalid { - message: format!( - "vindex vector byte length overflows: row_count={row_count}, dimension={dimension}" - ), - source: None, - }) -} - -fn checked_std_vector_bytes(row_count: usize, dimension: usize) -> std::io::Result { - row_count - .checked_mul(dimension) - .and_then(|values| values.checked_mul(std::mem::size_of::())) - .ok_or_else(|| { - std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "vindex vector byte length overflows usize", - ) - }) -} - -fn checked_training_vector_count(row_count: usize, ratio: f64) -> Result { - if row_count == 0 || !(ratio > 0.0 && ratio <= 1.0) { - return Err(Error::DataInvalid { - message: format!( - "Invalid vindex training sample: row_count={row_count}, ratio={ratio}; expected a positive row count and ratio in (0, 1]" - ), - source: None, - }); - } - Ok(((row_count as f64 * ratio).ceil() as usize).clamp(1, row_count)) -} - -fn checked_training_sample_index(sample: usize, rows: usize, samples: usize) -> Result { - sample - .checked_mul(rows / samples) - .and_then(|base| { - sample - .checked_mul(rows % samples) - .and_then(|remainder| base.checked_add(remainder / samples)) - }) - .ok_or_else(|| Error::DataInvalid { - message: "vindex training sample index overflows usize".to_string(), - source: None, - }) -} - -fn checked_i32(value: u64, context: &str) -> Result { - i32::try_from(value).map_err(|_| Error::DataInvalid { - message: format!("{context}: {value}"), - source: None, - }) -} - -fn checked_i64(value: u64, context: &str) -> Result { - i64::try_from(value).map_err(|_| Error::DataInvalid { - message: format!("{context}: {value}"), - source: None, - }) -} - -fn checked_row_count(row_range_start: i64, row_range_end: i64) -> Result { - if row_range_end < row_range_start { - return Err(Error::DataInvalid { - message: format!("Invalid vindex row range [{row_range_start}, {row_range_end}]"), - source: None, - }); - } - row_range_end - .checked_sub(row_range_start) - .and_then(|count| count.checked_add(1)) - .ok_or_else(|| Error::DataInvalid { - message: format!( - "Row count overflows for row range [{row_range_start}, {row_range_end}]" - ), - source: None, - }) } #[cfg(test)] -mod tests { - use super::*; - use crate::catalog::Identifier; - use crate::io::{FileIO, FileIOBuilder}; - use crate::spec::stats::BinaryTableStats; - use crate::spec::{ - ArrayType, FloatType, IndexManifest, IntType, ManifestEntry, Schema, TableSchema, - }; - use crate::table::TableWrite; - use crate::vindex::IVF_FLAT_IDENTIFIER; - use arrow_array::builder::{FixedSizeListBuilder, Float32Builder, Int64Builder, ListBuilder}; - use arrow_array::{ArrayRef, Int32Array}; - use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; - use chrono::{DateTime, Utc}; - use std::sync::Arc; - - fn data_file(name: &str, first_row_id: Option, row_count: i64) -> DataFileMeta { - DataFileMeta { - file_name: name.to_string(), - file_size: 128, - row_count, - min_key: vec![], - max_key: vec![], - key_stats: BinaryTableStats::new(vec![], vec![], vec![]), - value_stats: BinaryTableStats::new(vec![], vec![], vec![]), - min_sequence_number: 0, - max_sequence_number: 0, - schema_id: 0, - level: 0, - extra_files: vec![], - creation_time: Some( - "2024-09-06T07:45:55.039+00:00" - .parse::>() - .unwrap(), - ), - delete_row_count: None, - embedded_index: None, - first_row_id, - write_cols: None, - external_path: None, - file_source: None, - value_stats_cols: None, - column_max_sequence_numbers: None, - } - } - - fn manifest_entry(file: DataFileMeta) -> ManifestEntry { - ManifestEntry::new(FileKind::Add, vec![], 0, 1, file, 2) - } - - fn table_options(rows_per_shard: &str) -> HashMap { - HashMap::from([ - ("row-tracking.enabled".to_string(), "true".to_string()), - ("data-evolution.enabled".to_string(), "true".to_string()), - ("global-index.enabled".to_string(), "true".to_string()), - ( - "global-index.row-count-per-shard".to_string(), - rows_per_shard.to_string(), - ), - ]) - } - - fn test_table(options: HashMap) -> Table { - let schema = Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column( - "embedding", - DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), - ) - .options(options) - .build() - .unwrap(); - Table::new( - FileIOBuilder::new("memory").build().unwrap(), - Identifier::new("default", "test_table"), - "memory:/test_vindex_builder".to_string(), - TableSchema::new(0, &schema), - None, - ) - } - - fn plan(entries: Vec, rows_per_shard: i64) -> Result> { - plan_with_indexed(entries, rows_per_shard, &[]) - } - - fn plan_with_indexed( - entries: Vec, - rows_per_shard: i64, - indexed: &[RowRange], - ) -> Result> { - let table = test_table(table_options(&rows_per_shard.to_string())); - let core = CoreOptions::new(table.schema().options()); - plan_vindex_shards( - table.location(), - table.schema().partition_keys(), - table.schema().fields(), - &core, - 1, - entries, - rows_per_shard, - indexed, - ) - } - - #[test] - fn test_planner_splits_single_file_across_shards() { - let shards = plan(vec![manifest_entry(data_file("a", Some(0), 25))], 10).unwrap(); - - assert_eq!( - shards - .iter() - .map(|s| (s.row_range_start, s.row_range_end)) - .collect::>(), - vec![(0, 9), (10, 19), (20, 24)] - ); - } - - #[test] - fn test_planner_rejects_missing_first_row_id() { - let err = plan(vec![manifest_entry(data_file("a", None, 5))], 10) - .expect_err("missing first_row_id should fail"); - assert!( - matches!(err, Error::DataInvalid { message, .. } if message.contains("missing first_row_id")) - ); - } - - #[test] - fn test_validate_vector_field_accepts_array_float() { - let field = DataField::new( - 0, - "embedding".to_string(), - DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), - ); - assert!(validate_vector_field(&field).is_ok()); - } - - fn vector_batch(rows: Vec>>>, row_ids: Vec>) -> RecordBatch { - let mut vector_builder = ListBuilder::new(Float32Builder::new()); - for row in rows { - match row { - Some(values) => { - for value in values { - match value { - Some(value) => vector_builder.values().append_value(value), - None => vector_builder.values().append_null(), - } - } - vector_builder.append(true); - } - None => vector_builder.append(false), - } - } - let mut row_id_builder = Int64Builder::new(); - for row_id in row_ids { - match row_id { - Some(value) => row_id_builder.append_value(value), - None => row_id_builder.append_null(), - } - } - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new( - "embedding", - ArrowDataType::List(Arc::new(ArrowField::new( - "item", - ArrowDataType::Float32, - true, - ))), - true, - ), - ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), - ])); - RecordBatch::try_new( - schema, - vec![ - Arc::new(vector_builder.finish()) as ArrayRef, - Arc::new(row_id_builder.finish()) as ArrayRef, - ], - ) - .unwrap() - } - - fn extract_vectors_from_batches( - batches: &[RecordBatch], - index_column: &str, - dimension: i32, - row_range_start: i64, - expected_row_count: i64, - ) -> Result> { - let dimension = usize::try_from(dimension).map_err(|e| Error::DataInvalid { - message: format!("Invalid vindex dimension: {dimension}"), - source: Some(Box::new(e)), - })?; - let mut expected_row_id = row_range_start; - let mut vectors = Vec::new(); - for batch in batches { - vectors.extend_from_slice( - validate_vector_batch(batch, index_column, dimension, &mut expected_row_id)?.values, - ); - } - if expected_row_id - row_range_start != expected_row_count { - return Err(Error::DataInvalid { - message: format!( - "vindex vector extraction expected {expected_row_count} rows, got {}", - expected_row_id - row_range_start - ), - source: None, - }); - } - Ok(vectors) - } - - #[test] - fn test_extract_vectors_accepts_list_float32_and_row_ids() { - let batch = vector_batch( - vec![ - Some(vec![Some(1.0), Some(2.0)]), - Some(vec![Some(3.0), Some(4.0)]), - ], - vec![Some(10), Some(11)], - ); - - let vectors = extract_vectors_from_batches(&[batch], "embedding", 2, 10, 2).unwrap(); - - assert_eq!(vectors, vec![1.0, 2.0, 3.0, 4.0]); - } - - #[test] - fn test_extract_vectors_rejects_dimension_mismatch() { - let batch = vector_batch(vec![Some(vec![Some(1.0)])], vec![Some(0)]); - - let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) - .expect_err("dimension mismatch should fail"); - - assert!( - matches!(err, Error::DataInvalid { message, .. } if message.contains("dimension mismatch")) - ); - } - - #[test] - fn test_extract_vectors_handles_sliced_list_offsets() { - let batch = vector_batch( - vec![ - Some(vec![None, Some(0.0)]), - Some(vec![Some(1.0), Some(2.0)]), - Some(vec![Some(3.0), Some(4.0)]), - ], - vec![Some(9), Some(10), Some(11)], - ) - .slice(1, 2); - - let vectors = extract_vectors_from_batches(&[batch], "embedding", 2, 10, 2).unwrap(); - - assert_eq!(vectors, vec![1.0, 2.0, 3.0, 4.0]); - } - - #[test] - fn test_extract_vectors_handles_sliced_fixed_size_list() { - let mut vectors = FixedSizeListBuilder::new(Float32Builder::new(), 2); - for row in [[0.0, 0.0], [1.0, 2.0], [3.0, 4.0]] { - vectors.values().append_slice(&row); - vectors.append(true); - } - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new( - "embedding", - ArrowDataType::FixedSizeList( - Arc::new(ArrowField::new("item", ArrowDataType::Float32, true)), - 2, - ), - true, - ), - ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, false), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(vectors.finish()) as ArrayRef, - Arc::new(Int64Array::from(vec![9, 10, 11])) as ArrayRef, - ], - ) - .unwrap() - .slice(1, 2); - - let vectors = extract_vectors_from_batches(&[batch], "embedding", 2, 10, 2).unwrap(); - - assert_eq!(vectors, vec![1.0, 2.0, 3.0, 4.0]); - } - - #[test] - fn test_training_sample_count_and_indexes_match_java() { - assert_eq!(checked_training_vector_count(10, 0.01).unwrap(), 1); - assert_eq!(checked_training_vector_count(10, 0.25).unwrap(), 3); - assert_eq!(checked_training_vector_count(10, 1.0).unwrap(), 10); - assert_eq!(checked_training_vector_count(3, 0.9).unwrap(), 3); - assert_eq!( - (0..4) - .map(|sample| checked_training_sample_index(sample, 10, 4).unwrap()) - .collect::>(), - vec![0, 2, 5, 7] - ); - assert!(checked_vector_bytes(usize::MAX, 2).is_err()); - assert!(checked_training_sample_index(usize::MAX, usize::MAX, 1).is_err()); - } - - fn test_table_with_io(file_io: FileIO, table_path: &str, schema: Schema) -> Table { - Table::new( - file_io, - Identifier::new("default", "test_table"), - table_path.to_string(), - TableSchema::new(0, &schema), - None, - ) - } - - fn vindex_schema_builder(options: HashMap) -> crate::spec::SchemaBuilder { - Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column( - "embedding", - DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), - ) - .options(options) - } - - fn vindex_e2e_options(rows_per_shard: &str) -> HashMap { - let mut options = table_options(rows_per_shard); - // A small, valid IVF config so the (optional) native build can run; the - // no-op/incremental fix is exercised before or independently of it. - options.insert("ivf-flat.dimension".to_string(), "2".to_string()); - options.insert("ivf-flat.nlist".to_string(), "2".to_string()); - options - } - - fn vindex_e2e_table(table_path: &str, rows_per_shard: &str) -> Table { - test_table_with_io( - FileIOBuilder::new("memory").build().unwrap(), - table_path, - vindex_schema_builder(vindex_e2e_options(rows_per_shard)) - .build() - .unwrap(), - ) - } - - async fn setup_dirs(file_io: &FileIO, table_path: &str) { - file_io - .mkdirs(&format!("{table_path}/snapshot/")) - .await - .unwrap(); - file_io - .mkdirs(&format!("{table_path}/manifest/")) - .await - .unwrap(); - } - - fn build_vector_batch(ids: Vec, vectors: Vec>) -> RecordBatch { - let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); - let mut vector_builder = - ListBuilder::new(Float32Builder::new()).with_field(element_field.clone()); - for vector in vectors { - for value in vector { - vector_builder.values().append_value(value); - } - vector_builder.append(true); - } - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", ArrowDataType::Int32, false), - ArrowField::new("embedding", ArrowDataType::List(element_field), true), - ])); - RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from(ids)) as ArrayRef, - Arc::new(vector_builder.finish()) as ArrayRef, - ], - ) - .unwrap() - } - - async fn write_vectors(table: &Table, ids: Vec, vectors: Vec>) { - let mut table_write = TableWrite::new(table, "test-user".to_string()).unwrap(); - table_write - .write_arrow_batch(&build_vector_batch(ids, vectors)) - .await - .unwrap(); - let messages = table_write.prepare_commit().await.unwrap(); - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(messages) - .await - .unwrap(); - } - - /// Commit a synthetic vindex `IndexFileMeta` covering `[start, end]` for - /// `field_id` directly into the index manifest, without invoking the native - /// builder. Mirrors the Lumina/btree tests so the incremental gap logic can - /// be exercised without a trained vector index. Writes the same `index_type` - /// (`ivf-flat`) the builder-under-test uses, so the gap helper matches it. - async fn commit_synthetic_vindex_index(table: &Table, field_id: i32, start: i64, end: i64) { - let synthetic = IndexFileMeta { - index_type: IVF_FLAT_IDENTIFIER.to_string(), - file_name: format!("vector-ivf-flat-synthetic-{start}-{end}.index"), - file_size: 1, - row_count: end - start + 1, - deletion_vectors_ranges: None, - global_index_meta: Some(GlobalIndexMeta { - row_range_start: start, - row_range_end: end, - index_field_id: field_id, - extra_field_ids: None, - source_meta: None, - index_meta: None, - }), - }; - let mut message = CommitMessage::new(BinaryRow::new(0).to_serialized_bytes(), 0, vec![]); - message.new_index_files = vec![synthetic]; - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(vec![message]) - .await - .unwrap(); - } - - async fn latest_vindex_index_files(table: &Table) -> Vec { - let snapshot_manager = - SnapshotManager::new(table.file_io().clone(), table.location().to_string()); - let snapshot = snapshot_manager - .get_latest_snapshot() - .await - .unwrap() - .unwrap(); - let Some(index_manifest_name) = snapshot.index_manifest() else { - return Vec::new(); - }; - IndexManifest::read( - table.file_io(), - &snapshot_manager.manifest_path(index_manifest_name), - ) - .await - .unwrap() - .into_iter() - .filter(|entry| { - entry.kind == FileKind::Add && entry.index_file.index_type == IVF_FLAT_IDENTIFIER - }) - .map(|entry| entry.index_file) - .collect() - } - - /// Row-id coverage of the committed data files, read back from the data - /// manifest (never hard-coded) and merged into contiguous ranges. Mirrors - /// how `execute` gathers `manifest_entries`. - async fn data_row_id_coverage(table: &Table) -> Vec { - let snapshot_manager = - SnapshotManager::new(table.file_io().clone(), table.location().to_string()); - let snapshot = snapshot_manager - .get_latest_snapshot() - .await - .unwrap() - .unwrap(); - let entries = table - .new_read_builder() - .new_scan() - .with_scan_all_files() - .plan_manifest_entries(&snapshot) - .await - .unwrap(); - let ranges = entries - .iter() - .filter(|entry| *entry.kind() == FileKind::Add) - .filter_map(|entry| { - entry - .file() - .row_id_range() - .map(|(start, end)| RowRange::new(start, end)) - }) - .collect::>(); - crate::table::merge_row_ranges(ranges) - } - - /// Second build with the whole coverage already indexed must be a clean - /// no-op (returns 0), not an overlap error. Reaches `Ok(0)` before the - /// native build, so it runs in CI without a trained index. This is the core - /// bug fix: today the second call errors with the overlap message. - #[tokio::test] - async fn vindex_second_build_without_new_data_is_noop() { - let table_path = "memory:/test_vindex_second_build_noop"; - let table = vindex_e2e_table(table_path, "10"); - setup_dirs(table.file_io(), table_path).await; - - write_vectors(&table, vec![1, 2], vec![vec![1.0, 0.0], vec![0.0, 1.0]]).await; - write_vectors(&table, vec![3], vec![vec![1.0, 1.0]]).await; - - // Fully index the coverage via a synthetic manifest entry. - let coverage = data_row_id_coverage(&table).await; - assert_eq!(coverage.len(), 1, "data must be one contiguous range"); - let field_id = find_index_field(&table, "embedding").unwrap().id(); - commit_synthetic_vindex_index(&table, field_id, coverage[0].from(), coverage[0].to()).await; - - let names_before = latest_vindex_index_files(&table) - .await - .iter() - .map(|f| f.file_name.clone()) - .collect::>(); - assert!(!names_before.is_empty()); - - let built = table - .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) - .with_index_column("embedding") - .execute() - .await - .unwrap(); - assert_eq!(built, 0, "fully-indexed table must build nothing on re-run"); - - let names_after = latest_vindex_index_files(&table) - .await - .iter() - .map(|f| f.file_name.clone()) - .collect::>(); - assert_eq!( - names_before, names_after, - "re-run must not add or remove index manifest entries" - ); - } - - /// Real end-to-end incremental build. `paimon-vindex-core` is pure Rust and - /// trains/serializes an IVF-flat index in CI without a native lib, so this - /// asserts SUCCESS end-to-end (mirroring btree's incremental test): build #1 - /// indexes the initial rows, an appended batch is indexed by build #2, every - /// new index file's row range lies entirely in the appended gap `[n, ..]` - /// (`n` derived from the manifest, never hard-coded), and build-#1's index - /// files are retained untouched (append-only). No overlap error, no tolerated - /// native-build failure -- the build must actually succeed. - #[tokio::test] - async fn vindex_incremental_build_indexes_only_new_rows() { - let table_path = "memory:/test_vindex_incremental"; - let table = vindex_e2e_table(table_path, "10"); - setup_dirs(table.file_io(), table_path).await; - - // Build #1 over the initial batch via a real end-to-end build. - write_vectors( - &table, - vec![1, 2, 3], - vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]], - ) - .await; - let first_built = table - .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) - .with_index_column("embedding") - .with_options(HashMap::from([( - "ivf-flat.train.sample-ratio".to_string(), - "0.9".to_string(), - )])) - .execute() - .await - .unwrap(); - assert!(first_built > 0, "first build must index the initial rows"); - - // First appended row-id, derived from the data manifest (never hard-coded). - let indexed_coverage = data_row_id_coverage(&table).await; - assert_eq!(indexed_coverage.len(), 1); - let n = indexed_coverage[0].to() + 1; - - let first_names = latest_vindex_index_files(&table) - .await - .iter() - .map(|f| f.file_name.clone()) - .collect::>(); - assert_eq!(first_names.len(), 1, "one shard must write one index file"); - - // Append a second batch (new row-ids [n..]). - write_vectors( - &table, - vec![4, 5, 6], - vec![vec![2.0, 0.0], vec![0.0, 2.0], vec![2.0, 2.0]], - ) - .await; - - // End-to-end: build #2 must SUCCEED and index the appended rows. - let second_built = table - .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) - .with_index_column("embedding") - .execute() - .await - .unwrap(); - assert!(second_built > 0, "appended rows must be indexed"); - - let all_files = latest_vindex_index_files(&table).await; - let all_names = all_files - .iter() - .map(|f| f.file_name.clone()) - .collect::>(); - - // Every build-#1 file is still present (append-only, no rewrite/delete). - assert!( - first_names.iter().all(|name| all_names.contains(name)), - "build #1 index files must be retained untouched" - ); - - // Every build-#2 file covers only the appended gap [n, ..], never the - // already-indexed prefix. - let new_files = all_files - .iter() - .filter(|f| !first_names.contains(&f.file_name)) - .collect::>(); - assert!(!new_files.is_empty(), "build #2 must add new index files"); - for file in new_files { - let meta = file - .global_index_meta - .as_ref() - .expect("global index meta on new vindex file"); - assert!( - meta.row_range_start >= n, - "new index file range must start at or after {n}, got [{}, {}]", - meta.row_range_start, - meta.row_range_end - ); - } - } - - #[tokio::test] - async fn vindex_build_cleans_written_shards_when_later_shard_fails() { - let table_path = "memory:/test_vindex_abort_written_shard"; - let table = vindex_e2e_table(table_path, "2"); - setup_dirs(table.file_io(), table_path).await; - write_vectors( - &table, - vec![1, 2, 3], - vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0]], - ) - .await; - - let error = table - .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) - .with_index_column("embedding") - .execute() - .await - .expect_err("the second shard has an invalid vector dimension"); - - assert!(error.to_string().contains("dimension mismatch")); - assert!(table - .file_io() - .list_status(&format!("{table_path}/{INDEX_DIR}/")) - .await - .unwrap() - .is_empty()); - } - - /// A field that already carries a DIFFERENT index type (`lumina`) over an - /// overlapping row range must not block a vindex (`ivf-flat`) build on the - /// same field: the two indexes have distinct identities and coexist. Before - /// the full-identity fix, the overlap guard keyed only on field id + range - /// and spuriously rejected this build with the "overlaps requested row - /// range" error. - #[tokio::test] - async fn vindex_build_coexists_with_different_index_type_on_same_field() { - let table_path = "memory:/test_vindex_coexist_diff_type"; - let table = vindex_e2e_table(table_path, "10"); - setup_dirs(table.file_io(), table_path).await; - - write_vectors( - &table, - vec![1, 2, 3], - vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]], - ) - .await; - - // Pre-existing `lumina` index covering the full data range on the SAME - // field the vindex build will target. - let coverage = data_row_id_coverage(&table).await; - assert_eq!(coverage.len(), 1, "data must be one contiguous range"); - let field_id = find_index_field(&table, "embedding").unwrap().id(); - let lumina = IndexFileMeta { - index_type: "lumina".to_string(), - file_name: "lumina-synthetic-0.index".to_string(), - file_size: 1, - row_count: (coverage[0].to() - coverage[0].from() + 1) as i64, - deletion_vectors_ranges: None, - global_index_meta: Some(GlobalIndexMeta { - row_range_start: coverage[0].from(), - row_range_end: coverage[0].to(), - index_field_id: field_id, - extra_field_ids: None, - source_meta: None, - index_meta: None, - }), - }; - let mut message = CommitMessage::new(BinaryRow::new(0).to_serialized_bytes(), 0, vec![]); - message.new_index_files = vec![lumina]; - TableCommit::new(table.clone(), "test-user".to_string()) - .commit(vec![message]) - .await - .unwrap(); - - // Building `ivf-flat` on the same field must NOT trip the overlap guard. - // A native-build failure over the tiny synthetic dataset is tolerated; - // only the overlap error is forbidden. - let result = table - .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) - .with_index_column("embedding") - .execute() - .await; - match result { - Ok(_) => {} - Err(Error::DataInvalid { message, .. }) => { - assert!( - !message.contains("overlaps requested row range"), - "vindex build must coexist with a different-type index on the same field; got: {message}" - ); - } - Err(other) => panic!("unexpected error from vindex build: {other:?}"), - } - } - - /// Regression: a first build (no existing index) must equal the pre-change - /// full build -- subtracting an empty `indexed` yields full coverage. - #[test] - fn vindex_first_build_indexes_full_coverage() { - let full = plan(vec![manifest_entry(data_file("a", Some(0), 25))], 10).unwrap(); - let gapped = - plan_with_indexed(vec![manifest_entry(data_file("a", Some(0), 25))], 10, &[]).unwrap(); - // Empty `indexed` must not alter the shard layout. - assert_eq!( - full.iter() - .map(|s| (s.row_range_start, s.row_range_end)) - .collect::>(), - gapped - .iter() - .map(|s| (s.row_range_start, s.row_range_end)) - .collect::>() - ); - assert_eq!( - full.iter() - .map(|s| (s.row_range_start, s.row_range_end)) - .collect::>(), - vec![(0, 9), (10, 19), (20, 24)], - "first build must cover the full row range across shards" - ); - } - - /// Planner-level mid-coverage hole, mirroring btree/lumina: with a single - /// shard cell (rows_per_shard large enough to hold all data) the grid never - /// splits, so the only split is the indexed hole itself. An indexed range - /// strictly inside the data coverage must carve the build into exactly the - /// two contiguous segments on either side of the hole -- both bounds pinned, - /// and neither segment may span or touch the hole. - #[test] - fn vindex_plan_splits_gap_around_mid_coverage_indexed_hole() { - // Data row-ids [0, 9]; one shard cell [0, 99] so the grid never splits. - let n = 9; - let hole_start = 4; - let hole_end = 6; - let shards = plan_with_indexed( - vec![manifest_entry(data_file("a", Some(0), n + 1))], - 100, - &[RowRange::new(hole_start, hole_end)], - ) - .unwrap(); - - let ranges = shards - .iter() - .map(|s| (s.row_range_start, s.row_range_end)) - .collect::>(); - // Exactly the two contiguous segments around the hole. - assert_eq!( - ranges, - vec![(0, hole_start - 1), (hole_end + 1, n)], - "mid-coverage hole must split into exactly the two segments around it" - ); - // Every emitted range is contiguous and none spans or touches the hole. - for (start, end) in &ranges { - assert!(end >= start, "range must be non-empty: [{start}, {end}]"); - assert!( - *end < hole_start || *start > hole_end, - "shard [{start}, {end}] must not overlap indexed hole [{hole_start}, {hole_end}]" - ); - } - // Together the shards cover exactly coverage - indexed. - let expected = exclude_row_ranges( - &[RowRange::new(0, n)], - &[RowRange::new(hole_start, hole_end)], - ) - .into_iter() - .map(|r| (r.from(), r.to())) - .collect::>(); - assert_eq!( - ranges, expected, - "shards must cover exactly coverage minus the indexed hole" - ); - } - - /// Planner-level incremental prefix. Strengthens - /// `vindex_incremental_build_indexes_only_new_rows`, which asserts only a - /// one-sided lower bound (`row_range_start >= n`): an indexed prefix [0, k] - /// must leave EXACTLY the suffix [k+1, N] on both bounds, split along the - /// shard grid, with nothing re-indexed inside the prefix. - #[test] - fn vindex_plan_incremental_prefix_leaves_suffix() { - // Data row-ids [0, 24], rows_per_shard = 10 -> cells [0,9],[10,19],[20,29]. - // Indexed prefix [0, 9] fully fills the first cell, so the build must be - // exactly [10, 19] and [20, 24] (the suffix split along the grid). - let n = 24; - let k = 9; // prefix [0, k] == the first full shard cell - let shards = plan_with_indexed( - vec![manifest_entry(data_file("a", Some(0), n + 1))], - 10, - &[RowRange::new(0, k)], - ) - .unwrap(); - - let ranges = shards - .iter() - .map(|s| (s.row_range_start, s.row_range_end)) - .collect::>(); - assert_eq!( - ranges, - vec![(k + 1, 19), (20, n)], - "indexed prefix must leave exactly the suffix, split along the shard grid" - ); - // Both bounds pinned (this is what the one-sided existing check omits). - assert_eq!(ranges.first().unwrap().0, k + 1, "suffix must start at k+1"); - assert_eq!(ranges.last().unwrap().1, n, "suffix must end at N"); - // Contiguous, and no shard reaches back into the indexed prefix. - for pair in ranges.windows(2) { - assert_eq!( - pair[1].0, - pair[0].1 + 1, - "ranges must be contiguous: {:?} then {:?}", - pair[0], - pair[1] - ); - } - for (start, end) in &ranges { - assert!( - *start > k, - "shard [{start}, {end}] must not re-index the prefix [0, {k}]" - ); - } - } -} +mod tests; diff --git a/crates/paimon/src/table/vindex_index_build_builder/extraction.rs b/crates/paimon/src/table/vindex_index_build_builder/extraction.rs new file mode 100644 index 000000000..cf7e910e5 --- /dev/null +++ b/crates/paimon/src/table/vindex_index_build_builder/extraction.rs @@ -0,0 +1,190 @@ +// 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::planning::VindexIndexShard; +use super::validation::checked_vector_bytes; +use crate::spec::ROW_ID_FIELD_NAME; +use crate::table::{DataSplit, DataSplitBuilder, RowRange}; +use crate::{Error, Result}; +use arrow_array::{Array, FixedSizeListArray, Float32Array, Int64Array, ListArray, RecordBatch}; + +pub(super) fn data_split_for_shard(shard: &VindexIndexShard) -> Result { + DataSplitBuilder::new() + .with_snapshot(shard.snapshot_id) + .with_partition(shard.partition.clone()) + .with_bucket(shard.source_bucket) + .with_bucket_path(shard.bucket_path.clone()) + .with_total_buckets(shard.total_buckets) + .with_data_files(shard.files.clone()) + .with_row_ranges(vec![RowRange::new( + shard.row_range_start, + shard.row_range_end, + )]) + .build() +} + +pub(super) struct ValidatedVectorBatch<'a> { + pub(super) values: &'a [f32], + pub(super) bytes: &'a [u8], + pub(super) row_count: usize, +} + +pub(super) fn validate_vector_batch<'a>( + batch: &'a RecordBatch, + index_column: &str, + dimension: usize, + expected_row_id: &mut i64, +) -> Result> { + let vector_index = batch + .schema() + .index_of(index_column) + .map_err(|e| Error::DataInvalid { + message: format!("Vector column '{index_column}' not found in read batch: {e}"), + source: None, + })?; + let row_id_index = + batch + .schema() + .index_of(ROW_ID_FIELD_NAME) + .map_err(|e| Error::DataInvalid { + message: format!("_ROW_ID column not found in read batch: {e}"), + source: None, + })?; + let column = batch.column(vector_index); + let (values, start, end) = if let Some(array) = column.as_any().downcast_ref::() { + if array.null_count() != 0 { + return Err(Error::DataInvalid { + message: "vindex vector extraction found null vector row".to_string(), + source: None, + }); + } + let offsets = array.value_offsets(); + for offsets in offsets.windows(2) { + let actual = offsets[1] - offsets[0]; + if actual != dimension as i32 { + return Err(Error::DataInvalid { + message: format!( + "vindex vector dimension mismatch: expected {dimension}, got {actual}" + ), + source: None, + }); + } + } + let start = usize::try_from(offsets[0]).map_err(|e| Error::DataInvalid { + message: "vindex vector offset is negative".to_string(), + source: Some(Box::new(e)), + })?; + let end = usize::try_from(offsets[offsets.len() - 1]).map_err(|e| Error::DataInvalid { + message: "vindex vector offset is negative".to_string(), + source: Some(Box::new(e)), + })?; + (array.values(), start, end) + } else if let Some(array) = column.as_any().downcast_ref::() { + let actual = usize::try_from(array.value_length()).map_err(|e| Error::DataInvalid { + message: format!( + "Invalid vindex FixedSizeList dimension: {}", + array.value_length() + ), + source: Some(Box::new(e)), + })?; + if actual != dimension { + return Err(Error::DataInvalid { + message: format!( + "vindex vector dimension mismatch: expected {dimension}, got {actual}" + ), + source: None, + }); + } + if array.null_count() != 0 { + return Err(Error::DataInvalid { + message: "vindex vector extraction found null vector row".to_string(), + source: None, + }); + } + let end = batch + .num_rows() + .checked_mul(dimension) + .ok_or_else(|| Error::DataInvalid { + message: "vindex batch vector length overflows usize".to_string(), + source: None, + })?; + (array.values(), 0, end) + } else { + return Err(Error::DataInvalid { + message: + "vindex vector extraction requires Arrow List or FixedSizeList" + .to_string(), + source: None, + }); + }; + let values = values + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::DataInvalid { + message: "vindex vector extraction requires Float32 vector elements".to_string(), + source: None, + })?; + if values.null_count() != 0 + && values + .nulls() + .is_some_and(|nulls| nulls.slice(start, end - start).null_count() != 0) + { + return Err(Error::DataInvalid { + message: "vindex vector extraction found null vector element".to_string(), + source: None, + }); + } + let row_ids = batch + .column(row_id_index) + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::DataInvalid { + message: "vindex vector extraction requires non-null Int64 _ROW_ID".to_string(), + source: None, + })?; + if row_ids.null_count() != 0 { + return Err(Error::DataInvalid { + message: "vindex vector extraction found null _ROW_ID".to_string(), + source: None, + }); + } + for row_id in row_ids.values() { + if *row_id != *expected_row_id { + return Err(Error::DataInvalid { + message: format!( + "vindex vector extraction expected _ROW_ID {}, got {}", + expected_row_id, row_id + ), + source: None, + }); + } + *expected_row_id = expected_row_id + .checked_add(1) + .ok_or_else(|| Error::DataInvalid { + message: "vindex expected row id overflows i64".to_string(), + source: None, + })?; + } + + let byte_start = checked_vector_bytes(start, 1)?; + let byte_end = checked_vector_bytes(end, 1)?; + Ok(ValidatedVectorBatch { + values: &values.values()[start..end], + bytes: &values.values().inner().as_slice()[byte_start..byte_end], + row_count: batch.num_rows(), + }) +} diff --git a/crates/paimon/src/table/vindex_index_build_builder/planning.rs b/crates/paimon/src/table/vindex_index_build_builder/planning.rs new file mode 100644 index 000000000..3b4a1fb84 --- /dev/null +++ b/crates/paimon/src/table/vindex_index_build_builder/planning.rs @@ -0,0 +1,47 @@ +// 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 crate::spec::{CoreOptions, DataField, ManifestEntry}; +use crate::table::global_index_build_common::vector::{plan_vector_index_shards, VectorIndexShard}; +use crate::table::RowRange; +use crate::Result; + +pub(crate) type VindexIndexShard = VectorIndexShard; + +#[allow(clippy::too_many_arguments)] +pub(super) fn plan_vindex_shards( + table_location: &str, + partition_keys: &[String], + schema_fields: &[DataField], + core_options: &CoreOptions, + snapshot_id: i64, + entries: Vec, + rows_per_shard: i64, + indexed: &[RowRange], +) -> Result> { + plan_vector_index_shards( + table_location, + partition_keys, + schema_fields, + core_options, + snapshot_id, + entries, + rows_per_shard, + indexed, + "vindex", + ) +} diff --git a/crates/paimon/src/table/vindex_index_build_builder/tests.rs b/crates/paimon/src/table/vindex_index_build_builder/tests.rs new file mode 100644 index 000000000..f6ebe2658 --- /dev/null +++ b/crates/paimon/src/table/vindex_index_build_builder/tests.rs @@ -0,0 +1,863 @@ +// 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::validate_vector_batch; +use super::planning::{plan_vindex_shards, VindexIndexShard}; +use super::validation::{ + checked_training_sample_index, checked_training_vector_count, checked_vector_bytes, + find_index_field, validate_vector_field, +}; + +const INDEX_DIR: &str = "index"; +use crate::catalog::Identifier; +use crate::io::{FileIO, FileIOBuilder}; +use crate::spec::stats::BinaryTableStats; +use crate::spec::{ + ArrayType, BinaryRow, CoreOptions, DataField, DataFileMeta, DataType, FileKind, FloatType, + GlobalIndexMeta, IndexFileMeta, IndexManifest, IntType, ManifestEntry, Schema, TableSchema, + ROW_ID_FIELD_NAME, +}; +use crate::table::source::exclude_row_ranges; +use crate::table::{CommitMessage, RowRange, SnapshotManager, Table, TableCommit, TableWrite}; +use crate::vindex::IVF_FLAT_IDENTIFIER; +use crate::{Error, Result}; +use arrow_array::builder::{FixedSizeListBuilder, Float32Builder, Int64Builder, ListBuilder}; +use arrow_array::{ArrayRef, Int32Array, Int64Array, RecordBatch}; +use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; +use chrono::{DateTime, Utc}; +use std::collections::HashMap; +use std::sync::Arc; + +fn data_file(name: &str, first_row_id: Option, row_count: i64) -> DataFileMeta { + DataFileMeta { + file_name: name.to_string(), + file_size: 128, + row_count, + min_key: vec![], + max_key: vec![], + key_stats: BinaryTableStats::new(vec![], vec![], vec![]), + value_stats: BinaryTableStats::new(vec![], vec![], vec![]), + min_sequence_number: 0, + max_sequence_number: 0, + schema_id: 0, + level: 0, + extra_files: vec![], + creation_time: Some( + "2024-09-06T07:45:55.039+00:00" + .parse::>() + .unwrap(), + ), + delete_row_count: None, + embedded_index: None, + first_row_id, + write_cols: None, + external_path: None, + file_source: None, + value_stats_cols: None, + column_max_sequence_numbers: None, + } +} + +fn manifest_entry(file: DataFileMeta) -> ManifestEntry { + ManifestEntry::new(FileKind::Add, vec![], 0, 1, file, 2) +} + +fn table_options(rows_per_shard: &str) -> HashMap { + HashMap::from([ + ("row-tracking.enabled".to_string(), "true".to_string()), + ("data-evolution.enabled".to_string(), "true".to_string()), + ("global-index.enabled".to_string(), "true".to_string()), + ( + "global-index.row-count-per-shard".to_string(), + rows_per_shard.to_string(), + ), + ]) +} + +fn test_table(options: HashMap) -> Table { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ) + .options(options) + .build() + .unwrap(); + Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "test_table"), + "memory:/test_vindex_builder".to_string(), + TableSchema::new(0, &schema), + None, + ) +} + +fn plan(entries: Vec, rows_per_shard: i64) -> Result> { + plan_with_indexed(entries, rows_per_shard, &[]) +} + +fn plan_with_indexed( + entries: Vec, + rows_per_shard: i64, + indexed: &[RowRange], +) -> Result> { + let table = test_table(table_options(&rows_per_shard.to_string())); + let core = CoreOptions::new(table.schema().options()); + plan_vindex_shards( + table.location(), + table.schema().partition_keys(), + table.schema().fields(), + &core, + 1, + entries, + rows_per_shard, + indexed, + ) +} + +#[test] +fn test_planner_splits_single_file_across_shards() { + let shards = plan(vec![manifest_entry(data_file("a", Some(0), 25))], 10).unwrap(); + + assert_eq!( + shards + .iter() + .map(|s| (s.row_range_start, s.row_range_end)) + .collect::>(), + vec![(0, 9), (10, 19), (20, 24)] + ); +} + +#[test] +fn test_planner_rejects_missing_first_row_id() { + let err = plan(vec![manifest_entry(data_file("a", None, 5))], 10) + .expect_err("missing first_row_id should fail"); + assert!( + matches!(err, Error::DataInvalid { message, .. } if message.contains("missing first_row_id")) + ); +} + +#[test] +fn test_validate_vector_field_accepts_array_float() { + let field = DataField::new( + 0, + "embedding".to_string(), + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ); + assert!(validate_vector_field(&field).is_ok()); +} + +fn vector_batch(rows: Vec>>>, row_ids: Vec>) -> RecordBatch { + let mut vector_builder = ListBuilder::new(Float32Builder::new()); + for row in rows { + match row { + Some(values) => { + for value in values { + match value { + Some(value) => vector_builder.values().append_value(value), + None => vector_builder.values().append_null(), + } + } + vector_builder.append(true); + } + None => vector_builder.append(false), + } + } + let mut row_id_builder = Int64Builder::new(); + for row_id in row_ids { + match row_id { + Some(value) => row_id_builder.append_value(value), + None => row_id_builder.append_null(), + } + } + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "embedding", + ArrowDataType::List(Arc::new(ArrowField::new( + "item", + ArrowDataType::Float32, + true, + ))), + true, + ), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(vector_builder.finish()) as ArrayRef, + Arc::new(row_id_builder.finish()) as ArrayRef, + ], + ) + .unwrap() +} + +fn extract_vectors_from_batches( + batches: &[RecordBatch], + index_column: &str, + dimension: i32, + row_range_start: i64, + expected_row_count: i64, +) -> Result> { + let dimension = usize::try_from(dimension).map_err(|e| Error::DataInvalid { + message: format!("Invalid vindex dimension: {dimension}"), + source: Some(Box::new(e)), + })?; + let mut expected_row_id = row_range_start; + let mut vectors = Vec::new(); + for batch in batches { + vectors.extend_from_slice( + validate_vector_batch(batch, index_column, dimension, &mut expected_row_id)?.values, + ); + } + if expected_row_id - row_range_start != expected_row_count { + return Err(Error::DataInvalid { + message: format!( + "vindex vector extraction expected {expected_row_count} rows, got {}", + expected_row_id - row_range_start + ), + source: None, + }); + } + Ok(vectors) +} + +#[test] +fn test_extract_vectors_accepts_list_float32_and_row_ids() { + let batch = vector_batch( + vec![ + Some(vec![Some(1.0), Some(2.0)]), + Some(vec![Some(3.0), Some(4.0)]), + ], + vec![Some(10), Some(11)], + ); + + let vectors = extract_vectors_from_batches(&[batch], "embedding", 2, 10, 2).unwrap(); + + assert_eq!(vectors, vec![1.0, 2.0, 3.0, 4.0]); +} + +#[test] +fn test_extract_vectors_rejects_dimension_mismatch() { + let batch = vector_batch(vec![Some(vec![Some(1.0)])], vec![Some(0)]); + + let err = extract_vectors_from_batches(&[batch], "embedding", 2, 0, 1) + .expect_err("dimension mismatch should fail"); + + assert!( + matches!(err, Error::DataInvalid { message, .. } if message.contains("dimension mismatch")) + ); +} + +#[test] +fn test_extract_vectors_handles_sliced_list_offsets() { + let batch = vector_batch( + vec![ + Some(vec![None, Some(0.0)]), + Some(vec![Some(1.0), Some(2.0)]), + Some(vec![Some(3.0), Some(4.0)]), + ], + vec![Some(9), Some(10), Some(11)], + ) + .slice(1, 2); + + let vectors = extract_vectors_from_batches(&[batch], "embedding", 2, 10, 2).unwrap(); + + assert_eq!(vectors, vec![1.0, 2.0, 3.0, 4.0]); +} + +#[test] +fn test_extract_vectors_handles_sliced_fixed_size_list() { + let mut vectors = FixedSizeListBuilder::new(Float32Builder::new(), 2); + for row in [[0.0, 0.0], [1.0, 2.0], [3.0, 4.0]] { + vectors.values().append_slice(&row); + vectors.append(true); + } + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "embedding", + ArrowDataType::FixedSizeList( + Arc::new(ArrowField::new("item", ArrowDataType::Float32, true)), + 2, + ), + true, + ), + ArrowField::new(ROW_ID_FIELD_NAME, ArrowDataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(vectors.finish()) as ArrayRef, + Arc::new(Int64Array::from(vec![9, 10, 11])) as ArrayRef, + ], + ) + .unwrap() + .slice(1, 2); + + let vectors = extract_vectors_from_batches(&[batch], "embedding", 2, 10, 2).unwrap(); + + assert_eq!(vectors, vec![1.0, 2.0, 3.0, 4.0]); +} + +#[test] +fn test_training_sample_count_and_indexes_match_java() { + assert_eq!(checked_training_vector_count(10, 0.01).unwrap(), 1); + assert_eq!(checked_training_vector_count(10, 0.25).unwrap(), 3); + assert_eq!(checked_training_vector_count(10, 1.0).unwrap(), 10); + assert_eq!(checked_training_vector_count(3, 0.9).unwrap(), 3); + assert_eq!( + (0..4) + .map(|sample| checked_training_sample_index(sample, 10, 4).unwrap()) + .collect::>(), + vec![0, 2, 5, 7] + ); + assert!(checked_vector_bytes(usize::MAX, 2).is_err()); + assert!(checked_training_sample_index(usize::MAX, usize::MAX, 1).is_err()); +} + +fn test_table_with_io(file_io: FileIO, table_path: &str, schema: Schema) -> Table { + Table::new( + file_io, + Identifier::new("default", "test_table"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ) +} + +fn vindex_schema_builder(options: HashMap) -> crate::spec::SchemaBuilder { + Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ) + .options(options) +} + +fn vindex_e2e_options(rows_per_shard: &str) -> HashMap { + let mut options = table_options(rows_per_shard); + // A small, valid IVF config so the (optional) native build can run; the + // no-op/incremental fix is exercised before or independently of it. + options.insert("ivf-flat.dimension".to_string(), "2".to_string()); + options.insert("ivf-flat.nlist".to_string(), "2".to_string()); + options +} + +fn vindex_e2e_table(table_path: &str, rows_per_shard: &str) -> Table { + test_table_with_io( + FileIOBuilder::new("memory").build().unwrap(), + table_path, + vindex_schema_builder(vindex_e2e_options(rows_per_shard)) + .build() + .unwrap(), + ) +} + +async fn setup_dirs(file_io: &FileIO, table_path: &str) { + file_io + .mkdirs(&format!("{table_path}/snapshot/")) + .await + .unwrap(); + file_io + .mkdirs(&format!("{table_path}/manifest/")) + .await + .unwrap(); +} + +fn build_vector_batch(ids: Vec, vectors: Vec>) -> RecordBatch { + let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut vector_builder = + ListBuilder::new(Float32Builder::new()).with_field(element_field.clone()); + for vector in vectors { + for value in vector { + vector_builder.values().append_value(value); + } + vector_builder.append(true); + } + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("embedding", ArrowDataType::List(element_field), true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(ids)) as ArrayRef, + Arc::new(vector_builder.finish()) as ArrayRef, + ], + ) + .unwrap() +} + +async fn write_vectors(table: &Table, ids: Vec, vectors: Vec>) { + let mut table_write = TableWrite::new(table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&build_vector_batch(ids, vectors)) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); +} + +/// Commit a synthetic vindex `IndexFileMeta` covering `[start, end]` for +/// `field_id` directly into the index manifest, without invoking the native +/// builder. Mirrors the Lumina/btree tests so the incremental gap logic can +/// be exercised without a trained vector index. Writes the same `index_type` +/// (`ivf-flat`) the builder-under-test uses, so the gap helper matches it. +async fn commit_synthetic_vindex_index(table: &Table, field_id: i32, start: i64, end: i64) { + let synthetic = IndexFileMeta { + index_type: IVF_FLAT_IDENTIFIER.to_string(), + file_name: format!("vector-ivf-flat-synthetic-{start}-{end}.index"), + file_size: 1, + row_count: end - start + 1, + deletion_vectors_ranges: None, + global_index_meta: Some(GlobalIndexMeta { + row_range_start: start, + row_range_end: end, + index_field_id: field_id, + extra_field_ids: None, + source_meta: None, + index_meta: None, + }), + }; + let mut message = CommitMessage::new(BinaryRow::new(0).to_serialized_bytes(), 0, vec![]); + message.new_index_files = vec![synthetic]; + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(vec![message]) + .await + .unwrap(); +} + +async fn latest_vindex_index_files(table: &Table) -> Vec { + let snapshot_manager = + SnapshotManager::new(table.file_io().clone(), table.location().to_string()); + let snapshot = snapshot_manager + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let Some(index_manifest_name) = snapshot.index_manifest() else { + return Vec::new(); + }; + IndexManifest::read( + table.file_io(), + &snapshot_manager.manifest_path(index_manifest_name), + ) + .await + .unwrap() + .into_iter() + .filter(|entry| { + entry.kind == FileKind::Add && entry.index_file.index_type == IVF_FLAT_IDENTIFIER + }) + .map(|entry| entry.index_file) + .collect() +} + +/// Row-id coverage of the committed data files, read back from the data +/// manifest (never hard-coded) and merged into contiguous ranges. Mirrors +/// how `execute` gathers `manifest_entries`. +async fn data_row_id_coverage(table: &Table) -> Vec { + let snapshot_manager = + SnapshotManager::new(table.file_io().clone(), table.location().to_string()); + let snapshot = snapshot_manager + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let entries = table + .new_read_builder() + .new_scan() + .with_scan_all_files() + .plan_manifest_entries(&snapshot) + .await + .unwrap(); + let ranges = entries + .iter() + .filter(|entry| *entry.kind() == FileKind::Add) + .filter_map(|entry| { + entry + .file() + .row_id_range() + .map(|(start, end)| RowRange::new(start, end)) + }) + .collect::>(); + crate::table::merge_row_ranges(ranges) +} + +/// Second build with the whole coverage already indexed must be a clean +/// no-op (returns 0), not an overlap error. Reaches `Ok(0)` before the +/// native build, so it runs in CI without a trained index. This is the core +/// bug fix: today the second call errors with the overlap message. +#[tokio::test] +async fn vindex_second_build_without_new_data_is_noop() { + let table_path = "memory:/test_vindex_second_build_noop"; + let table = vindex_e2e_table(table_path, "10"); + setup_dirs(table.file_io(), table_path).await; + + write_vectors(&table, vec![1, 2], vec![vec![1.0, 0.0], vec![0.0, 1.0]]).await; + write_vectors(&table, vec![3], vec![vec![1.0, 1.0]]).await; + + // Fully index the coverage via a synthetic manifest entry. + let coverage = data_row_id_coverage(&table).await; + assert_eq!(coverage.len(), 1, "data must be one contiguous range"); + let field_id = find_index_field(&table, "embedding").unwrap().id(); + commit_synthetic_vindex_index(&table, field_id, coverage[0].from(), coverage[0].to()).await; + + let names_before = latest_vindex_index_files(&table) + .await + .iter() + .map(|f| f.file_name.clone()) + .collect::>(); + assert!(!names_before.is_empty()); + + let built = table + .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) + .with_index_column("embedding") + .execute() + .await + .unwrap(); + assert_eq!(built, 0, "fully-indexed table must build nothing on re-run"); + + let names_after = latest_vindex_index_files(&table) + .await + .iter() + .map(|f| f.file_name.clone()) + .collect::>(); + assert_eq!( + names_before, names_after, + "re-run must not add or remove index manifest entries" + ); +} + +/// Real end-to-end incremental build. `paimon-vindex-core` is pure Rust and +/// trains/serializes an IVF-flat index in CI without a native lib, so this +/// asserts SUCCESS end-to-end (mirroring btree's incremental test): build #1 +/// indexes the initial rows, an appended batch is indexed by build #2, every +/// new index file's row range lies entirely in the appended gap `[n, ..]` +/// (`n` derived from the manifest, never hard-coded), and build-#1's index +/// files are retained untouched (append-only). No overlap error, no tolerated +/// native-build failure -- the build must actually succeed. +#[tokio::test] +async fn vindex_incremental_build_indexes_only_new_rows() { + let table_path = "memory:/test_vindex_incremental"; + let table = vindex_e2e_table(table_path, "10"); + setup_dirs(table.file_io(), table_path).await; + + // Build #1 over the initial batch via a real end-to-end build. + write_vectors( + &table, + vec![1, 2, 3], + vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]], + ) + .await; + let first_built = table + .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) + .with_index_column("embedding") + .with_options(HashMap::from([( + "ivf-flat.train.sample-ratio".to_string(), + "0.9".to_string(), + )])) + .execute() + .await + .unwrap(); + assert!(first_built > 0, "first build must index the initial rows"); + + // First appended row-id, derived from the data manifest (never hard-coded). + let indexed_coverage = data_row_id_coverage(&table).await; + assert_eq!(indexed_coverage.len(), 1); + let n = indexed_coverage[0].to() + 1; + + let first_names = latest_vindex_index_files(&table) + .await + .iter() + .map(|f| f.file_name.clone()) + .collect::>(); + assert_eq!(first_names.len(), 1, "one shard must write one index file"); + + // Append a second batch (new row-ids [n..]). + write_vectors( + &table, + vec![4, 5, 6], + vec![vec![2.0, 0.0], vec![0.0, 2.0], vec![2.0, 2.0]], + ) + .await; + + // End-to-end: build #2 must SUCCEED and index the appended rows. + let second_built = table + .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) + .with_index_column("embedding") + .execute() + .await + .unwrap(); + assert!(second_built > 0, "appended rows must be indexed"); + + let all_files = latest_vindex_index_files(&table).await; + let all_names = all_files + .iter() + .map(|f| f.file_name.clone()) + .collect::>(); + + // Every build-#1 file is still present (append-only, no rewrite/delete). + assert!( + first_names.iter().all(|name| all_names.contains(name)), + "build #1 index files must be retained untouched" + ); + + // Every build-#2 file covers only the appended gap [n, ..], never the + // already-indexed prefix. + let new_files = all_files + .iter() + .filter(|f| !first_names.contains(&f.file_name)) + .collect::>(); + assert!(!new_files.is_empty(), "build #2 must add new index files"); + for file in new_files { + let meta = file + .global_index_meta + .as_ref() + .expect("global index meta on new vindex file"); + assert!( + meta.row_range_start >= n, + "new index file range must start at or after {n}, got [{}, {}]", + meta.row_range_start, + meta.row_range_end + ); + } +} + +#[tokio::test] +async fn vindex_build_cleans_written_shards_when_later_shard_fails() { + let table_path = "memory:/test_vindex_abort_written_shard"; + let table = vindex_e2e_table(table_path, "2"); + setup_dirs(table.file_io(), table_path).await; + write_vectors( + &table, + vec![1, 2, 3], + vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0]], + ) + .await; + + let error = table + .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) + .with_index_column("embedding") + .execute() + .await + .expect_err("the second shard has an invalid vector dimension"); + + assert!(error.to_string().contains("dimension mismatch")); + assert!(table + .file_io() + .list_status(&format!("{table_path}/{INDEX_DIR}/")) + .await + .unwrap() + .is_empty()); +} + +/// A field that already carries a DIFFERENT index type (`lumina`) over an +/// overlapping row range must not block a vindex (`ivf-flat`) build on the +/// same field: the two indexes have distinct identities and coexist. Before +/// the full-identity fix, the overlap guard keyed only on field id + range +/// and spuriously rejected this build with the "overlaps requested row +/// range" error. +#[tokio::test] +async fn vindex_build_coexists_with_different_index_type_on_same_field() { + let table_path = "memory:/test_vindex_coexist_diff_type"; + let table = vindex_e2e_table(table_path, "10"); + setup_dirs(table.file_io(), table_path).await; + + write_vectors( + &table, + vec![1, 2, 3], + vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]], + ) + .await; + + // Pre-existing `lumina` index covering the full data range on the SAME + // field the vindex build will target. + let coverage = data_row_id_coverage(&table).await; + assert_eq!(coverage.len(), 1, "data must be one contiguous range"); + let field_id = find_index_field(&table, "embedding").unwrap().id(); + let lumina = IndexFileMeta { + index_type: "lumina".to_string(), + file_name: "lumina-synthetic-0.index".to_string(), + file_size: 1, + row_count: (coverage[0].to() - coverage[0].from() + 1) as i64, + deletion_vectors_ranges: None, + global_index_meta: Some(GlobalIndexMeta { + row_range_start: coverage[0].from(), + row_range_end: coverage[0].to(), + index_field_id: field_id, + extra_field_ids: None, + source_meta: None, + index_meta: None, + }), + }; + let mut message = CommitMessage::new(BinaryRow::new(0).to_serialized_bytes(), 0, vec![]); + message.new_index_files = vec![lumina]; + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(vec![message]) + .await + .unwrap(); + + // Building `ivf-flat` on the same field must NOT trip the overlap guard. + // A native-build failure over the tiny synthetic dataset is tolerated; + // only the overlap error is forbidden. + let result = table + .new_vindex_index_build_builder(IVF_FLAT_IDENTIFIER) + .with_index_column("embedding") + .execute() + .await; + match result { + Ok(_) => {} + Err(Error::DataInvalid { message, .. }) => { + assert!( + !message.contains("overlaps requested row range"), + "vindex build must coexist with a different-type index on the same field; got: {message}" + ); + } + Err(other) => panic!("unexpected error from vindex build: {other:?}"), + } +} + +/// Regression: a first build (no existing index) must equal the pre-change +/// full build -- subtracting an empty `indexed` yields full coverage. +#[test] +fn vindex_first_build_indexes_full_coverage() { + let full = plan(vec![manifest_entry(data_file("a", Some(0), 25))], 10).unwrap(); + let gapped = + plan_with_indexed(vec![manifest_entry(data_file("a", Some(0), 25))], 10, &[]).unwrap(); + // Empty `indexed` must not alter the shard layout. + assert_eq!( + full.iter() + .map(|s| (s.row_range_start, s.row_range_end)) + .collect::>(), + gapped + .iter() + .map(|s| (s.row_range_start, s.row_range_end)) + .collect::>() + ); + assert_eq!( + full.iter() + .map(|s| (s.row_range_start, s.row_range_end)) + .collect::>(), + vec![(0, 9), (10, 19), (20, 24)], + "first build must cover the full row range across shards" + ); +} + +/// Planner-level mid-coverage hole, mirroring btree/lumina: with a single +/// shard cell (rows_per_shard large enough to hold all data) the grid never +/// splits, so the only split is the indexed hole itself. An indexed range +/// strictly inside the data coverage must carve the build into exactly the +/// two contiguous segments on either side of the hole -- both bounds pinned, +/// and neither segment may span or touch the hole. +#[test] +fn vindex_plan_splits_gap_around_mid_coverage_indexed_hole() { + // Data row-ids [0, 9]; one shard cell [0, 99] so the grid never splits. + let n = 9; + let hole_start = 4; + let hole_end = 6; + let shards = plan_with_indexed( + vec![manifest_entry(data_file("a", Some(0), n + 1))], + 100, + &[RowRange::new(hole_start, hole_end)], + ) + .unwrap(); + + let ranges = shards + .iter() + .map(|s| (s.row_range_start, s.row_range_end)) + .collect::>(); + // Exactly the two contiguous segments around the hole. + assert_eq!( + ranges, + vec![(0, hole_start - 1), (hole_end + 1, n)], + "mid-coverage hole must split into exactly the two segments around it" + ); + // Every emitted range is contiguous and none spans or touches the hole. + for (start, end) in &ranges { + assert!(end >= start, "range must be non-empty: [{start}, {end}]"); + assert!( + *end < hole_start || *start > hole_end, + "shard [{start}, {end}] must not overlap indexed hole [{hole_start}, {hole_end}]" + ); + } + // Together the shards cover exactly coverage - indexed. + let expected = exclude_row_ranges( + &[RowRange::new(0, n)], + &[RowRange::new(hole_start, hole_end)], + ) + .into_iter() + .map(|r| (r.from(), r.to())) + .collect::>(); + assert_eq!( + ranges, expected, + "shards must cover exactly coverage minus the indexed hole" + ); +} + +/// Planner-level incremental prefix. Strengthens +/// `vindex_incremental_build_indexes_only_new_rows`, which asserts only a +/// one-sided lower bound (`row_range_start >= n`): an indexed prefix [0, k] +/// must leave EXACTLY the suffix [k+1, N] on both bounds, split along the +/// shard grid, with nothing re-indexed inside the prefix. +#[test] +fn vindex_plan_incremental_prefix_leaves_suffix() { + // Data row-ids [0, 24], rows_per_shard = 10 -> cells [0,9],[10,19],[20,29]. + // Indexed prefix [0, 9] fully fills the first cell, so the build must be + // exactly [10, 19] and [20, 24] (the suffix split along the grid). + let n = 24; + let k = 9; // prefix [0, k] == the first full shard cell + let shards = plan_with_indexed( + vec![manifest_entry(data_file("a", Some(0), n + 1))], + 10, + &[RowRange::new(0, k)], + ) + .unwrap(); + + let ranges = shards + .iter() + .map(|s| (s.row_range_start, s.row_range_end)) + .collect::>(); + assert_eq!( + ranges, + vec![(k + 1, 19), (20, n)], + "indexed prefix must leave exactly the suffix, split along the shard grid" + ); + // Both bounds pinned (this is what the one-sided existing check omits). + assert_eq!(ranges.first().unwrap().0, k + 1, "suffix must start at k+1"); + assert_eq!(ranges.last().unwrap().1, n, "suffix must end at N"); + // Contiguous, and no shard reaches back into the indexed prefix. + for pair in ranges.windows(2) { + assert_eq!( + pair[1].0, + pair[0].1 + 1, + "ranges must be contiguous: {:?} then {:?}", + pair[0], + pair[1] + ); + } + for (start, end) in &ranges { + assert!( + *start > k, + "shard [{start}, {end}] must not re-index the prefix [0, {k}]" + ); + } +} diff --git a/crates/paimon/src/table/vindex_index_build_builder/timing.rs b/crates/paimon/src/table/vindex_index_build_builder/timing.rs new file mode 100644 index 000000000..3221dc4c1 --- /dev/null +++ b/crates/paimon/src/table/vindex_index_build_builder/timing.rs @@ -0,0 +1,102 @@ +// 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 std::sync::OnceLock; +use std::time::Duration; + +const VECTOR_INDEX_BUILD_TIMING_ENV: &str = "PAIMON_LOG_VECTOR_INDEX_BUILD_TIMING"; + +pub(super) fn vector_index_build_timing_enabled() -> bool { + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| { + std::env::var_os(VECTOR_INDEX_BUILD_TIMING_ENV).is_some_and(|value| value == "1") + }) +} + +pub(super) struct VectorIndexBuildTiming { + pub(super) total_without_commit: Duration, + pub(super) source_batch_wait: Duration, + pub(super) oss_read: Duration, + pub(super) parquet_decode: Duration, + pub(super) file_schema_open: Duration, + pub(super) first_batch_wait: Duration, + pub(super) remaining_batch_wait: Duration, + pub(super) parquet_row_group_count: u64, + pub(super) parquet_projected_bytes_min: u64, + pub(super) parquet_projected_bytes_max: u64, + pub(super) parquet_projected_bytes_total: u64, + pub(super) parquet_peak_inflight_row_groups: usize, + pub(super) raw_temp_write: Duration, + pub(super) train_finish: Duration, + pub(super) raw_temp_reread: Duration, + pub(super) index_add: Duration, + pub(super) serialize_upload: Duration, + pub(super) rows: usize, + pub(super) training_rows_seen: usize, + pub(super) training_rows_retained: usize, + pub(super) batch_count: usize, + pub(super) raw_temp_bytes: usize, + pub(super) index_bytes: u64, + pub(super) data_file_count: usize, + pub(super) file_name: String, +} + +impl VectorIndexBuildTiming { + pub(super) fn log(self, index_type: &str, commit: Duration) { + let total = self.total_without_commit.saturating_add(commit); + let accounted = self + .source_batch_wait + .saturating_add(self.raw_temp_write) + .saturating_add(self.train_finish) + .saturating_add(self.raw_temp_reread) + .saturating_add(self.index_add) + .saturating_add(self.serialize_upload) + .saturating_add(commit); + let unattributed = total.saturating_sub(accounted); + eprintln!( + "event=paimon_vector_index_build index_type={} file={} rows={} training_rows_seen={} training_rows_retained={} batch_count={} raw_temp_bytes={} index_bytes={} source_batch_wait_ms={:.3} oss_read_ms={:.3} parquet_decode_ms={:.3} file_schema_open_ms={:.3} first_batch_wait_ms={:.3} remaining_batch_wait_ms={:.3} parquet_row_group_count={} parquet_projected_bytes_min={} parquet_projected_bytes_max={} parquet_projected_bytes_total={} parquet_peak_inflight_row_groups={} raw_temp_write_ms={:.3} train_finish_ms={:.3} raw_temp_reread_ms={:.3} index_add_ms={:.3} serialize_upload_ms={:.3} commit_ms={:.3} sample_read_ms=0.000 full_scan_add_ms=0.000 pipeline_blocked_ms=0.000 producer_blocked_ms=0.000 consumer_add_ms=0.000 data_file_count={} data_file_read_concurrency=1 peak_ready_batches=0 total_ms={:.3} unattributed_ms={:.3}", + index_type, + self.file_name, + self.rows, + self.training_rows_seen, + self.training_rows_retained, + self.batch_count, + self.raw_temp_bytes, + self.index_bytes, + self.source_batch_wait.as_secs_f64() * 1000.0, + self.oss_read.as_secs_f64() * 1000.0, + self.parquet_decode.as_secs_f64() * 1000.0, + self.file_schema_open.as_secs_f64() * 1000.0, + self.first_batch_wait.as_secs_f64() * 1000.0, + self.remaining_batch_wait.as_secs_f64() * 1000.0, + self.parquet_row_group_count, + self.parquet_projected_bytes_min, + self.parquet_projected_bytes_max, + self.parquet_projected_bytes_total, + self.parquet_peak_inflight_row_groups, + self.raw_temp_write.as_secs_f64() * 1000.0, + self.train_finish.as_secs_f64() * 1000.0, + self.raw_temp_reread.as_secs_f64() * 1000.0, + self.index_add.as_secs_f64() * 1000.0, + self.serialize_upload.as_secs_f64() * 1000.0, + commit.as_secs_f64() * 1000.0, + self.data_file_count, + total.as_secs_f64() * 1000.0, + unattributed.as_secs_f64() * 1000.0, + ); + } +} diff --git a/crates/paimon/src/table/vindex_index_build_builder/validation.rs b/crates/paimon/src/table/vindex_index_build_builder/validation.rs new file mode 100644 index 000000000..f1c5270dc --- /dev/null +++ b/crates/paimon/src/table/vindex_index_build_builder/validation.rs @@ -0,0 +1,155 @@ +// 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 crate::spec::{CoreOptions, DataField}; +use crate::table::global_index_build_common::vector::{ + find_index_field as find_vector_index_field, + validate_vector_field as validate_common_vector_field, +}; +use crate::table::Table; +use crate::{Error, Result}; + +pub(super) fn validate_table_options(table: &Table, core_options: &CoreOptions) -> Result<()> { + if !core_options.row_tracking_enabled() { + return Err(Error::DataInvalid { + message: "vindex index build requires 'row-tracking.enabled' = 'true'".to_string(), + source: None, + }); + } + if !core_options.data_evolution_enabled() { + return Err(Error::DataInvalid { + message: "vindex index build requires 'data-evolution.enabled' = 'true'".to_string(), + source: None, + }); + } + if !core_options.global_index_enabled() { + return Err(Error::DataInvalid { + message: "vindex index build requires 'global-index.enabled' = 'true'".to_string(), + source: None, + }); + } + if !table.schema().primary_keys().is_empty() { + return Err(Error::Unsupported { + message: "vindex index build does not support primary-key tables".to_string(), + }); + } + if core_options.deletion_vectors_enabled() { + return Err(Error::Unsupported { + message: + "vindex index build does not support tables with deletion-vectors.enabled=true" + .to_string(), + }); + } + Ok(()) +} + +pub(super) fn find_index_field<'a>(table: &'a Table, column: &str) -> Result<&'a DataField> { + find_vector_index_field(table, column) +} + +pub(super) fn validate_vector_field(field: &DataField) -> Result<()> { + validate_common_vector_field(field, "vindex") +} + +pub(super) fn checked_vector_bytes(row_count: usize, dimension: usize) -> Result { + row_count + .checked_mul(dimension) + .and_then(|values| values.checked_mul(std::mem::size_of::())) + .ok_or_else(|| Error::DataInvalid { + message: format!( + "vindex vector byte length overflows: row_count={row_count}, dimension={dimension}" + ), + source: None, + }) +} + +pub(super) fn checked_std_vector_bytes( + row_count: usize, + dimension: usize, +) -> std::io::Result { + row_count + .checked_mul(dimension) + .and_then(|values| values.checked_mul(std::mem::size_of::())) + .ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "vindex vector byte length overflows usize", + ) + }) +} + +pub(super) fn checked_training_vector_count(row_count: usize, ratio: f64) -> Result { + if row_count == 0 || !(ratio > 0.0 && ratio <= 1.0) { + return Err(Error::DataInvalid { + message: format!( + "Invalid vindex training sample: row_count={row_count}, ratio={ratio}; expected a positive row count and ratio in (0, 1]" + ), + source: None, + }); + } + Ok(((row_count as f64 * ratio).ceil() as usize).clamp(1, row_count)) +} + +pub(super) fn checked_training_sample_index( + sample: usize, + rows: usize, + samples: usize, +) -> Result { + sample + .checked_mul(rows / samples) + .and_then(|base| { + sample + .checked_mul(rows % samples) + .and_then(|remainder| base.checked_add(remainder / samples)) + }) + .ok_or_else(|| Error::DataInvalid { + message: "vindex training sample index overflows usize".to_string(), + source: None, + }) +} + +pub(super) fn checked_i32(value: u64, context: &str) -> Result { + i32::try_from(value).map_err(|_| Error::DataInvalid { + message: format!("{context}: {value}"), + source: None, + }) +} + +pub(super) fn checked_i64(value: u64, context: &str) -> Result { + i64::try_from(value).map_err(|_| Error::DataInvalid { + message: format!("{context}: {value}"), + source: None, + }) +} + +pub(super) fn checked_row_count(row_range_start: i64, row_range_end: i64) -> Result { + if row_range_end < row_range_start { + return Err(Error::DataInvalid { + message: format!("Invalid vindex row range [{row_range_start}, {row_range_end}]"), + source: None, + }); + } + row_range_end + .checked_sub(row_range_start) + .and_then(|count| count.checked_add(1)) + .ok_or_else(|| Error::DataInvalid { + message: format!( + "Row count overflows for row range [{row_range_start}, {row_range_end}]" + ), + source: None, + }) +} diff --git a/crates/paimon/src/table/vindex_index_build_builder/writer.rs b/crates/paimon/src/table/vindex_index_build_builder/writer.rs new file mode 100644 index 000000000..fd516466e --- /dev/null +++ b/crates/paimon/src/table/vindex_index_build_builder/writer.rs @@ -0,0 +1,442 @@ +// 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, validate_vector_batch}; +use super::planning::VindexIndexShard; +use super::timing::{vector_index_build_timing_enabled, VectorIndexBuildTiming}; +use super::validation::{ + checked_i64, checked_row_count, checked_std_vector_bytes, checked_training_sample_index, + checked_training_vector_count, checked_vector_bytes, +}; +use super::VindexIndexBuildBuilder; +use crate::spec::{GlobalIndexMeta, IndexFileMeta, ROW_ID_FIELD_NAME}; +use crate::table::data_file_reader::DataFileReadTiming; +use crate::table::table_read::configured_parquet_read_budget; +use crate::vindex::VindexVectorIndexOptions; +use crate::{Error, Result}; +use arrow_buffer::MutableBuffer; +use futures::TryStreamExt; +use paimon_vindex_core::autotune::default_training_vector_count; +use paimon_vindex_core::index::{VectorIndexTrainer, VectorIndexWriter}; +use paimon_vindex_core::io::PosWriter; +use std::io::{Read, Seek, SeekFrom}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::io::AsyncWriteExt; +use tokio_util::io::SyncIoBridge; + +const INDEX_DIR: &str = "index"; +const VECTOR_BUFFER_BYTES: usize = 8 * 1024 * 1024; +pub(super) struct BuiltIndexFile { + pub(super) meta: IndexFileMeta, + pub(super) timing: Option, +} + +impl<'a> VindexIndexBuildBuilder<'a> { + pub(super) async fn build_index_file( + &self, + shard: &VindexIndexShard, + index_column: &str, + dimension: i32, + index_field_id: i32, + options: &VindexVectorIndexOptions, + index_meta: Vec, + ) -> Result { + let timing_enabled = vector_index_build_timing_enabled(); + let total_start = timing_enabled.then(Instant::now); + let mut source_batch_wait = Duration::ZERO; + let mut raw_temp_write = Duration::ZERO; + let read_timing = timing_enabled.then(|| Arc::new(DataFileReadTiming::default())); + let parquet_read_budget = if timing_enabled { + let budget = configured_parquet_read_budget(self.table)?; + budget.enable_diagnostics(); + Some(budget) + } else { + None + }; + let mut batch_count = 0usize; + let row_count = checked_row_count(shard.row_range_start, shard.row_range_end)?; + let row_count_usize = usize::try_from(row_count).map_err(|e| Error::DataInvalid { + message: format!("Invalid vindex row count: {row_count}"), + source: Some(Box::new(e)), + })?; + let dimension_usize = usize::try_from(dimension).map_err(|e| Error::DataInvalid { + message: format!("Invalid vindex dimension: {dimension}"), + source: Some(Box::new(e)), + })?; + if dimension_usize == 0 { + return Err(Error::DataInvalid { + message: "vindex vector dimension must be positive".to_string(), + source: None, + }); + } + let expected_bytes = checked_vector_bytes(row_count_usize, dimension_usize)?; + let training_vector_count = + checked_training_vector_count(row_count_usize, options.train_sample_ratio)?; + let training_buffer_rows = + (VECTOR_BUFFER_BYTES / checked_vector_bytes(1, dimension_usize)?).max(1); + let training_buffer_floats = training_buffer_rows + .checked_mul(dimension_usize) + .ok_or_else(|| Error::DataInvalid { + message: "vindex training buffer length overflows usize".to_string(), + source: None, + })?; + + let mut trainer = + VectorIndexTrainer::new(options.config.clone()).map_err(|e| Error::DataInvalid { + message: format!("Failed to initialize vindex trainer: {e}"), + source: Some(Box::new(e)), + })?; + let raw_file = tempfile::tempfile().map_err(|e| Error::UnexpectedError { + message: format!("Failed to create temporary vindex vector file: {e}"), + source: Some(Box::new(e)), + })?; + let mut raw_file = tokio::fs::File::from_std(raw_file); + let split = data_split_for_shard(shard)?; + let mut read_builder = self.table.new_read_builder(); + read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; + let read = read_builder.new_read()?; + let read = match read_timing.as_ref() { + Some(timing) => read.with_data_file_read_timing(Arc::clone(timing)), + None => read, + }; + let read = match parquet_read_budget.as_ref() { + Some(budget) => read.with_parquet_read_budget(Arc::clone(budget)), + None => read, + }; + let mut batches = read.to_arrow(&[split])?; + let mut expected_row_id = shard.row_range_start; + let mut rows_seen = 0usize; + let mut bytes_written = 0usize; + let mut next_training_sample = 0usize; + let mut training_buffer = Vec::with_capacity(training_buffer_floats); + + loop { + let source_start = timing_enabled.then(Instant::now); + let batch = batches.try_next().await?; + if let Some(source_start) = source_start { + source_batch_wait = source_batch_wait.saturating_add(source_start.elapsed()); + } + let Some(batch) = batch else { break }; + batch_count += 1; + let vectors = + validate_vector_batch(&batch, index_column, dimension_usize, &mut expected_row_id)?; + let batch_end = + rows_seen + .checked_add(vectors.row_count) + .ok_or_else(|| Error::DataInvalid { + message: "vindex streamed row count overflows usize".to_string(), + source: None, + })?; + + if training_vector_count == row_count_usize { + trainer + .add_training_vectors_mut(vectors.values, vectors.row_count) + .map_err(|e| Error::DataInvalid { + message: format!("Failed to add vindex training vectors: {e}"), + source: Some(Box::new(e)), + })?; + } else { + while next_training_sample < training_vector_count { + let sample_row = checked_training_sample_index( + next_training_sample, + row_count_usize, + training_vector_count, + )?; + if sample_row >= batch_end { + break; + } + let start = (sample_row - rows_seen) * dimension_usize; + training_buffer + .extend_from_slice(&vectors.values[start..start + dimension_usize]); + next_training_sample += 1; + if training_buffer.len() == training_buffer_floats { + trainer + .add_training_vectors_mut( + &training_buffer, + training_buffer.len() / dimension_usize, + ) + .map_err(|e| Error::DataInvalid { + message: format!("Failed to add vindex training vectors: {e}"), + source: Some(Box::new(e)), + })?; + training_buffer.clear(); + } + } + } + + let raw_write_start = timing_enabled.then(Instant::now); + raw_file + .write_all(vectors.bytes) + .await + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to spill vindex vectors: {e}"), + source: Some(Box::new(e)), + })?; + if let Some(raw_write_start) = raw_write_start { + raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); + } + bytes_written = bytes_written + .checked_add(vectors.bytes.len()) + .ok_or_else(|| Error::DataInvalid { + message: "vindex spilled byte count overflows usize".to_string(), + source: None, + })?; + rows_seen = batch_end; + } + + if !training_buffer.is_empty() { + trainer + .add_training_vectors_mut(&training_buffer, training_buffer.len() / dimension_usize) + .map_err(|e| Error::DataInvalid { + message: format!("Failed to add vindex training vectors: {e}"), + source: Some(Box::new(e)), + })?; + } + if rows_seen != row_count_usize + || expected_row_id + != shard + .row_range_end + .checked_add(1) + .ok_or_else(|| Error::DataInvalid { + message: "vindex row range end overflows i64".to_string(), + source: None, + })? + || (training_vector_count != row_count_usize + && next_training_sample != training_vector_count) + || bytes_written != expected_bytes + { + return Err(Error::DataInvalid { + message: format!( + "vindex streamed data mismatch: rows={rows_seen}/{row_count_usize}, training={next_training_sample}/{training_vector_count}, bytes={bytes_written}/{expected_bytes}" + ), + source: None, + }); + } + let raw_write_start = timing_enabled.then(Instant::now); + raw_file.flush().await.map_err(|e| Error::UnexpectedError { + message: format!("Failed to flush temporary vindex vector file: {e}"), + source: Some(Box::new(e)), + })?; + if let Some(raw_write_start) = raw_write_start { + raw_temp_write = raw_temp_write.saturating_add(raw_write_start.elapsed()); + } + let raw_file_len = raw_file + .metadata() + .await + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to inspect temporary vindex vector file: {e}"), + source: Some(Box::new(e)), + })? + .len(); + if raw_file_len != expected_bytes as u64 { + return Err(Error::DataInvalid { + message: format!( + "temporary vindex vector file size mismatch: {raw_file_len}/{expected_bytes}" + ), + source: None, + }); + } + let raw_file = raw_file.into_std().await; + // Diagnostics only: never fail the build for a timing log field. + let training_rows_retained = if timing_enabled { + default_training_vector_count(training_vector_count, options.config.nlist()) + .unwrap_or(0) + } else { + 0 + }; + + let (writer, train_finish, raw_temp_reread, index_add) = tokio::task::spawn_blocking( + move || -> std::io::Result<(VectorIndexWriter, Duration, Duration, Duration)> { + let train_start = timing_enabled.then(Instant::now); + let training = trainer.finish()?; + let train_finish = train_start.map_or(Duration::ZERO, |start| start.elapsed()); + let mut writer = VectorIndexWriter::new(training); + let mut raw_temp_reread = Duration::ZERO; + let mut index_add = Duration::ZERO; + let mut raw_file = raw_file; + let reread_start = timing_enabled.then(Instant::now); + raw_file.seek(SeekFrom::Start(0))?; + if let Some(start) = reread_start { + raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); + } + let batch_rows = training_buffer_rows.min(row_count_usize); + let batch_bytes = checked_std_vector_bytes(batch_rows, dimension_usize)?; + let mut buffer = MutableBuffer::new(batch_bytes); + let mut ids = Vec::with_capacity(batch_rows); + let mut rows_added = 0usize; + while rows_added < row_count_usize { + let rows = batch_rows.min(row_count_usize - rows_added); + buffer.resize(checked_std_vector_bytes(rows, dimension_usize)?, 0); + let reread_start = timing_enabled.then(Instant::now); + raw_file.read_exact(buffer.as_slice_mut())?; + if let Some(start) = reread_start { + raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); + } + ids.clear(); + for row in rows_added..rows_added + rows { + ids.push(i64::try_from(row).map_err(|_| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + "vindex row id does not fit i64", + ) + })?); + } + let add_start = timing_enabled.then(Instant::now); + writer.add_vectors(&ids, buffer.typed_data::(), rows)?; + if let Some(start) = add_start { + index_add = index_add.saturating_add(start.elapsed()); + } + rows_added += rows; + } + let mut trailing = [0u8; 1]; + let reread_start = timing_enabled.then(Instant::now); + if raw_file.read(&mut trailing)? != 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "temporary vindex vector file contains trailing bytes", + )); + } + if let Some(start) = reread_start { + raw_temp_reread = raw_temp_reread.saturating_add(start.elapsed()); + } + Ok((writer, train_finish, raw_temp_reread, index_add)) + }, + ) + .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 or add vectors to vindex index: {e}"), + source: Some(Box::new(e)), + })?; + + let serialize_upload_start = timing_enabled.then(Instant::now); + self.table + .file_io() + .mkdirs(&format!( + "{}/{INDEX_DIR}/", + self.table.location().trim_end_matches('/') + )) + .await?; + let file_name = format!( + "vector-{}-global-index-{}.index", + self.index_type, + uuid::Uuid::new_v4() + ); + let index_path = format!( + "{}/{INDEX_DIR}/{}", + self.table.location().trim_end_matches('/'), + file_name + ); + let write_result = async { + let async_writer = self + .table + .file_io() + .new_output(&index_path)? + .async_writer() + .await?; + let mut output = SyncIoBridge::new(async_writer); + tokio::task::spawn_blocking(move || -> std::io::Result<()> { + let mut writer = writer; + writer.write(&mut PosWriter::new(&mut output))?; + output.shutdown() + }) + .await + .map_err(|e| Error::UnexpectedError { + message: format!("vindex serialization task failed: {e}"), + source: None, + })? + .map_err(|e| Error::UnexpectedError { + message: format!("Failed to stream vindex index: {e}"), + source: Some(Box::new(e)), + })?; + self.table.file_io().get_status(&index_path).await + } + .await; + let status = match write_result { + Ok(status) => status, + Err(error) => { + let _ = self.table.file_io().delete_file(&index_path).await; + return Err(error); + } + }; + let serialize_upload = + serialize_upload_start.map_or(Duration::ZERO, |start| start.elapsed()); + let meta = IndexFileMeta { + index_type: self.index_type.clone(), + file_name: file_name.clone(), + file_size: checked_i64( + status.size, + "Index file is too large for Rust IndexFileMeta", + )?, + row_count, + deletion_vectors_ranges: None, + global_index_meta: Some(GlobalIndexMeta { + row_range_start: shard.row_range_start, + row_range_end: shard.row_range_end, + index_field_id, + extra_field_ids: None, + source_meta: None, + index_meta: Some(index_meta), + }), + }; + let (oss_read, parquet_decode) = read_timing + .as_ref() + .map_or((Duration::ZERO, Duration::ZERO), |timing| { + (timing.file_read(), timing.parquet_decode()) + }); + let (file_schema_open, first_batch_wait, remaining_batch_wait) = read_timing + .as_ref() + .map_or((Duration::ZERO, Duration::ZERO, Duration::ZERO), |timing| { + timing.file_waits() + }); + let parquet_diagnostics = parquet_read_budget + .as_ref() + .map_or_else(Default::default, |budget| budget.diagnostics()); + let timing = total_start.map(|start| VectorIndexBuildTiming { + total_without_commit: start.elapsed(), + source_batch_wait, + oss_read, + parquet_decode, + file_schema_open, + first_batch_wait, + remaining_batch_wait, + parquet_row_group_count: parquet_diagnostics.row_group_count, + parquet_projected_bytes_min: parquet_diagnostics.projected_bytes_min, + parquet_projected_bytes_max: parquet_diagnostics.projected_bytes_max, + parquet_projected_bytes_total: parquet_diagnostics.projected_bytes_total, + parquet_peak_inflight_row_groups: parquet_diagnostics.peak_inflight, + raw_temp_write, + train_finish, + raw_temp_reread, + index_add, + serialize_upload, + rows: row_count_usize, + training_rows_seen: training_vector_count, + training_rows_retained, + batch_count, + raw_temp_bytes: bytes_written, + index_bytes: status.size, + data_file_count: shard.files.len(), + file_name, + }); + Ok(BuiltIndexFile { meta, timing }) + } +} From ccc8761d7712e7c4219ab1cd1a5c34ee8a5e69c6 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Tue, 1 Sep 2026 10:54:42 +0800 Subject: [PATCH 2/2] fix: satisfy Rust 1.98 suffix array lint --- crates/paimon/src/fm_index/suffix_array.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/paimon/src/fm_index/suffix_array.rs b/crates/paimon/src/fm_index/suffix_array.rs index aa18983a6..b40f4493b 100644 --- a/crates/paimon/src/fm_index/suffix_array.rs +++ b/crates/paimon/src/fm_index/suffix_array.rs @@ -297,10 +297,10 @@ mod tests { for length in 1..=80 { for _ in 0..30 { let alphabet = random.gen_range(1..=12); - let mut text = (0..length) + let text = (0..length) .map(|_| random.gen_range(1..=alphabet) as u16) + .chain(std::iter::once(0)) .collect::>(); - text.push(0); assert_eq!(build(&text, alphabet).unwrap(), naive(&text)); } }