diff --git a/Cargo.lock b/Cargo.lock index 2723915836c..87d84101e6e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4550,6 +4550,7 @@ dependencies = [ "serde", "serde_json", "serial_test", + "sha2 0.10.9", "snafu", "tempfile", "test-log", diff --git a/Cargo.toml b/Cargo.toml index 11de38ea2d6..ed2a419ba06 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -215,6 +215,7 @@ serde = { version = "^1" } serde_json = { version = "1" } semver = "1.0" serial_test = "3" +sha2 = "0.10" snafu = "0.9" syn = { version = "2.0.37", features = ["full"] } lindera = { version = "3.0.7" } diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock index 2b1502f1c45..7a7986a7150 100644 --- a/java/lance-jni/Cargo.lock +++ b/java/lance-jni/Cargo.lock @@ -3888,6 +3888,7 @@ dependencies = [ "semver", "serde", "serde_json", + "sha2 0.10.9", "snafu", "tokio", "tokio-stream", diff --git a/python/Cargo.lock b/python/Cargo.lock index c3656d9f9fe..2bea9dd1e0b 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -4056,6 +4056,7 @@ dependencies = [ "semver", "serde", "serde_json", + "sha2 0.10.9", "snafu", "tokio", "tokio-stream", diff --git a/rust/lance-index-core/src/lib.rs b/rust/lance-index-core/src/lib.rs index 4e59dce39d6..0d7bae434b0 100644 --- a/rust/lance-index-core/src/lib.rs +++ b/rust/lance-index-core/src/lib.rs @@ -224,6 +224,11 @@ impl IndexType { /// Indices which higher version than these will be ignored for compatibility, /// This would happen when creating index in a newer version of Lance, /// but then opening the index in older version of Lance + /// + /// For [`Self::IvfRq`] this is only the highest version this build reads: + /// an IVF_RQ writer commits the version of the row layout it writes + /// (`lance_index::ivf_rq_index_version`), so a column-layout index keeps + /// the version older readers accept. pub fn version(&self) -> i32 { match self { Self::Scalar => 0, @@ -253,7 +258,8 @@ impl IndexType { | Self::IvfHnswSq | Self::IvfHnswPq | Self::IvfHnswFlat => 1, - Self::IvfRq => 2, + // 2: the column layout; 3: the plane-row layout. + Self::IvfRq => 3, } } diff --git a/rust/lance-index/benches/rq.rs b/rust/lance-index/benches/rq.rs index 03197a40c6a..5d29431633d 100644 --- a/rust/lance-index/benches/rq.rs +++ b/rust/lance-index/benches/rq.rs @@ -253,6 +253,7 @@ fn ex_code_storage_load(c: &mut Criterion) { num_bits: NUM_BITS, packed: true, layered: false, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, }; let code_len = LOAD_DIM / 8; @@ -678,6 +679,7 @@ fn build_transform(c: &mut Criterion) { num_bits, packed: false, layered: false, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, }; let rq = RabitQuantizer::try_from( diff --git a/rust/lance-index/src/lib.rs b/rust/lance-index/src/lib.rs index 7da42115f46..01c77eaef99 100644 --- a/rust/lance-index/src/lib.rs +++ b/rust/lance-index/src/lib.rs @@ -43,8 +43,23 @@ pub const INDEX_METADATA_SCHEMA_KEY: &str = "lance:index"; /// Most vector indices should use this version unless they need to bump for a /// format change. pub const VECTOR_INDEX_VERSION: u32 = 1; -/// Version for IVF_RQ indices. -pub const IVF_RQ_INDEX_VERSION: u32 = 2; +/// Version of IVF_RQ indices in the plane-row layout, and the highest IVF_RQ +/// version this build reads. +pub const IVF_RQ_INDEX_VERSION: u32 = 3; +/// Version of IVF_RQ indices in the column layout, which readers that predate +/// the plane-row layout accept. +pub const IVF_RQ_COLUMNS_INDEX_VERSION: u32 = 2; + +/// The index version an IVF_RQ segment whose auxiliary file is in `row_layout` +/// is committed with. Every writer (build, append, optimize, conversion) takes +/// it from the layout it writes, so a column-layout index never takes the +/// plane-row version, which older readers leave out of their usable indexes. +pub fn ivf_rq_index_version(row_layout: vector::bq::RQRowLayout) -> i32 { + match row_layout { + vector::bq::RQRowLayout::Columns => IVF_RQ_COLUMNS_INDEX_VERSION as i32, + vector::bq::RQRowLayout::PlaneRows => IVF_RQ_INDEX_VERSION as i32, + } +} /// The factor of threshold to trigger split / join for vector index. /// @@ -108,6 +123,27 @@ mod tests { assert_eq!(IndexType::max_vector_version(), IVF_RQ_INDEX_VERSION); } + #[test] + fn test_ivf_rq_index_version_follows_row_layout() { + use crate::vector::bq::RQRowLayout; + assert_eq!( + ivf_rq_index_version(RQRowLayout::Columns), + IVF_RQ_COLUMNS_INDEX_VERSION as i32 + ); + assert_eq!( + ivf_rq_index_version(RQRowLayout::PlaneRows), + IVF_RQ_INDEX_VERSION as i32 + ); + // Plane rows are the read ceiling; columns keep the version readers + // without plane-row support accept, which still exceeds IVF_PQ's. + assert_eq!( + ivf_rq_index_version(RQRowLayout::PlaneRows), + IndexType::IvfRq.version() + ); + assert!(ivf_rq_index_version(RQRowLayout::Columns) < IndexType::IvfRq.version()); + assert!(ivf_rq_index_version(RQRowLayout::Columns) > IndexType::IvfPq.version()); + } + #[test] fn test_ivf_rq_target_partition_size() { assert_eq!(IndexType::IvfRq.target_partition_size(), 4096); diff --git a/rust/lance-index/src/vector.rs b/rust/lance-index/src/vector.rs index 5dd3205f77d..468db59da16 100644 --- a/rust/lance-index/src/vector.rs +++ b/rust/lance-index/src/vector.rs @@ -28,6 +28,7 @@ use v3::subindex::SubIndexType; pub mod bq; pub mod distributed; +pub(crate) mod exact_buffers; pub mod flat; pub mod graph; pub mod hnsw; diff --git a/rust/lance-index/src/vector/bq.rs b/rust/lance-index/src/vector/bq.rs index 0a6abe98866..b957b13e69e 100644 --- a/rust/lance-index/src/vector/bq.rs +++ b/rust/lance-index/src/vector/bq.rs @@ -14,6 +14,7 @@ use lance_core::{Error, Result}; use num_traits::Float; use serde::{Deserialize, Serialize}; +pub use crate::vector::bq::storage::RQRowLayout; use crate::vector::bq::storage::RabitQuantizationMetadata; use crate::vector::quantizer::QuantizerBuildParams; @@ -27,6 +28,7 @@ mod lazy_tests; pub(crate) mod pairwise; pub mod partition_codes; mod plane_cache; +pub mod plane_rows; pub mod prune; pub mod raw_body; pub mod resident; @@ -133,6 +135,10 @@ pub struct RQBuildParams { pub rotation: Option, /// Store prefix code planes; supported for 3 through 9 bits. Default off. pub layered: bool, + /// How the auxiliary file stores each row's fields. Default columns; the + /// plane-row layout is committed as index version 3, which readers + /// without plane-row support leave out of their usable indexes. + pub row_layout: RQRowLayout, } pub fn validate_rq_num_bits(num_bits: u8) -> Result<()> { @@ -175,12 +181,20 @@ impl RQBuildParams { self } + /// Write the auxiliary file in `row_layout`. Plane rows require a reader + /// with plane-row support. + pub fn with_row_layout(mut self, row_layout: RQRowLayout) -> Self { + self.row_layout = row_layout; + self + } + pub fn new(num_bits: u8) -> Self { Self { num_bits, rotation_type: RQRotationType::default(), rotation: None, layered: false, + row_layout: RQRowLayout::Columns, } } @@ -190,6 +204,7 @@ impl RQBuildParams { rotation_type, rotation: None, layered: false, + row_layout: RQRowLayout::Columns, } } } @@ -200,7 +215,10 @@ impl From<&RQBuildParams> for RabitQuantization { Self { num_bits: value.num_bits as u32, layered: value.layered, - row_layout: RowLayout::Columns as i32, + row_layout: match value.row_layout { + RQRowLayout::Columns => RowLayout::Columns as i32, + RQRowLayout::PlaneRows => RowLayout::PlaneRows as i32, + }, rotation_type: match value.rotation_type { RQRotationType::Fast => RotationType::Fast as i32, RQRotationType::Matrix => RotationType::Matrix as i32, @@ -222,6 +240,7 @@ impl Default for RQBuildParams { rotation_type: RQRotationType::default(), rotation: None, layered: false, + row_layout: RQRowLayout::Columns, } } } @@ -268,6 +287,21 @@ mod tests { assert_eq!(RQBuildParams::default().num_bits, 5); } + #[test] + fn test_rq_build_params_row_layout_reaches_details() { + use crate::pb::vector_index_details::rabit_quantization::RowLayout; + assert_eq!(RQBuildParams::default().row_layout, RQRowLayout::Columns); + assert_eq!(RQBuildParams::new(7).row_layout, RQRowLayout::Columns); + let columns = RabitQuantization::from(&RQBuildParams::new(7)); + assert_eq!(columns.row_layout(), RowLayout::Columns); + let params = RQBuildParams::new(7) + .with_layered(true) + .with_row_layout(RQRowLayout::PlaneRows); + let plane_rows = RabitQuantization::from(¶ms); + assert_eq!(plane_rows.row_layout(), RowLayout::PlaneRows); + assert!(plane_rows.layered); + } + #[test] fn test_rabit_num_bits_validation() { validate_rq_num_bits(1).unwrap(); diff --git a/rust/lance-index/src/vector/bq/builder.rs b/rust/lance-index/src/vector/bq/builder.rs index ade108a2cf6..8d14be17781 100644 --- a/rust/lance-index/src/vector/bq/builder.rs +++ b/rust/lance-index/src/vector/bq/builder.rs @@ -17,8 +17,8 @@ use rand_distr::Distribution; use rayon::prelude::*; use crate::vector::bq::storage::{ - RABIT_CODE_COLUMN, RABIT_METADATA_KEY, RabitQuantizationMetadata, RabitQuantizationStorage, - RabitQueryEstimator, rabit_binary_code_field, rabit_ex_code_field, + RABIT_CODE_COLUMN, RABIT_METADATA_KEY, RQRowLayout, RabitQuantizationMetadata, + RabitQuantizationStorage, RabitQueryEstimator, rabit_binary_code_field, rabit_ex_code_field, }; use crate::vector::bq::transform::{ ADD_FACTORS_FIELD, ERROR_FACTORS_FIELD, EX_ADD_FACTORS_FIELD, EX_SCALE_FACTORS_FIELD, @@ -578,6 +578,7 @@ impl RabitQuantizer { num_bits, packed: false, layered: false, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, } } @@ -590,6 +591,7 @@ impl RabitQuantizer { num_bits, packed: false, layered: false, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, }, }; @@ -1061,6 +1063,7 @@ impl Quantization for RabitQuantizer { } if let Some(mut q) = Self::from_supplied_rotation(params, dim)? { q.metadata.layered = params.layered; + q.metadata.row_layout = params.row_layout; return Ok(q); } @@ -1088,6 +1091,7 @@ impl Quantization for RabitQuantizer { } }; q.metadata.layered = params.layered; + q.metadata.row_layout = params.row_layout; Ok(q) } @@ -1168,25 +1172,38 @@ impl Quantization for RabitQuantizer { } fn extra_fields(&self) -> Vec { - let mut fields = vec![ADD_FACTORS_FIELD.clone(), SCALE_FACTORS_FIELD.clone()]; - if self.metadata.query_estimator == RabitQueryEstimator::RawQuery { - fields.push(ERROR_FACTORS_FIELD.clone()); - } - if self.metadata.layered { - return super::layered::storage_fields(self.code_dim(), self.metadata.num_bits, fields) - .expect("layered layout validated at build"); - } - if let Some(ex_code_field) = rabit_ex_code_field(self.code_dim(), self.metadata.num_bits) - .expect("RabitQ num_bits should be validated") - { - fields.push(ex_code_field); - fields.push(EX_ADD_FACTORS_FIELD.clone()); - fields.push(EX_SCALE_FACTORS_FIELD.clone()); - } + let mut fields = rabit_storage_fields(&self.metadata) + .expect("RabitQ num_bits and layered layout validated at build"); + // The first is the binary code field, `Self::field`. + fields.remove(0); fields } } +/// The fields an IVF_RQ storage file with `metadata` stores after the row id +/// in the column layout, in their order: the binary codes and their factors, +/// then the ex codes and factors, as plain or layered planes. +pub fn rabit_storage_fields(metadata: &RabitQuantizationMetadata) -> Result> { + let rotated_dim = metadata.rotated_dim(); + let mut fields = vec![ + rabit_binary_code_field(rotated_dim), + ADD_FACTORS_FIELD.clone(), + SCALE_FACTORS_FIELD.clone(), + ]; + if metadata.query_estimator == RabitQueryEstimator::RawQuery { + fields.push(ERROR_FACTORS_FIELD.clone()); + } + if metadata.layered { + return super::layered::storage_fields(rotated_dim, metadata.num_bits, fields); + } + if let Some(ex_code_field) = rabit_ex_code_field(rotated_dim, metadata.num_bits)? { + fields.push(ex_code_field); + fields.push(EX_ADD_FACTORS_FIELD.clone()); + fields.push(EX_SCALE_FACTORS_FIELD.clone()); + } + Ok(fields) +} + impl TryFrom for RabitQuantizer { type Error = Error; @@ -1693,6 +1710,27 @@ mod tests { assert_eq!(metadata.num_bits, 3); assert_eq!(metadata.rotation_type, RQRotationType::Fast); assert_eq!(metadata.fast_rotation_signs, supplied_signs); + assert_eq!(metadata.row_layout, RQRowLayout::Columns); + } + + /// The build params' row layout reaches the quantizer's metadata, which + /// the writer reads it from, whether the rotation is new or supplied. + #[test] + fn test_rabit_quantizer_takes_row_layout_from_params() { + let vectors = Float32Array::from(vec![0.0f32; 4 * 32]); + let fsl = FixedSizeListArray::try_new_from_values(vectors, 32).unwrap(); + let params = RQBuildParams::new(7) + .with_layered(true) + .with_row_layout(RQRowLayout::PlaneRows); + let quantizer = RabitQuantizer::build(&fsl, DistanceType::L2, ¶ms).unwrap(); + assert_eq!(quantizer.metadata_ref().row_layout, RQRowLayout::PlaneRows); + assert_eq!(quantizer.metadata(None).row_layout, RQRowLayout::PlaneRows); + + // A rotation supplied from a plane-row index takes the params' layout. + let mut params = RQBuildParams::new(7).with_layered(true); + params.rotation = Some(quantizer.metadata(None)); + let quantizer = RabitQuantizer::build(&fsl, DistanceType::L2, ¶ms).unwrap(); + assert_eq!(quantizer.metadata_ref().row_layout, RQRowLayout::Columns); } #[test] diff --git a/rust/lance-index/src/vector/bq/layered_stats.rs b/rust/lance-index/src/vector/bq/layered_stats.rs index 1bfb1ee4d93..a08cdde1de4 100644 --- a/rust/lance-index/src/vector/bq/layered_stats.rs +++ b/rust/lance-index/src/vector/bq/layered_stats.rs @@ -10,7 +10,8 @@ //! counters are the exception: they advance when an IVF_RQ index, layered or //! not, opens with, loads or reads through its resident columns. So is //! `storage_construct_repacks`, which advances when any RaBitQ storage is -//! built from codes it must rewrite. +//! built from codes it must rewrite, and `plane_rows_unpack_*`, which advance +//! on every read of a plane-row IVF_RQ file. //! Readers take deltas with [`snapshot_and_reset`]. use std::sync::atomic::{AtomicU64, Ordering}; @@ -267,8 +268,9 @@ layered_lazy_counters! { lazy_rank0_scored_queries, /// Time from the start of each of those scans until it published the /// scoring of its first probe (rank 0) to the gathers: once the probe - /// was scored or, earlier, once its rows filled the heap. No gather - /// that waits for the threshold or its turn is released before. + /// was scored or, with `LANCE_RQ_LAZY_PARTIAL_PUBLISH=on`, earlier, + /// once its rows filled the heap. No gather that waits for the + /// threshold or its turn is released before. lazy_rank0_scored_ns, /// Lazy scans that issued a gather of a probe that is not /// `certain_dense`, whose issue waits on scoring progress; see @@ -326,6 +328,11 @@ layered_lazy_counters! { resident_attach_bytes, /// Time spent assembling them, copies included. resident_attach_ns, + /// Bytes of packed row columns that reads of plane-row IVF_RQ files + /// unpacked into the column layout's fields. + plane_rows_unpack_bytes, + /// Time spent unpacking them. + plane_rows_unpack_ns, } ranked { /// Probes gathered lazily. @@ -382,6 +389,20 @@ layered_lazy_counters! { stage2_pruned_live, /// Exactly scored survivors outside the query's distance range. stage2_exact_rejected, + /// Lazy probes whose survivors first filled the heap to `k` rows + /// while they were scored, and that published the heap's top then, + /// before the probe was scored whole, which only + /// `LANCE_RQ_LAZY_PARTIAL_PUBLISH=on` does. That threshold is the + /// `k`-th best of the rows scored so far, looser than the probe's + /// final one. + mid_probe_full_publishes, + /// Gathers that selected their survivors against such a mid-probe + /// threshold: issued after it was published and before its probe + /// was scored whole. + partial_threshold_issues, + /// Survivors of those gathers whose ex rows were gathered, part of + /// `rows_fetched`. + partial_threshold_rows, } maxima { /// Largest number of probes still unscored at a gather's issue. diff --git a/rust/lance-index/src/vector/bq/plane_rows.rs b/rust/lance-index/src/vector/bq/plane_rows.rs new file mode 100644 index 00000000000..8c2939a8ab1 --- /dev/null +++ b/rust/lance-index/src/vector/bq/plane_rows.rs @@ -0,0 +1,1445 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! The plane-row layout of IVF_RQ auxiliary files. +//! +//! The column layout stores each field of an IVF_RQ storage row (row id, +//! codes, factors, bounds) in a column of its own, so reading one plane of a +//! partition reads one byte range per field. The plane-row layout packs the +//! fields of each plane into one fixed-width row column instead, so a plane of +//! a partition, or a run of its rows, is one contiguous byte range. Both hold +//! the same values; readers unpack plane rows into the column layout's fields, +//! and nothing downstream depends on the layout. See the format +//! documentation's "RaBitQ plane-row layout" for the binding specification. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use std::time::Instant; + +use arrow_array::cast::AsArray; +use arrow_array::types::{Float32Type, UInt8Type, UInt64Type}; +use arrow_array::{ + Array, ArrayRef, FixedSizeListArray, Float32Array, RecordBatch, RecordBatchOptions, UInt8Array, + UInt64Array, new_empty_array, +}; +use arrow_schema::{DataType, Field, FieldRef, Schema, SchemaRef}; +use lance_core::{Error, ROW_ID, ROW_ID_FIELD, Result}; +use lance_encoding::constants::{ + COMPRESSION_META_KEY, STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, +}; +use lance_file::version::ConcreteFileVersion; + +use super::builder::rabit_storage_fields; +use super::layered::{ + FULL_BOUNDS_COLUMN, HIGH_ADD_FACTORS_COLUMN, HIGH_BOUNDS_COLUMN, HIGH_SCALE_FACTORS_COLUMN, +}; +use super::layered_stats; +use super::storage::{ + RABIT_BLOCKED_EX_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_LO_COLUMN, RABIT_CODE_COLUMN, + RABIT_EX_CODE_COLUMN, RabitQuantizationMetadata, +}; +use super::transform::{ + ADD_FACTORS_COLUMN, ERROR_FACTORS_COLUMN, EX_ADD_FACTORS_COLUMN, EX_SCALE_FACTORS_COLUMN, + SCALE_FACTORS_COLUMN, +}; + +/// A native index's packed column: the whole row. +pub const RQ_ROWS_COLUMN: &str = "__rq_rows"; +/// A layered index's sign plane: row id, sign codes and their factors. +pub const RQ_SIGN_ROWS_COLUMN: &str = "__rq_sign_rows"; +/// A layered index's estimator bounds. +pub const RQ_BOUNDS_ROWS_COLUMN: &str = "__rq_bounds_rows"; +/// A layered index's high ex plane. +pub const RQ_HIGH_ROWS_COLUMN: &str = "__rq_high_rows"; +/// A layered index's low ex plane. +pub const RQ_LOW_ROWS_COLUMN: &str = "__rq_low_rows"; + +/// Every packed column name. A column-layout file has none of them. +pub const PACKED_COLUMNS: [&str; 5] = [ + RQ_ROWS_COLUMN, + RQ_SIGN_ROWS_COLUMN, + RQ_BOUNDS_ROWS_COLUMN, + RQ_HIGH_ROWS_COLUMN, + RQ_LOW_ROWS_COLUMN, +]; + +/// A packed column and the fields it packs, in their order within a row. +/// A field the index does not store is left out. +type PackedGroup = (&'static str, &'static [&'static str]); + +/// A native index's one packed column. +const NATIVE_GROUPS: [PackedGroup; 1] = [( + RQ_ROWS_COLUMN, + &[ + ROW_ID, + RABIT_CODE_COLUMN, + ADD_FACTORS_COLUMN, + SCALE_FACTORS_COLUMN, + ERROR_FACTORS_COLUMN, + RABIT_BLOCKED_EX_CODE_COLUMN, + EX_ADD_FACTORS_COLUMN, + EX_SCALE_FACTORS_COLUMN, + ], +)]; + +/// A layered index's packed columns, one per plane, in file order. +const LAYERED_GROUPS: [PackedGroup; 4] = [ + ( + RQ_SIGN_ROWS_COLUMN, + &[ + ROW_ID, + RABIT_CODE_COLUMN, + ADD_FACTORS_COLUMN, + SCALE_FACTORS_COLUMN, + ERROR_FACTORS_COLUMN, + ], + ), + ( + RQ_BOUNDS_ROWS_COLUMN, + &[HIGH_BOUNDS_COLUMN, FULL_BOUNDS_COLUMN], + ), + ( + RQ_HIGH_ROWS_COLUMN, + &[ + RABIT_BLOCKED_EX_CODE_COLUMN, + HIGH_ADD_FACTORS_COLUMN, + HIGH_SCALE_FACTORS_COLUMN, + ], + ), + ( + RQ_LOW_ROWS_COLUMN, + &[ + RABIT_BLOCKED_EX_CODE_LO_COLUMN, + EX_ADD_FACTORS_COLUMN, + EX_SCALE_FACTORS_COLUMN, + ], + ), +]; + +/// Field metadata value that turns off general compression. +const NO_COMPRESSION: &str = "none"; + +/// Whether `name` is an internal column of an IVF_RQ storage file in either +/// layout, which a plane-row file cannot carry unpacked. +fn is_internal_column(name: &str) -> bool { + name == RABIT_EX_CODE_COLUMN + || PACKED_COLUMNS.contains(&name) + || NATIVE_GROUPS + .iter() + .chain(&LAYERED_GROUPS) + .any(|(_, fields)| fields.contains(&name)) +} + +/// Whether writers of Lance file `version` request the full-zip structural +/// encoding for packed columns: 2.1 and later, where rows narrower than the +/// full-zip cutoff would otherwise be written as mini-blocks, which are read +/// whole. +pub fn requests_fullzip(version: ConcreteFileVersion) -> bool { + !matches!(version, ConcreteFileVersion::V1 | ConcreteFileVersion::V2_0) +} + +/// How a packed field's values are encoded in a row. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ValueKind { + /// A little-endian `uint64`. + U64, + /// A little-endian `float32`. + F32, + /// A `list[n]`: its `n` bytes. + U8List(usize), + /// A `list[n]`: `n` little-endian `float32` values. + F32List(usize), +} + +impl ValueKind { + fn of(field: &Field) -> Result { + let kind = match field.data_type() { + DataType::UInt64 => Some(Self::U64), + DataType::Float32 => Some(Self::F32), + DataType::FixedSizeList(item, size) => { + let size = usize::try_from(*size).ok(); + match (item.data_type(), size) { + (DataType::UInt8, Some(size)) => Some(Self::U8List(size)), + (DataType::Float32, Some(size)) => Some(Self::F32List(size)), + _ => None, + } + } + _ => None, + }; + kind.ok_or_else(|| { + Error::invalid_input(format!( + "column {} of type {} has no plane-row form", + field.name(), + field.data_type() + )) + }) + } + + /// Bytes the value takes in a row. + fn width(self) -> usize { + match self { + Self::U64 => size_of::(), + Self::F32 => size_of::(), + Self::U8List(size) => size, + Self::F32List(size) => size * size_of::(), + } + } +} + +/// One field of a packed column. +#[derive(Debug, Clone)] +struct PackedMember { + /// The field as the column layout stores it. + field: FieldRef, + /// Where its value starts in a row. + offset: usize, + kind: ValueKind, +} + +impl PackedMember { + fn type_error(&self, array: &dyn Array) -> Error { + Error::invalid_input(format!( + "column {} is {}, which does not pack as {}", + self.field.name(), + array.data_type(), + self.field.data_type() + )) + } + + /// Write `array`'s values into the rows of `out`, rows of `stride` bytes. + fn pack_into(&self, array: &dyn Array, out: &mut [u8], stride: usize) -> Result<()> { + let name = self.field.name(); + let rows = out.len() / stride; + if array.len() != rows { + return Err(Error::invalid_input(format!( + "column {name} has {} rows, the batch {rows}", + array.len() + ))); + } + if array.null_count() > 0 { + return Err(Error::invalid_input(format!( + "column {name} has nulls, which the plane-row layout cannot store" + ))); + } + let range = self.offset..self.offset + self.kind.width(); + let rows_out = out.chunks_exact_mut(stride); + match self.kind { + ValueKind::U64 => { + let values = array + .as_primitive_opt::() + .ok_or_else(|| self.type_error(array))?; + for (row, value) in rows_out.zip(values.values()) { + row[range.clone()].copy_from_slice(&value.to_le_bytes()); + } + } + ValueKind::F32 => { + let values = array + .as_primitive_opt::() + .ok_or_else(|| self.type_error(array))?; + for (row, value) in rows_out.zip(values.values()) { + row[range.clone()].copy_from_slice(&value.to_le_bytes()); + } + } + ValueKind::U8List(size) => { + let values = self.list_values(array, size)?; + let values = values + .as_primitive_opt::() + .ok_or_else(|| self.type_error(array))? + .values(); + for (row, value) in rows_out.zip(values.chunks_exact(size)) { + row[range.clone()].copy_from_slice(value); + } + } + ValueKind::F32List(size) => { + let values = self.list_values(array, size)?; + let values = values + .as_primitive_opt::() + .ok_or_else(|| self.type_error(array))? + .values(); + for (row, value) in rows_out.zip(values.chunks_exact(size)) { + for (dst, item) in row[range.clone()] + .as_chunks_mut::<{ size_of::() }>() + .0 + .iter_mut() + .zip(value) + { + dst.copy_from_slice(&item.to_le_bytes()); + } + } + } + } + Ok(()) + } + + /// The items of list column `array`, `size` per row and without nulls. + fn list_values<'a>(&self, array: &'a dyn Array, size: usize) -> Result<&'a ArrayRef> { + let list = array + .as_fixed_size_list_opt() + .filter(|list| list.value_length() as usize == size) + .ok_or_else(|| self.type_error(array))?; + let values = list.values(); + if values.len() != list.len() * size || values.null_count() > 0 { + return Err(Error::invalid_input(format!( + "column {} has null or misaligned list items", + self.field.name() + ))); + } + Ok(values) + } + + /// The field's values in `packed`, slices of `rows` rows of `stride` bytes + /// each, as the column layout stores them: in buffers of exactly their + /// size. + fn unpack(&self, packed: &[&[u8]], stride: usize, rows: usize) -> Result { + let range = self.offset..self.offset + self.kind.width(); + let rows_in = || packed.iter().flat_map(|bytes| bytes.chunks_exact(stride)); + Ok(match self.kind { + ValueKind::U64 => { + let mut values = Vec::with_capacity(rows); + values + .extend(rows_in().map(|row| u64::from_le_bytes(le_bytes(&row[range.clone()])))); + Arc::new(UInt64Array::from(values)) + } + ValueKind::F32 => { + let mut values = Vec::with_capacity(rows); + values + .extend(rows_in().map(|row| f32::from_le_bytes(le_bytes(&row[range.clone()])))); + Arc::new(Float32Array::from(values)) + } + ValueKind::U8List(size) => { + let mut values = Vec::with_capacity(rows * size); + for row in rows_in() { + values.extend_from_slice(&row[range.clone()]); + } + self.list_array(size, Arc::new(UInt8Array::from(values)))? + } + ValueKind::F32List(size) => { + let mut values = Vec::with_capacity(rows * size); + for row in rows_in() { + values.extend( + row[range.clone()] + .as_chunks::<{ size_of::() }>() + .0 + .iter() + .copied() + .map(f32::from_le_bytes), + ); + } + self.list_array(size, Arc::new(Float32Array::from(values)))? + } + }) + } + + /// A list column of this field's type over `values`. + fn list_array(&self, size: usize, values: ArrayRef) -> Result { + let DataType::FixedSizeList(item, _) = self.field.data_type() else { + return Err(Error::internal(format!( + "packed field {} is not a list", + self.field.name() + ))); + }; + Ok(Arc::new(FixedSizeListArray::try_new( + item.clone(), + size as i32, + values, + None, + )?)) + } +} + +/// Whether `read` is `expected` but for the nullability of a list's items, +/// which a Lance file's schema does not keep: the packed columns are written +/// with non-nullable items and read back with nullable ones, so reads reject +/// null items instead (see [`PackedColumn`]). +fn same_type_but_item_nullability(read: &DataType, expected: &DataType) -> bool { + match (read, expected) { + (DataType::FixedSizeList(read_item, read_size), DataType::FixedSizeList(item, size)) => { + read_size == size && read_item.data_type() == item.data_type() + } + _ => read == expected, + } +} + +/// `bytes`, a slice of exactly `N` bytes, as an array. +fn le_bytes(bytes: &[u8]) -> [u8; N] { + let mut out = [0; N]; + out.copy_from_slice(bytes); + out +} + +/// One packed column of a plane-row file: `list[stride]`, each row +/// the concatenation of its fields' values, without padding. +#[derive(Debug, Clone)] +pub struct PackedColumn { + name: &'static str, + stride: usize, + members: Vec, +} + +impl PackedColumn { + pub fn name(&self) -> &'static str { + self.name + } + + /// Bytes per row. + pub fn stride(&self) -> usize { + self.stride + } + + /// The fields it packs, in their order within a row. + pub fn fields(&self) -> impl Iterator { + self.members.iter().map(|member| &member.field) + } + + fn data_type(&self) -> DataType { + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::UInt8, false)), + self.stride as i32, + ) + } + + /// The column's field in a file. With `fullzip`, it requests the + /// full-zip structural encoding and no general compression, so a page + /// holds `rows * stride` bytes and any run of rows is one byte range. + fn file_field(&self, fullzip: bool) -> Field { + let field = Field::new(self.name, self.data_type(), false); + if !fullzip { + return field; + } + field.with_metadata(HashMap::from([ + ( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_FULLZIP.to_string(), + ), + (COMPRESSION_META_KEY.to_string(), NO_COMPRESSION.to_string()), + ])) + } + + /// The rows of this column in `batch`, `stride` bytes each. + fn rows_of<'a>(&self, batch: &'a RecordBatch) -> Result<&'a [u8]> { + let column = batch.column_by_name(self.name).ok_or_else(|| { + Error::internal(format!("a plane-row read is missing column {}", self.name)) + })?; + let rows = column + .as_fixed_size_list_opt() + .filter(|list| list.value_length() as usize == self.stride && list.null_count() == 0) + .and_then(|list| list.values().as_primitive_opt::()) + .filter(|values| values.null_count() == 0) + .ok_or_else(|| { + Error::index(format!( + "plane-row column {} read as {}, expected non-null {}", + self.name, + column.data_type(), + self.data_type() + )) + })? + .values(); + if rows.len() != batch.num_rows() * self.stride { + return Err(Error::index(format!( + "plane-row column {} holds {} bytes for {} rows of {} bytes", + self.name, + rows.len(), + batch.num_rows(), + self.stride + ))); + } + Ok(rows) + } +} + +/// Where a logical field is stored in a plane-row file. +#[derive(Debug, Clone, Copy)] +enum Location { + /// Member `member` of packed column `column`. + Packed { column: usize, member: usize }, + /// Carried column `index`, stored unpacked. + Carried(usize), +} + +/// The plane-row layout of one IVF_RQ storage file: its packed columns, the +/// column layout's fields each packs, and the columns it carries unpacked. +#[derive(Debug, Clone)] +pub struct PlaneRowsSpec { + packed: Vec, + carried: Vec, + /// The column layout's fields in its order, then the carried fields, + /// with the file's schema metadata: what a read unpacks into. + logical: SchemaRef, + locations: HashMap, +} + +impl PlaneRowsSpec { + /// The plane rows of an IVF_RQ storage with `metadata`: the column + /// layout's fields, the row id and [`rabit_storage_fields`], packed as + /// the layout packs them, with no carried columns. + pub fn try_new(metadata: &RabitQuantizationMetadata) -> Result { + let mut fields = vec![Arc::new(ROW_ID_FIELD.clone())]; + fields.extend(rabit_storage_fields(metadata)?.into_iter().map(Arc::new)); + Self::from_parts(fields, metadata.layered, Vec::new(), HashMap::new()) + } + + /// The plane rows of a plane-row file with `metadata` and `file_schema`, + /// which must start with the packed columns the metadata implies, in + /// order and with their strides; every later column is carried. A file + /// whose schema does not match is an invalid index. + pub fn for_file(metadata: &RabitQuantizationMetadata, file_schema: &Schema) -> Result { + let spec = Self::try_new(metadata)?; + let packed = spec.packed.len(); + let carried: Vec = file_schema.fields().iter().skip(packed).cloned().collect(); + let spec = Self::from_parts( + spec.column_fields().cloned().collect(), + metadata.layered, + carried, + file_schema.metadata().clone(), + ) + .map_err(|error| Error::index(format!("invalid plane-row IVF_RQ file: {error}")))?; + spec.validate_file_schema(file_schema)?; + Ok(spec) + } + + /// The plane rows of the column-layout file with `metadata` and + /// `column_schema`, to convert it: the file must store the column + /// layout's fields exactly as [`Self::try_new`] derives them, in that + /// order; every other column is carried. + pub fn for_column_file( + metadata: &RabitQuantizationMetadata, + column_schema: &Schema, + ) -> Result { + let spec = Self::try_new(metadata)?; + let fields: Vec<&FieldRef> = spec.column_fields().collect(); + let stored: Vec<&FieldRef> = column_schema + .fields() + .iter() + .filter(|field| spec.locations.contains_key(field.name())) + .collect(); + if stored != fields { + return Err(Error::invalid_input(format!( + "the column layout's fields are {:?}, the plane-row layout packs {:?}", + stored, fields + ))); + } + let carried = column_schema + .fields() + .iter() + .filter(|field| !spec.locations.contains_key(field.name())) + .cloned() + .collect(); + Self::from_parts( + fields.into_iter().cloned().collect(), + metadata.layered, + carried, + column_schema.metadata().clone(), + ) + } + + fn from_parts( + column_fields: Vec, + layered: bool, + carried: Vec, + metadata: HashMap, + ) -> Result { + let groups: &[PackedGroup] = if layered { + &LAYERED_GROUPS + } else { + &NATIVE_GROUPS + }; + let mut locations = HashMap::new(); + let mut packed = Vec::with_capacity(groups.len()); + for (name, group) in groups { + let mut members = Vec::new(); + let mut stride = 0; + for &member in *group { + let Some(field) = column_fields.iter().find(|field| field.name() == member) else { + continue; + }; + let kind = ValueKind::of(field)?; + locations.insert( + member.to_string(), + Location::Packed { + column: packed.len(), + member: members.len(), + }, + ); + members.push(PackedMember { + field: field.clone(), + offset: stride, + kind, + }); + stride += kind.width(); + } + if members.is_empty() { + return Err(Error::invalid_input(format!( + "packed column {name} would pack no field" + ))); + } + packed.push(PackedColumn { + name, + stride, + members, + }); + } + if let Some(field) = column_fields + .iter() + .find(|field| !locations.contains_key(field.name())) + { + return Err(Error::invalid_input(format!( + "column {} has no plane-row form", + field.name() + ))); + } + for (index, field) in carried.iter().enumerate() { + if is_internal_column(field.name()) { + return Err(Error::invalid_input(format!( + "column {} cannot be carried unpacked in the plane-row layout", + field.name() + ))); + } + if locations + .insert(field.name().clone(), Location::Carried(index)) + .is_some() + { + return Err(Error::invalid_input(format!( + "column {} appears twice", + field.name() + ))); + } + } + let logical = Arc::new(Schema::new_with_metadata( + column_fields + .into_iter() + .chain(carried.iter().cloned()) + .collect::>(), + metadata, + )); + Ok(Self { + packed, + carried, + logical, + locations, + }) + } + + /// The column layout's fields, without the carried ones. + fn column_fields(&self) -> impl Iterator { + self.logical + .fields() + .iter() + .take(self.logical.fields().len() - self.carried.len()) + } + + /// The packed columns, in file order. + pub fn packed_columns(&self) -> &[PackedColumn] { + &self.packed + } + + /// The columns stored unpacked after the packed ones. + pub fn carried_fields(&self) -> &[FieldRef] { + &self.carried + } + + /// What reads unpack into: the column layout's fields in its order, then + /// the carried fields, with the file's schema metadata. + pub fn logical_schema(&self) -> &SchemaRef { + &self.logical + } + + /// The fields of a plane-row file: the packed columns, requesting the + /// full-zip encoding when `fullzip` (see [`requests_fullzip`]), then the + /// carried columns. + pub fn file_fields(&self, fullzip: bool) -> Vec { + self.packed + .iter() + .map(|column| column.file_field(fullzip)) + .chain(self.carried.iter().map(|field| field.as_ref().clone())) + .collect() + } + + /// Check that `schema`, a file's, is this layout's: the packed columns, + /// non-nullable lists of bytes of their strides, then the carried + /// columns. Field metadata is not checked, as a version 2.0 file requests + /// no encoding. + pub fn validate_file_schema(&self, schema: &Schema) -> Result<()> { + let expected = self.file_fields(false); + let matches = schema.fields().len() == expected.len() + && schema + .fields() + .iter() + .zip(&expected) + .all(|(field, expected)| { + field.name() == expected.name() + && same_type_but_item_nullability(field.data_type(), expected.data_type()) + && field.is_nullable() == expected.is_nullable() + }); + if matches { + return Ok(()); + } + let describe = |fields: Vec<(&String, &DataType, bool)>| { + fields + .into_iter() + .map(|(name, data_type, nullable)| { + format!("{name}: {data_type}{}", if nullable { "?" } else { "" }) + }) + .collect::>() + .join(", ") + }; + Err(Error::index(format!( + "invalid plane-row IVF_RQ file: its columns are [{}], the layout needs [{}]", + describe( + schema + .fields() + .iter() + .map(|field| (field.name(), field.data_type(), field.is_nullable())) + .collect() + ), + describe( + expected + .iter() + .map(|field| (field.name(), field.data_type(), field.is_nullable())) + .collect() + ), + ))) + } + + /// The file columns a read of the column layout's `columns` reads: the + /// packed columns that hold them, in file order, then the carried ones. + /// Every packed column read must be wanted whole, since the plane-row + /// layout stores a plane only as all its fields. + pub fn packed_for(&self, columns: &[&str]) -> Result> { + let mut wanted = vec![0usize; self.packed.len()]; + let mut carried = Vec::new(); + let mut seen = HashSet::new(); + for &name in columns { + if !seen.insert(name) { + continue; + } + match self.location(name)? { + Location::Packed { column, .. } => wanted[column] += 1, + Location::Carried(index) => carried.push(index), + } + } + let mut read = Vec::new(); + for (column, count) in self.packed.iter().zip(wanted) { + if count == 0 { + continue; + } + if count != column.members.len() { + return Err(Error::invalid_input(format!( + "a read of {columns:?} wants only some of the fields packed in {}, which holds {:?}", + column.name, + column + .fields() + .map(|field| field.name()) + .collect::>() + ))); + } + read.push(column.name); + } + carried.sort_unstable(); + read.extend( + carried + .into_iter() + .map(|index| self.carried[index].name().as_str()), + ); + Ok(read) + } + + /// The column layout's `columns` in that order, as a read of them + /// returns them, with the file's schema metadata. + pub fn projection(&self, columns: &[&str]) -> Result { + let fields = columns + .iter() + .map(|name| { + self.location(name)?; + Ok(self.logical.field_with_name(name)?.clone()) + }) + .collect::>>()?; + Ok(Arc::new(Schema::new_with_metadata( + fields, + self.logical.metadata().clone(), + ))) + } + + fn location(&self, name: &str) -> Result { + self.locations.get(name).copied().ok_or_else(|| { + Error::invalid_input(format!("the plane-row layout has no column {name}")) + }) + } + + /// `batch`, rows in the column layout, as plane rows: the packed + /// columns, then the carried ones. The batch must hold exactly this + /// layout's fields, by name, without nulls. + pub fn pack(&self, batch: &RecordBatch) -> Result { + if let Some(field) = batch + .schema() + .fields() + .iter() + .find(|field| !self.locations.contains_key(field.name())) + { + return Err(Error::invalid_input(format!( + "column {} has no place in the plane-row layout", + field.name() + ))); + } + let rows = batch.num_rows(); + let mut columns: Vec = Vec::with_capacity(self.packed.len() + self.carried.len()); + for column in &self.packed { + let mut bytes = vec![0u8; rows * column.stride]; + for member in &column.members { + let array = batch.column_by_name(member.field.name()).ok_or_else(|| { + Error::invalid_input(format!( + "packing {}: the batch has no column {}", + column.name, + member.field.name() + )) + })?; + member.pack_into(array.as_ref(), &mut bytes, column.stride)?; + } + columns.push(Arc::new(FixedSizeListArray::try_new( + Arc::new(Field::new("item", DataType::UInt8, false)), + column.stride as i32, + Arc::new(UInt8Array::from(bytes)), + None, + )?)); + } + for field in &self.carried { + let array = batch.column_by_name(field.name()).ok_or_else(|| { + Error::invalid_input(format!("the batch has no carried column {}", field.name())) + })?; + columns.push(array.clone()); + } + let schema = Arc::new(Schema::new(self.file_fields(false))); + Ok(RecordBatch::try_new_with_options( + schema, + columns, + &RecordBatchOptions::new().with_row_count(Some(rows)), + )?) + } + + /// `schema`'s columns, fields of [`Self::logical_schema`], from + /// `batches`, consecutive reads of the packed (and carried) columns + /// that hold them ([`Self::packed_for`]): values in buffers of exactly + /// their size, as the column layout returns them. + pub fn unpack(&self, batches: &[RecordBatch], schema: &SchemaRef) -> Result { + let started = Instant::now(); + let rows: usize = batches.iter().map(RecordBatch::num_rows).sum(); + let mut unpacked = vec![false; self.packed.len()]; + let mut columns = Vec::with_capacity(schema.fields().len()); + for field in schema.fields() { + match self.location(field.name())? { + Location::Packed { column, member } => { + let packed = &self.packed[column]; + let bytes = batches + .iter() + .map(|batch| packed.rows_of(batch)) + .collect::>>()?; + columns.push(packed.members[member].unpack(&bytes, packed.stride, rows)?); + unpacked[column] = true; + } + Location::Carried(_) => { + let arrays = batches + .iter() + .map(|batch| { + batch.column_by_name(field.name()).ok_or_else(|| { + Error::internal(format!( + "a plane-row read is missing carried column {}", + field.name() + )) + }) + }) + .collect::>>()?; + columns.push(match arrays.as_slice() { + [] => new_empty_array(field.data_type()), + [array] => (*array).clone(), + arrays => arrow_select::concat::concat( + &arrays + .iter() + .map(|array| array.as_ref()) + .collect::>(), + )?, + }); + } + } + } + let batch = RecordBatch::try_new_with_options( + schema.clone(), + columns, + &RecordBatchOptions::new().with_row_count(Some(rows)), + )?; + let stats = layered_stats::counters(); + stats.plane_rows_unpack_bytes.add( + self.packed + .iter() + .zip(unpacked) + .filter(|(_, unpacked)| *unpacked) + .map(|(column, _)| (rows * column.stride) as u64) + .sum(), + ); + stats.plane_rows_unpack_ns.add_elapsed(started); + Ok(batch) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow_array::builder::Float32Builder; + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + use rstest::rstest; + + use crate::vector::bq::RQRotationType; + use crate::vector::bq::layered::{SignBounds, plane_columns}; + use crate::vector::bq::storage::{ + RQRowLayout, RabitQueryEstimator, pack_codes, take_packed_codes, + }; + use lance_arrow::RecordBatchExt; + use lance_core::deepsize::{Context, DeepSizeOf}; + + fn metadata( + num_bits: u8, + dim: usize, + layered: bool, + query_estimator: RabitQueryEstimator, + ) -> RabitQuantizationMetadata { + RabitQuantizationMetadata { + rotate_mat: None, + rotate_mat_position: None, + fast_rotation_signs: None, + rotation_type: RQRotationType::Fast, + code_dim: dim as u32, + num_bits, + packed: true, + layered, + row_layout: RQRowLayout::PlaneRows, + query_estimator, + } + } + + /// `rows` rows of `schema` with random values, none null. + fn random_batch(schema: &SchemaRef, rows: usize, seed: u64) -> RecordBatch { + let mut rng = StdRng::seed_from_u64(seed); + let columns = schema + .fields() + .iter() + .map(|field| -> ArrayRef { + match field.data_type() { + DataType::UInt64 => Arc::new(UInt64Array::from_iter_values( + (0..rows).map(|_| rng.random()), + )), + DataType::Float32 => Arc::new(Float32Array::from_iter_values( + (0..rows).map(|_| rng.random_range(-10.0f32..10.0)), + )), + DataType::FixedSizeList(item, size) => { + let len = rows * *size as usize; + let values: ArrayRef = match item.data_type() { + DataType::UInt8 => Arc::new(UInt8Array::from_iter_values( + (0..len).map(|_| rng.random()), + )), + DataType::Float32 => Arc::new(Float32Array::from_iter_values( + (0..len).map(|_| rng.random_range(-10.0f32..10.0)), + )), + other => panic!("unexpected item type {other}"), + }; + Arc::new( + FixedSizeListArray::try_new(item.clone(), *size, values, None).unwrap(), + ) + } + other => panic!("unexpected type {other}"), + } + }) + .collect(); + RecordBatch::try_new(schema.clone(), columns).unwrap() + } + + /// Every stride as the format documentation gives it. + fn documented_strides(metadata: &RabitQuantizationMetadata) -> Vec<(&'static str, usize)> { + let dim = metadata.code_dim as usize; + let sign = dim.div_ceil(8); + let padded = 64 * dim.div_ceil(64); + if metadata.layered { + let layout = crate::vector::bq::layered::RQLayout::try_new(metadata.num_bits).unwrap(); + vec![ + (RQ_SIGN_ROWS_COLUMN, 20 + sign), + (RQ_BOUNDS_ROWS_COLUMN, 24), + ( + RQ_HIGH_ROWS_COLUMN, + padded * layout.high_bits as usize / 8 + 8, + ), + ( + RQ_LOW_ROWS_COLUMN, + padded * layout.low_bits as usize / 8 + 8, + ), + ] + } else { + let factors = match metadata.query_estimator { + RabitQueryEstimator::RawQuery => 3, + RabitQueryEstimator::ResidualQuery => 2, + }; + let mut stride = 8 + sign + 4 * factors; + if metadata.num_bits > 1 { + stride += padded * (metadata.num_bits as usize - 1) / 8 + 8; + } + vec![(RQ_ROWS_COLUMN, stride)] + } + } + + fn all_layouts() -> Vec { + let mut layouts = Vec::new(); + for dim in [64, 128, 1024] { + for num_bits in [1, 3, 4, 5, 6, 7, 8, 9] { + for estimator in [ + RabitQueryEstimator::RawQuery, + RabitQueryEstimator::ResidualQuery, + ] { + layouts.push(metadata(num_bits, dim, false, estimator)); + } + if num_bits > 1 { + layouts.push(metadata(num_bits, dim, true, RabitQueryEstimator::RawQuery)); + } + } + } + layouts + } + + /// Packing then unpacking every native and layered layout gives the + /// column layout's batch back, with buffers of exactly their size, in one + /// batch or across reads of several, with the documented strides. + #[test] + fn plane_rows_round_trip_every_layout() { + for metadata in all_layouts() { + let context = format!( + "num_bits={} dim={} layered={} {:?}", + metadata.num_bits, metadata.code_dim, metadata.layered, metadata.query_estimator + ); + let spec = PlaneRowsSpec::try_new(&metadata).unwrap(); + let strides: Vec<_> = spec + .packed_columns() + .iter() + .map(|column| (column.name(), column.stride())) + .collect(); + assert_eq!(strides, documented_strides(&metadata), "{context}"); + for rows in [0, 1, 37] { + let batch = random_batch(spec.logical_schema(), rows, rows as u64); + let packed = spec.pack(&batch).unwrap(); + assert_eq!(packed.num_rows(), rows, "{context}"); + let unpacked = spec + .unpack(std::slice::from_ref(&packed), spec.logical_schema()) + .unwrap(); + assert_eq!(unpacked, batch, "{context}"); + // The arrays plus exactly the values' bytes: a packed row's + // per row. + let row_bytes: usize = strides.iter().map(|(_, stride)| stride).sum(); + let arrays = spec + .unpack(&[], spec.logical_schema()) + .unwrap() + .deep_size_of_children(&mut Context::new()); + assert_eq!( + unpacked.deep_size_of_children(&mut Context::new()), + arrays + rows * row_bytes, + "{context}" + ); + if rows > 2 { + let split = [packed.slice(0, 2), packed.slice(2, rows - 2)]; + assert_eq!( + spec.unpack(&split, spec.logical_schema()).unwrap(), + batch, + "{context}" + ); + } + } + } + } + + /// The documented example: d = 1024, 7 bits. + #[test] + fn plane_rows_strides_match_the_format_example() { + let layered = + PlaneRowsSpec::try_new(&metadata(7, 1024, true, RabitQueryEstimator::RawQuery)) + .unwrap(); + let strides: Vec<_> = layered + .packed_columns() + .iter() + .map(|column| column.stride()) + .collect(); + assert_eq!(strides, [148, 24, 520, 264]); + let native = + PlaneRowsSpec::try_new(&metadata(7, 1024, false, RabitQueryEstimator::RawQuery)) + .unwrap(); + assert_eq!(native.packed_columns()[0].stride(), 924); + } + + /// A packed row is its fields' little-endian values concatenated in + /// order, without padding. + #[test] + fn plane_rows_bytes_equal_hand_built_rows() { + let metadata = metadata(3, 64, false, RabitQueryEstimator::RawQuery); + let spec = PlaneRowsSpec::try_new(&metadata).unwrap(); + let batch = random_batch(spec.logical_schema(), 2, 7); + let packed = spec.pack(&batch).unwrap(); + let rows = packed[RQ_ROWS_COLUMN].as_fixed_size_list(); + let rows = rows.values().as_primitive::().values(); + let stride = spec.packed_columns()[0].stride(); + for row in 0..2 { + let mut expected = Vec::new(); + expected.extend_from_slice( + &batch[ROW_ID] + .as_primitive::() + .value(row) + .to_le_bytes(), + ); + let list_bytes = |name: &str| { + batch[name] + .as_fixed_size_list() + .value(row) + .as_primitive::() + .values() + .to_vec() + }; + let factor = |name: &str| { + batch[name] + .as_primitive::() + .value(row) + .to_le_bytes() + }; + expected.extend(list_bytes(RABIT_CODE_COLUMN)); + for name in [ + ADD_FACTORS_COLUMN, + SCALE_FACTORS_COLUMN, + ERROR_FACTORS_COLUMN, + ] { + expected.extend(factor(name)); + } + expected.extend(list_bytes(RABIT_BLOCKED_EX_CODE_COLUMN)); + for name in [EX_ADD_FACTORS_COLUMN, EX_SCALE_FACTORS_COLUMN] { + expected.extend(factor(name)); + } + assert_eq!(expected.len(), stride); + assert_eq!(&rows[row * stride..(row + 1) * stride], expected.as_slice()); + } + // The bounds pack as three little-endian floats each. + let layered = PlaneRowsSpec::try_new(&metadata_layered()).unwrap(); + let batch = random_batch(layered.logical_schema(), 1, 3); + let packed = layered.pack(&batch).unwrap(); + let bounds = packed[RQ_BOUNDS_ROWS_COLUMN].as_fixed_size_list(); + let bounds = bounds.values().as_primitive::().values(); + let expected: Vec = [HIGH_BOUNDS_COLUMN, FULL_BOUNDS_COLUMN] + .iter() + .flat_map(|name| { + batch[*name] + .as_fixed_size_list() + .values() + .as_primitive::() + .values() + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>() + }) + .collect(); + assert_eq!(bounds, expected.as_slice()); + } + + fn metadata_layered() -> RabitQuantizationMetadata { + metadata(7, 128, true, RabitQueryEstimator::RawQuery) + } + + /// Sign codes keep their partition-local transposition: the bytes of + /// the packed buffer survive, and selecting rows from them still agrees. + #[test] + fn plane_rows_keep_transposed_sign_codes() { + let spec = PlaneRowsSpec::try_new(&metadata_layered()).unwrap(); + let mut batch = random_batch(spec.logical_schema(), 45, 11); + let transposed = pack_codes(batch[RABIT_CODE_COLUMN].as_fixed_size_list()); + batch = batch + .replace_column_by_name(RABIT_CODE_COLUMN, Arc::new(transposed.clone())) + .unwrap(); + let unpacked = spec + .unpack(&[spec.pack(&batch).unwrap()], spec.logical_schema()) + .unwrap(); + let codes = unpacked[RABIT_CODE_COLUMN].as_fixed_size_list(); + assert_eq!(codes, &transposed); + let rows = [0, 3, 31, 32, 44]; + assert_eq!( + take_packed_codes(codes, &rows).unwrap(), + take_packed_codes(&transposed, &rows).unwrap() + ); + } + + #[test] + fn plane_rows_reject_nulls_bad_widths_and_types() { + let spec = PlaneRowsSpec::try_new(&metadata_layered()).unwrap(); + let batch = random_batch(spec.logical_schema(), 4, 5); + + let mut nulls = Float32Builder::new(); + nulls.append_value(1.0); + nulls.append_null(); + nulls.append_value(1.0); + nulls.append_value(1.0); + let with_nulls = RecordBatch::try_new( + Arc::new(Schema::new( + batch + .schema() + .fields() + .iter() + .map(|field| field.as_ref().clone().with_nullable(true)) + .collect::>(), + )), + batch.columns().to_vec(), + ) + .unwrap() + .replace_column_by_name(ADD_FACTORS_COLUMN, Arc::new(nulls.finish())) + .unwrap(); + let err = spec.pack(&with_nulls).unwrap_err().to_string(); + assert!(err.contains("nulls"), "{err}"); + + let code_bytes = metadata_layered().binary_code_bytes(); + let wide_codes = FixedSizeListArray::try_new( + Arc::new(Field::new("item", DataType::UInt8, true)), + code_bytes as i32 + 1, + Arc::new(UInt8Array::from(vec![0u8; 4 * (code_bytes + 1)])), + None, + ) + .unwrap(); + let mut fields: Vec = batch + .schema() + .fields() + .iter() + .map(|field| field.as_ref().clone()) + .collect(); + let mut columns = batch.columns().to_vec(); + let code_index = batch.schema().index_of(RABIT_CODE_COLUMN).unwrap(); + fields[code_index] = Field::new(RABIT_CODE_COLUMN, wide_codes.data_type().clone(), true); + columns[code_index] = Arc::new(wide_codes); + let wide = + RecordBatch::try_new(Arc::new(Schema::new(fields.clone())), columns.clone()).unwrap(); + assert!(spec.pack(&wide).is_err()); + + let row_index = batch.schema().index_of(ROW_ID).unwrap(); + fields[code_index] = batch.schema().field(code_index).clone(); + columns[code_index] = batch.column(code_index).clone(); + fields[row_index] = Field::new(ROW_ID, DataType::Int64, true); + columns[row_index] = Arc::new(arrow_array::Int64Array::from(vec![1i64; 4])); + let wrong_type = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); + assert!(spec.pack(&wrong_type).is_err()); + + let extra = batch + .clone() + .try_with_column( + Field::new("other", DataType::UInt64, true), + Arc::new(UInt64Array::from(vec![1u64; 4])), + ) + .unwrap(); + assert!(spec.pack(&extra).is_err()); + assert!( + spec.pack(&batch.drop_column(HIGH_BOUNDS_COLUMN).unwrap()) + .is_err() + ); + + // A field type without a plane-row form. + let field = Field::new("utf8", DataType::Utf8, true); + assert!(ValueKind::of(&field).is_err()); + } + + /// A read maps whole planes to their packed columns, and no read can + /// take only part of a plane. + #[test] + fn plane_rows_packed_for_planes_and_rejects_partial_covers() { + let spec = PlaneRowsSpec::try_new(&metadata_layered()).unwrap(); + let cases = [ + (0, SignBounds::Lazy, vec![RQ_SIGN_ROWS_COLUMN]), + ( + 0, + SignBounds::Eager, + vec![RQ_SIGN_ROWS_COLUMN, RQ_BOUNDS_ROWS_COLUMN], + ), + (1, SignBounds::Lazy, vec![RQ_HIGH_ROWS_COLUMN]), + (2, SignBounds::Eager, vec![RQ_LOW_ROWS_COLUMN]), + (3, SignBounds::Lazy, vec![RQ_BOUNDS_ROWS_COLUMN]), + ]; + for (plane, bounds, expected) in cases { + let columns = plane_columns(plane, bounds); + assert_eq!(spec.packed_for(columns).unwrap(), expected); + let projection = spec.projection(columns).unwrap(); + let names: Vec<_> = projection + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect(); + assert_eq!(names, columns); + } + let all: Vec<&str> = spec + .logical_schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect(); + assert_eq!(spec.packed_for(&all).unwrap().len(), 4); + + let err = spec + .packed_for(&[ROW_ID, RABIT_CODE_COLUMN]) + .unwrap_err() + .to_string(); + assert!(err.contains(RQ_SIGN_ROWS_COLUMN), "{err}"); + assert!(spec.packed_for(&[RABIT_BLOCKED_EX_CODE_COLUMN]).is_err()); + assert!(spec.packed_for(&["missing"]).is_err()); + assert!(spec.projection(&["missing"]).is_err()); + } + + #[test] + fn plane_rows_file_fields_request_fullzip() { + let spec = PlaneRowsSpec::try_new(&metadata_layered()).unwrap(); + for field in spec.file_fields(true) { + assert!(!field.is_nullable()); + let DataType::FixedSizeList(item, _) = field.data_type() else { + panic!("{field:?}"); + }; + assert_eq!(item.data_type(), &DataType::UInt8); + assert!(!item.is_nullable()); + assert_eq!( + field.metadata().get(STRUCTURAL_ENCODING_META_KEY), + Some(&STRUCTURAL_ENCODING_FULLZIP.to_string()) + ); + assert_eq!( + field.metadata().get(COMPRESSION_META_KEY), + Some(&NO_COMPRESSION.to_string()) + ); + } + assert!( + spec.file_fields(false) + .iter() + .all(|field| field.metadata().is_empty()) + ); + assert!(!requests_fullzip(ConcreteFileVersion::V2_0)); + assert!(requests_fullzip(ConcreteFileVersion::V2_1)); + assert!(requests_fullzip(ConcreteFileVersion::V2_2)); + } + + /// A plane-row file must have exactly the packed columns its metadata + /// implies, then carried columns; anything else is an invalid index. + #[rstest] + #[case::native(false)] + #[case::layered(true)] + fn plane_rows_validate_file_schemas(#[case] layered: bool) { + let metadata = metadata(7, 128, layered, RabitQueryEstimator::RawQuery); + let spec = PlaneRowsSpec::try_new(&metadata).unwrap(); + let carried = Field::new("carried", DataType::Utf8, true); + let mut fields = spec.file_fields(true); + fields.push(carried.clone()); + let file = Schema::new_with_metadata( + fields.clone(), + HashMap::from([("key".to_string(), "value".to_string())]), + ); + let read = PlaneRowsSpec::for_file(&metadata, &file).unwrap(); + // A Lance file reads list items back as nullable. + let read_back: Vec = fields + .iter() + .map(|field| match field.data_type() { + DataType::FixedSizeList(item, size) => Field::new( + field.name(), + DataType::FixedSizeList( + Arc::new(item.as_ref().clone().with_nullable(true)), + *size, + ), + field.is_nullable(), + ), + _ => field.clone(), + }) + .collect(); + PlaneRowsSpec::for_file(&metadata, &Schema::new(read_back)).unwrap(); + assert_eq!(read.carried_fields().len(), 1); + assert_eq!(read.logical_schema().metadata(), file.metadata()); + assert_eq!( + read.logical_schema().fields().last().unwrap().as_ref(), + &carried + ); + assert_eq!( + read.packed_for(&["carried"]).unwrap(), + vec!["carried"], + "a carried column reads alone" + ); + + let mut wrong_stride = fields.clone(); + wrong_stride[0] = Field::new( + wrong_stride[0].name(), + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::UInt8, false)), 3), + false, + ); + let mut nullable = fields.clone(); + nullable[0] = nullable[0].clone().with_nullable(true); + let mut column_layout: Vec = spec + .logical_schema() + .fields() + .iter() + .map(|field| field.as_ref().clone()) + .collect(); + column_layout.push(carried); + let mut legacy = fields.clone(); + legacy.push(Field::new(RABIT_CODE_COLUMN, DataType::UInt8, true)); + for (case, fields) in [ + ("wrong stride", wrong_stride), + ("nullable", nullable), + ("column layout", column_layout), + ("an internal column carried", legacy), + ("no columns", Vec::new()), + ] { + let err = PlaneRowsSpec::for_file(&metadata, &Schema::new(fields)).unwrap_err(); + assert!(matches!(err, Error::Index { .. }), "{case}: {err:?}"); + } + } + + /// Converting needs the column layout's fields exactly as the metadata + /// implies them; other columns are carried, and the legacy sequential ex + /// codes have no plane-row form. + #[test] + fn plane_rows_for_column_files() { + let metadata = metadata(7, 128, false, RabitQueryEstimator::RawQuery); + let spec = PlaneRowsSpec::try_new(&metadata).unwrap(); + let mut fields: Vec = spec + .logical_schema() + .fields() + .iter() + .map(|field| field.as_ref().clone()) + .collect(); + let carried = Field::new("carried", DataType::Float32, true); + fields.push(carried.clone()); + let read = PlaneRowsSpec::for_column_file(&metadata, &Schema::new(fields.clone())).unwrap(); + assert_eq!(read.carried_fields()[0].as_ref(), &carried); + let batch = random_batch(read.logical_schema(), 5, 9); + let packed = read.pack(&batch).unwrap(); + assert_eq!(packed.schema().fields().len(), 2); + assert_eq!( + read.unpack(&[packed], read.logical_schema()).unwrap(), + batch + ); + + let mut legacy = fields.clone(); + let ex = legacy + .iter() + .position(|field| field.name() == RABIT_BLOCKED_EX_CODE_COLUMN) + .unwrap(); + legacy[ex] = Field::new(RABIT_EX_CODE_COLUMN, legacy[ex].data_type().clone(), true); + assert!(PlaneRowsSpec::for_column_file(&metadata, &Schema::new(legacy)).is_err()); + + let mut retyped = fields.clone(); + retyped[0] = Field::new(ROW_ID, DataType::UInt64, false); + assert!(PlaneRowsSpec::for_column_file(&metadata, &Schema::new(retyped)).is_err()); + + let mut reordered = fields; + reordered.swap(2, 3); + assert!(PlaneRowsSpec::for_column_file(&metadata, &Schema::new(reordered)).is_err()); + } + + /// The unpack counters advance by the packed bytes each read unpacks. + #[test] + fn plane_rows_unpack_counts_packed_bytes() { + let spec = PlaneRowsSpec::try_new(&metadata_layered()).unwrap(); + let batch = random_batch(spec.logical_schema(), 10, 1); + let packed = spec.pack(&batch).unwrap(); + let before = layered_stats::counters().plane_rows_unpack_bytes.get(); + let columns = plane_columns(1, SignBounds::Lazy); + spec.unpack(&[packed], &spec.projection(columns).unwrap()) + .unwrap(); + let high = spec.packed_columns()[2].stride() * 10; + assert!(layered_stats::counters().plane_rows_unpack_bytes.get() >= before + high as u64); + } +} diff --git a/rust/lance-index/src/vector/bq/resident.rs b/rust/lance-index/src/vector/bq/resident.rs index d66fef0cb72..928c99078ba 100644 --- a/rust/lance-index/src/vector/bq/resident.rs +++ b/rust/lance-index/src/vector/bq/resident.rs @@ -26,10 +26,8 @@ use std::collections::HashMap; use std::sync::{Arc, LazyLock, Mutex, Weak}; use std::time::Instant; -use arrow::compute::concat_batches; use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, UInt64Array, new_empty_array}; use arrow_schema::{Field, Schema, SchemaRef}; -use arrow_select::concat::concat; use arrow_select::take::take; use futures::TryStreamExt; use lance_core::cache::{ @@ -49,6 +47,7 @@ use super::storage::{ RABIT_BLOCKED_EX_CODE_COLUMN, RABIT_BLOCKED_EX_CODE_LO_COLUMN, RABIT_CODE_COLUMN, RABIT_EX_CODE_COLUMN, }; +use crate::vector::exact_buffers::{exact_array, exact_batch}; use crate::vector::storage::{IndexFileKey, ResidentLifetime, WeakRegistry, shared_by_key}; /// Columns that reads always fetch from the file: the codes, which hold most @@ -459,7 +458,9 @@ impl ResidentColumnStore { .await? .try_collect::>() .await?; - Some(concat_batches(&projected, batches.iter())?) + // Buffers of exactly the values, so the cache charges the store + // what `resident_columns_bytes` sizes it at without reading it. + Some(exact_batch(&projected, &batches, false)?) }; let mut columns = HashMap::with_capacity(fields.len()); let mut bytes = 0u64; @@ -659,11 +660,11 @@ pub(crate) struct ResidentColumn { } impl ResidentColumn { - /// A copy of the ascending file rows `rows`, laid out as a file read of - /// them so that a cache entry holding it is charged the same bytes: the - /// reader decodes each page a read touches into a buffer of exactly its - /// rows and concatenates the buffers when the read touches several - /// pages. A slice would instead keep, and be charged, the whole column. + /// A copy of the ascending file rows `rows` in a buffer of exactly their + /// values, as a whole read of the file returns them (see + /// `read_projected`), so that a cache entry holding it is charged the + /// same bytes. A slice would instead keep, and be charged, the whole + /// column. pub(crate) fn copy_rows(&self, rows: &UInt64Array) -> Result { let offsets = rows.values(); // The page split below relies on the order. @@ -695,10 +696,7 @@ impl ResidentColumn { match pieces.len() { 0 => Ok(new_empty_array(self.values.data_type())), 1 => Ok(pieces.swap_remove(0)), - _ => { - let pieces: Vec<&dyn Array> = pieces.iter().map(|piece| piece.as_ref()).collect(); - Ok(concat(&pieces)?) - } + _ => exact_array(&pieces.iter().collect::>(), false), } } } @@ -1090,8 +1088,11 @@ mod tests { } } + /// Rows across pages are copied into one buffer of exactly their + /// values, as a whole read of the file returns them wherever its pages + /// end. #[test] - fn copy_rows_across_pages_matches_a_file_read() { + fn copy_rows_across_pages_holds_exactly_the_rows() { let column = column(); for rows in [ (30..60).collect::>(), @@ -1101,18 +1102,13 @@ mod tests { let copied = copy(&column, &rows); assert_eq!(values(&copied), as_values(&rows)); assert!(copied.nulls().is_none()); - // The reader concatenates the pages it decodes, each into a - // buffer of exactly its rows. let (first, second): (Vec, Vec) = rows.iter().partition(|&&row| row < 40); let pages = [first, second] .map(|page| take(column.values.as_ref(), &UInt64Array::from(page), None).unwrap()); - let read = concat(&[pages[0].as_ref(), pages[1].as_ref()]).unwrap(); + let read = + arrow_select::concat::concat(&[pages[0].as_ref(), pages[1].as_ref()]).unwrap(); assert_eq!(copied.as_ref(), read.as_ref(), "{rows:?}"); - assert_eq!( - copied.get_buffer_memory_size(), - read.get_buffer_memory_size(), - "{rows:?}" - ); + assert_eq!(copied.get_buffer_memory_size(), rows.len() * 4, "{rows:?}"); } } } diff --git a/rust/lance-index/src/vector/bq/storage.rs b/rust/lance-index/src/vector/bq/storage.rs index 982388e932f..4c683668fb4 100644 --- a/rust/lance-index/src/vector/bq/storage.rs +++ b/rust/lance-index/src/vector/bq/storage.rs @@ -314,6 +314,55 @@ pub fn rabit_ex_code_field(rotated_dim: usize, num_bits: u8) -> Result &'static str { + match self { + Self::Columns => "columns", + Self::PlaneRows => "plane_rows", + } + } + + /// Whether this is the column layout, which the storage metadata leaves + /// out so that column-layout metadata stays as older writers wrote it. + pub fn is_columns(&self) -> bool { + *self == Self::Columns + } +} + +impl std::fmt::Display for RQRowLayout { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +impl std::str::FromStr for RQRowLayout { + type Err = Error; + + fn from_str(value: &str) -> Result { + match value { + "columns" => Ok(Self::Columns), + "plane_rows" => Ok(Self::PlaneRows), + _ => Err(Error::invalid_input(format!( + "unknown IVF_RQ row layout {value:?}; expected columns or plane_rows" + ))), + } + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct RabitQuantizationMetadata { // this rotate matrix is large, and lance index would store all metadata in schema metadata, @@ -334,6 +383,9 @@ pub struct RabitQuantizationMetadata { /// Whether the ex codes use fixed high and low prefix planes. #[serde(default)] pub layered: bool, + /// How the auxiliary file stores each row's fields. + #[serde(default, skip_serializing_if = "RQRowLayout::is_columns")] + pub row_layout: RQRowLayout, #[serde(default = "default_query_estimator_compat")] pub query_estimator: RabitQueryEstimator, } @@ -501,6 +553,10 @@ impl QuantizerMetadata for RabitQuantizationMetadata { self.rotate_mat_position = Some(index); } + fn row_layout(&self) -> RQRowLayout { + self.row_layout + } + fn parse_buffer(&mut self, bytes: Bytes) -> Result<()> { if self.rotation_type != RQRotationType::Matrix { return Ok(()); @@ -4269,6 +4325,39 @@ mod tests { assert_eq!(metadata.query_estimator, RabitQueryEstimator::ResidualQuery); } + /// Metadata without `row_layout` is the column layout, which is never + /// written, so column-layout metadata keeps the bytes older writers + /// produced; the plane-row layout is written and read back by name. + #[test] + fn test_rabit_metadata_row_layout_serde() { + let legacy = r#"{"rotate_mat_position":0,"rotation_type":"matrix","code_dim":64,"num_bits":1,"packed":true,"layered":false,"query_estimator":"raw_query"}"#; + let metadata: RabitQuantizationMetadata = serde_json::from_str(legacy).unwrap(); + assert_eq!(metadata.row_layout, RQRowLayout::Columns); + assert_eq!(metadata.row_layout(), RQRowLayout::Columns); + let json = serde_json::to_string(&metadata).unwrap(); + assert!(!json.contains("row_layout"), "{json}"); + + let mut plane_rows = metadata; + plane_rows.row_layout = RQRowLayout::PlaneRows; + let json = serde_json::to_string(&plane_rows).unwrap(); + assert!(json.contains(r#""row_layout":"plane_rows""#), "{json}"); + let read: RabitQuantizationMetadata = serde_json::from_str(&json).unwrap(); + assert_eq!(read.row_layout, RQRowLayout::PlaneRows); + + for layout in [RQRowLayout::Columns, RQRowLayout::PlaneRows] { + assert_eq!(layout.as_str().parse::().unwrap(), layout); + assert_eq!(layout.to_string(), layout.as_str()); + } + assert!("rows".parse::().is_err()); + assert!( + serde_json::from_str::(&legacy.replace( + r#""layered":false"#, + r#""layered":false,"row_layout":"rows""# + )) + .is_err() + ); + } + #[test] fn test_new_rabit_metadata_uses_raw_query_estimator() { let metadata = make_test_metadata(64); @@ -4439,6 +4528,7 @@ mod tests { num_bits: 2, packed: false, layered: false, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, }; let codes = @@ -4566,6 +4656,7 @@ mod tests { num_bits, packed: false, layered: false, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, }; let codes = FixedSizeListArray::try_new_from_values( @@ -4731,6 +4822,7 @@ mod tests { num_bits: 2, packed: false, layered: false, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, }; let codes = @@ -4905,6 +4997,7 @@ mod tests { num_bits, packed: false, layered: false, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, }; let code_len = rabit_binary_code_bytes(code_dim); @@ -5110,6 +5203,7 @@ mod tests { num_bits, packed: false, layered: false, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, }; let codes = FixedSizeListArray::try_new_from_values( @@ -5195,6 +5289,7 @@ mod tests { num_bits, packed: false, layered: false, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, }; let codes = FixedSizeListArray::try_new_from_values( @@ -5515,6 +5610,7 @@ mod tests { num_bits: 1, packed: false, layered: false, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, }; let codes = @@ -5550,6 +5646,7 @@ mod tests { num_bits: 2, packed: false, layered: false, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, }; let codes = diff --git a/rust/lance-index/src/vector/distributed/index_merger.rs b/rust/lance-index/src/vector/distributed/index_merger.rs index 8227a1e9d36..90c238505e3 100755 --- a/rust/lance-index/src/vector/distributed/index_merger.rs +++ b/rust/lance-index/src/vector/distributed/index_merger.rs @@ -21,8 +21,8 @@ use crate::IndexMetadata as IndexMetaSchema; use crate::pb; use crate::scalar::OldIndexDataFilter; use crate::vector::bq::storage::{ - RABIT_CODE_COLUMN, RABIT_METADATA_KEY, RabitQuantizationMetadata, RabitQueryEstimator, - pack_codes, rabit_binary_code_field, rabit_ex_code_field, + RABIT_CODE_COLUMN, RABIT_METADATA_KEY, RQRowLayout, RabitQuantizationMetadata, + RabitQueryEstimator, pack_codes, rabit_binary_code_field, rabit_ex_code_field, }; use crate::vector::bq::transform::{ ADD_FACTORS_FIELD, ERROR_FACTORS_FIELD, EX_ADD_FACTORS_FIELD, EX_SCALE_FACTORS_FIELD, @@ -361,6 +361,17 @@ pub async fn init_writer_for_sq( Ok(w) } +/// The merger reads and writes the column layout only: a plane-row shard, or +/// plane-row output, is not supported until it handles packed row columns. +fn reject_plane_rows(rq_meta: &RabitQuantizationMetadata, what: &str) -> Result<()> { + if rq_meta.row_layout == RQRowLayout::PlaneRows { + return Err(Error::not_supported(format!( + "Distributed RQ merge: {what} uses the plane-row layout, which the merger does not support; merge column-layout segments and convert the result" + ))); + } + Ok(()) +} + /// Create and initialize a unified writer for RQ storage. pub async fn init_writer_for_rq( object_store: &lance_io::object_store::ObjectStore, @@ -369,6 +380,7 @@ pub async fn init_writer_for_rq( rq_meta: &RabitQuantizationMetadata, format_version: ConcreteFileVersion, ) -> Result { + reject_plane_rows(rq_meta, "the merged auxiliary file")?; let mut fields = vec![ (*ROW_ID_FIELD).clone(), rabit_binary_code_field(rq_meta.rotated_dim()), @@ -1148,6 +1160,7 @@ async fn merge_partial_vector_auxiliary_files_inner( }; let mut rq_meta_parsed: RabitQuantizationMetadata = serde_json::from_str(&rq_json) .map_err(|e| Error::index(format!("RQ metadata parse error: {}", e)))?; + reject_plane_rows(&rq_meta_parsed, &format!("source shard {idx}"))?; if rq_meta_parsed.rotation_type == crate::vector::bq::RQRotationType::Matrix && rq_meta_parsed.rotate_mat.is_none() && let Some(buf_idx) = rq_meta_parsed.buffer_index() @@ -2568,6 +2581,13 @@ mod tests { if metadata.layered { columns.push(Arc::new(Float32Array::from(vec![100.; total_rows]))); columns.push(Arc::new(Float32Array::from(vec![10.; total_rows]))); + // The high and full estimator bounds, three floats each. + for bound in [1.0f32, 2.0] { + columns.push(Arc::new(FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![bound; total_rows * 3]), + 3, + )?)); + } } } let batch = RecordBatch::try_new(Arc::new(arrow_schema), columns)?; @@ -2791,6 +2811,7 @@ mod tests { num_bits: 1, packed: false, layered: false, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, }; @@ -2924,6 +2945,7 @@ mod tests { num_bits: 1, packed: true, layered: false, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, }; @@ -2965,6 +2987,65 @@ mod tests { } } + /// The merger handles the column layout only: a shard whose metadata + /// declares plane rows is not supported, whether its codes are packed + /// or not, and it never writes plane-row metadata over columns. + #[rstest::rstest] + #[tokio::test] + async fn test_merge_ivf_rq_rejects_plane_row_shard(#[values(false, true)] packed: bool) { + let object_store = ObjectStore::memory(); + let index_dir = Path::from("index/uuid_rq_plane_rows"); + let aux0 = index_dir + .clone() + .join("partial_0") + .join(INDEX_AUXILIARY_FILE_NAME); + let rq_meta = RabitQuantizationMetadata { + rotate_mat: None, + rotate_mat_position: None, + fast_rotation_signs: Some(vec![0xAA; 2]), + rotation_type: RQRotationType::Fast, + code_dim: 16, + num_bits: 1, + packed, + layered: false, + row_layout: RQRowLayout::PlaneRows, + query_estimator: RabitQueryEstimator::RawQuery, + }; + let writer = init_writer_for_rq( + &object_store, + &aux0, + DistanceType::L2, + &rq_meta, + ConcreteFileVersion::V2_0, + ) + .await; + match writer { + Err(Error::NotSupported { .. }) => {} + Err(other) => panic!("expected Error::NotSupported, got {other:?}"), + Ok(_) => panic!("the merger wrote plane-row metadata over columns"), + } + // The shard's rows are written in columns; only its metadata matters. + write_rq_partial_aux(&object_store, &aux0, &rq_meta, &[2, 1], 0, DistanceType::L2) + .await + .unwrap(); + + let res = merge_partial_vector_auxiliary_files( + &object_store, + std::slice::from_ref(&aux0), + &index_dir, + crate::progress::noop_progress(), + ) + .await; + match res { + Err(Error::NotSupported { source, .. }) => { + let message = source.to_string(); + assert!(message.contains("source shard 0"), "{message}"); + assert!(message.contains("plane-row layout"), "{message}"); + } + other => panic!("expected Error::NotSupported for a plane-row shard, got {other:?}"), + } + } + #[rstest::rstest] #[case::single(4, false)] #[case::rq5(5, true)] @@ -2995,6 +3076,7 @@ mod tests { num_bits, packed: false, layered, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, }; @@ -3138,6 +3220,7 @@ mod tests { num_bits: 5, packed: false, layered: false, + row_layout: RQRowLayout::Columns, query_estimator: RabitQueryEstimator::RawQuery, }; write_rq_partial_aux(&store, &a, &metadata, &[2], 0, DistanceType::L2) diff --git a/rust/lance-index/src/vector/exact_buffers.rs b/rust/lance-index/src/vector/exact_buffers.rs new file mode 100644 index 00000000000..0030df4d1fd --- /dev/null +++ b/rust/lance-index/src/vector/exact_buffers.rs @@ -0,0 +1,286 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Index cache entries in buffers of exactly their values' size. +//! +//! The index cache charges an entry what its arrays allocate. A plane-row +//! read of an IVF_RQ file unpacks its values into buffers of exactly their +//! size, while a read of the column layout can keep more: a decoder that +//! joins the pages a read spans rounds its buffers up, and a slice keeps the +//! whole buffer it was cut from. Reads that become cache entries are copied +//! here into exact buffers where they are not already, so the entries of +//! both layouts, and of native and layered indexes alike, weigh their +//! values. + +use std::sync::Arc; + +use arrow::array::ArrayData; +use arrow::buffer::{Buffer, MutableBuffer}; +use arrow::compute::concat_batches; +use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, make_array}; +use arrow_schema::{DataType, SchemaRef}; +use arrow_select::concat::concat; +use lance_arrow::deepcopy::deep_copy_array_sliced; +use lance_core::Result; + +/// Whether `data` holds exactly its values: a fixed-width column, or a +/// fixed-size list of one, without nulls, whose buffer allocates its values' +/// bytes and no more. Other layouts are never exact here. +fn is_exact(data: &ArrayData) -> bool { + if data.nulls().is_some() { + return false; + } + match data.data_type() { + DataType::FixedSizeList(_, size) => match data.child_data() { + [child] => { + data.offset() == 0 && child.len() == data.len() * *size as usize && is_exact(child) + } + _ => false, + }, + data_type => match (data_type.primitive_width(), data.buffers()) { + (Some(width), [values]) => values.capacity() == data.len() * width, + _ => false, + }, + } +} + +/// `pieces`, consecutive parts of one column, as one array in buffers of +/// exactly its values' size; `None` for a layout [`is_exact`] does not size +/// or for pieces with nulls. +fn exact_concat(pieces: &[ArrayData]) -> Option { + let data_type = pieces.first()?.data_type(); + if pieces.iter().any(|piece| piece.nulls().is_some()) { + return None; + } + let len = pieces.iter().map(ArrayData::len).sum(); + match data_type { + DataType::FixedSizeList(_, size) => { + let size = *size as usize; + let children = pieces + .iter() + .map(|piece| match piece.child_data() { + [child] => Some(child.slice(piece.offset() * size, piece.len() * size)), + _ => None, + }) + .collect::>>()?; + let child = exact_concat(&children)?; + ArrayData::builder(data_type.clone()) + .len(len) + .child_data(vec![child]) + .build() + .ok() + } + _ => { + let width = data_type.primitive_width()?; + let mut values = MutableBuffer::from_len_zeroed(len * width); + let mut written = 0; + for piece in pieces { + let [buffer] = piece.buffers() else { + return None; + }; + let start = piece.offset() * width; + let bytes = buffer.as_slice().get(start..start + piece.len() * width)?; + values.as_slice_mut()[written..written + bytes.len()].copy_from_slice(bytes); + written += bytes.len(); + } + ArrayData::builder(data_type.clone()) + .len(len) + .add_buffer(Buffer::from(values)) + .build() + .ok() + } + } +} + +/// `pieces`, consecutive parts of one column, as one array whose +/// fixed-width values sit in a buffer of exactly their size. A single piece +/// that already does is kept without a copy, unless `copy` asks for one so +/// that the array never shares the buffer of a larger read. A layout +/// [`exact_concat`] does not size is concatenated, or copied, as is. +pub fn exact_array(pieces: &[&ArrayRef], copy: bool) -> Result { + let data: Vec = pieces.iter().map(|piece| piece.to_data()).collect(); + if let [single] = data.as_slice() + && !copy + && is_exact(single) + { + return Ok(pieces[0].clone()); + } + if let Some(exact) = exact_concat(&data) { + return Ok(make_array(exact)); + } + Ok(match pieces { + [single] if copy => deep_copy_array_sliced(single.as_ref()), + [single] => (*single).clone(), + pieces => { + let pieces: Vec<&dyn Array> = pieces.iter().map(|piece| piece.as_ref()).collect(); + concat(&pieces)? + } + }) +} + +/// `batches`, consecutive reads of `schema`'s columns, as one batch whose +/// columns are [`exact_array`]s. +pub fn exact_batch(schema: &SchemaRef, batches: &[RecordBatch], copy: bool) -> Result { + if batches.is_empty() { + return Ok(concat_batches(schema, batches)?); + } + let rows = batches.iter().map(RecordBatch::num_rows).sum(); + let columns = (0..schema.fields().len()) + .map(|column| { + let pieces: Vec<&ArrayRef> = batches.iter().map(|batch| batch.column(column)).collect(); + exact_array(&pieces, copy) + }) + .collect::>>()?; + Ok(RecordBatch::try_new_with_options( + Arc::clone(schema), + columns, + &RecordBatchOptions::new().with_row_count(Some(rows)), + )?) +} + +/// Bytes `array`'s buffers allocate, as the index cache charges them. +#[cfg(test)] +pub fn allocated_bytes(array: &dyn Array) -> usize { + fn walk(data: &ArrayData) -> usize { + data.buffers().iter().map(Buffer::capacity).sum::() + + data + .nulls() + .map_or(0, |nulls| nulls.inner().inner().capacity()) + + data.child_data().iter().map(walk).sum::() + } + walk(&array.to_data()) +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow_array::types::{Float32Type, UInt8Type, UInt64Type}; + use arrow_array::{ + FixedSizeListArray, Float32Array, StringArray, UInt8Array, UInt64Array, cast::AsArray, + }; + use arrow_schema::{Field, Schema}; + use lance_arrow::FixedSizeListArrayExt; + + /// `rows` lists of `width` bytes, in a buffer of exactly their size. + fn codes(rows: usize, width: i32) -> ArrayRef { + let values: Vec = (0..rows * width as usize).map(|v| v as u8).collect(); + Arc::new(FixedSizeListArray::try_new_from_values(UInt8Array::from(values), width).unwrap()) + } + + /// Joined pieces and slices are copied into buffers of exactly their + /// values, which they equal; exact columns are kept, and copied only + /// when asked. + #[test] + fn exact_array_sizes_fixed_width_columns() { + let rows: ArrayRef = Arc::new(UInt64Array::from((0..100).collect::>())); + let factors: ArrayRef = Arc::new(Float32Array::from( + (0..100).map(|v| v as f32).collect::>(), + )); + let codes = codes(100, 3); + for array in [&rows, &factors, &codes] { + assert!(is_exact(&array.to_data()), "{array:?}"); + assert!(Arc::ptr_eq(&exact_array(&[array], false).unwrap(), array)); + let copied = exact_array(&[array], true).unwrap(); + assert_eq!(copied.as_ref(), array.as_ref()); + assert_eq!( + allocated_bytes(copied.as_ref()), + allocated_bytes(array.as_ref()) + ); + + // A slice keeps the whole buffer; its exact copy holds its rows. + let slice = array.slice(10, 30); + assert!(!is_exact(&slice.to_data())); + let exact = exact_array(&[&slice], false).unwrap(); + assert_eq!(exact.as_ref(), slice.as_ref()); + assert!(is_exact(&exact.to_data())); + assert_eq!( + allocated_bytes(exact.as_ref()), + allocated_bytes(array.as_ref()) * 30 / 100 + ); + + // Pieces join into one exact buffer, where a concatenation + // rounds its buffer up. + let pieces = [array.slice(0, 33), array.slice(33, 67)]; + let joined = exact_array(&[&pieces[0], &pieces[1]], false).unwrap(); + assert_eq!(joined.as_ref(), array.as_ref()); + assert_eq!( + allocated_bytes(joined.as_ref()), + allocated_bytes(array.as_ref()) + ); + } + assert_eq!( + exact_array(&[&codes.slice(5, 2)], false) + .unwrap() + .as_fixed_size_list() + .values() + .as_primitive::() + .values() + .as_ref(), + &[15, 16, 17, 18, 19, 20] + ); + assert_eq!( + exact_array(&[&rows.slice(98, 2)], false) + .unwrap() + .as_primitive::() + .values() + .as_ref(), + &[98, 99] + ); + assert_eq!( + exact_array(&[&factors.slice(0, 1)], true) + .unwrap() + .as_primitive::() + .value(0), + 0.0 + ); + } + + /// Columns with nulls or without a fixed width keep the decoder's + /// layout: kept, or concatenated, or copied as a slice. + #[test] + fn exact_array_keeps_other_layouts() { + let nullable: ArrayRef = Arc::new(Float32Array::from(vec![Some(1.0), None, Some(3.0)])); + let names: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c"])); + for array in [&nullable, &names] { + assert!(!is_exact(&array.to_data())); + assert_eq!( + exact_array(&[array], false).unwrap().as_ref(), + array.as_ref() + ); + assert_eq!( + exact_array(&[array], true).unwrap().as_ref(), + array.as_ref() + ); + let pieces = [array.slice(0, 1), array.slice(1, 2)]; + assert_eq!( + exact_array(&[&pieces[0], &pieces[1]], false) + .unwrap() + .as_ref(), + array.as_ref() + ); + } + } + + #[test] + fn exact_batch_joins_reads_of_a_schema() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", arrow_schema::DataType::UInt64, false), + Field::new("codes", codes(1, 4).data_type().clone(), true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(UInt64Array::from((0..10).collect::>())), + codes(10, 4), + ], + ) + .unwrap(); + let joined = exact_batch(&schema, &[batch.slice(0, 4), batch.slice(4, 6)], false).unwrap(); + assert_eq!(joined, batch); + for column in joined.columns() { + assert!(is_exact(&column.to_data())); + } + assert_eq!(exact_batch(&schema, &[], false).unwrap().num_rows(), 0); + } +} diff --git a/rust/lance-index/src/vector/quantizer.rs b/rust/lance-index/src/vector/quantizer.rs index b1f434453ed..a2d11b1e2e0 100644 --- a/rust/lance-index/src/vector/quantizer.rs +++ b/rust/lance-index/src/vector/quantizer.rs @@ -27,6 +27,7 @@ use super::{ivf::storage::IvfModel, sq::ScalarQuantizer, storage::VectorStore}; use crate::frag_reuse::FragReuseIndex; use crate::scalar::RowIdRemapper; use crate::vector::bq::builder::RabitQuantizer; +use crate::vector::bq::storage::RQRowLayout; use crate::{INDEX_METADATA_SCHEMA_KEY, IndexMetadata}; pub trait Quantization: @@ -253,6 +254,12 @@ pub trait QuantizerMetadata: Ok(None) } + /// How the storage file lays out each row's fields. Only an IVF_RQ file + /// can store plane rows; every other storage holds one column per field. + fn row_layout(&self) -> RQRowLayout { + RQRowLayout::Columns + } + async fn load(reader: &V1FileReader) -> Result; } diff --git a/rust/lance-index/src/vector/storage.rs b/rust/lance-index/src/vector/storage.rs index 410cc8fe060..84b82c8125b 100644 --- a/rust/lance-index/src/vector/storage.rs +++ b/rust/lance-index/src/vector/storage.rs @@ -10,10 +10,11 @@ use crate::vector::bq::layered::{ }; use crate::vector::bq::layered_stats; use crate::vector::bq::partition_codes::{PartitionCodes, PartitionCodesKey}; +use crate::vector::bq::plane_rows::{PACKED_COLUMNS, PlaneRowsSpec}; use crate::vector::bq::resident::{ResidentColumnStore, is_resident}; -use crate::vector::bq::storage::normalize_entry_codes; +use crate::vector::bq::storage::{RQRowLayout, normalize_entry_codes}; use crate::vector::bq::transform::ERROR_FACTORS_COLUMN; -use crate::vector::quantizer::{QuantizationMetadata, QuantizerStorage}; +use crate::vector::quantizer::{QuantizationMetadata, QuantizationType, QuantizerStorage}; use arrow::compute::concat_batches; use arrow_array::{ Array, ArrayRef, FixedSizeListArray, RecordBatch, UInt8Array, UInt32Array, UInt64Array, @@ -59,6 +60,7 @@ use crate::{ }, }; +use super::exact_buffers::exact_batch; use super::graph::OrderedFloat; use super::graph::OrderedNode; use super::pairwise::{ @@ -109,25 +111,30 @@ where spawn_cpu(materialize).await } +/// `columns` of `batch`, in that order, sharing its arrays. +fn project_columns(batch: &RecordBatch, columns: &[&str]) -> Result { + let schema = batch.schema(); + let indices = columns + .iter() + .map(|column| schema.index_of(column)) + .collect::, _>>()?; + Ok(batch.project(&indices)?) +} + +/// A prewarm's rows of one partition, slices of a read of several, copied +/// once into buffers of their own, so that the entry keeps neither the +/// whole read nor more than its values (see [`exact_batch`]). fn compact_prewarm_batches(batches: Vec) -> Result { let schema = batches .first() .ok_or_else(|| Error::internal("prewarm partition has no storage batches"))? .schema(); - if batches.len() == 1 { - let batch = batches.into_iter().next().ok_or_else(|| { - Error::internal("prewarm partition storage batch unexpectedly missing") - })?; - if batch.num_rows() == 0 { - Ok(batch) - } else { - Ok(batch.shrink_to_fit()?) - } - } else { - // Concatenation allocates compact output buffers already; do not - // deep-copy them a second time with `shrink_to_fit`. - Ok(concat_batches(&schema, batches.iter())?) + if let [batch] = batches.as_slice() + && batch.num_rows() == 0 + { + return Ok(batch.clone()); } + exact_batch(&schema, &batches, true) } ///
@@ -929,6 +936,20 @@ pub const LAZY_FAR_WINDOW_ENV: &str = "LANCE_RQ_LAZY_FAR_WINDOW"; /// ordinary window under [`LAZY_FAR_WINDOW_ENV`] whose reads have not /// returned (default 64, at least 1). The pool is sized once per index. pub const LAZY_FAR_INFLIGHT_ENV: &str = "LANCE_RQ_LAZY_FAR_INFLIGHT"; +/// When the lazy scan publishes the heap's threshold to the gathers once a +/// lazy probe's survivors first fill the heap to `k` rows: `off` (default) +/// only once the probe is scored whole, as every other probe publishes, or +/// `on` as soon as they do, partway through scoring the probe. A gather +/// waiting for the threshold selects its survivors against the first one +/// published; a mid-probe threshold is the `k`-th best of the rows scored so +/// far, looser than the probe's final one, so such a gather reads more rows. +/// On S3 at 32 concurrent queries, the gathers of probe ranks 1-15 selected +/// about 80x (k=10) and 30x (k=100) as many rows with `on` as with `off`, +/// which served 1.49x and 1.32x the queries per second; at one query in +/// flight the two did not differ. Results are the same either way: the +/// heap's top never increases, so a later threshold still covers every row +/// stage 2, pruning on the live heap, scores. +pub const LAZY_PARTIAL_PUBLISH_ENV: &str = "LANCE_RQ_LAZY_PARTIAL_PUBLISH"; /// Benchmark knob (`0` or `1`, default `0`): `1` loads a layered partition's /// sign plane before its ex planes on backends that do not gate plane /// admission too, as gated backends always do, instead of loading all three @@ -1506,6 +1527,9 @@ pub struct LayeredLazyConfig { pub far_window: usize, /// See [`LAZY_FAR_INFLIGHT_ENV`]. pub far_inflight: usize, + /// See [`LAZY_PARTIAL_PUBLISH_ENV`]: `true` publishes the threshold as + /// soon as a probe's survivors fill the heap. + pub partial_publish: bool, } impl Default for LayeredLazyConfig { @@ -1525,6 +1549,7 @@ impl Default for LayeredLazyConfig { dense_to_eager: DenseToEager::Origin, far_window: DEFAULT_LAZY_FAR_WINDOW, far_inflight: DEFAULT_LAZY_FAR_INFLIGHT, + partial_publish: false, } } } @@ -1646,6 +1671,13 @@ impl LayeredLazyConfig { ) })?; } + if let Some(value) = lookup(LAZY_PARTIAL_PUBLISH_ENV) { + config.partial_publish = match value.trim() { + "on" => true, + "off" => false, + _ => return Err(invalid(LAZY_PARTIAL_PUBLISH_ENV, &value, "on or off")), + }; + } Ok(config) } @@ -1860,6 +1892,10 @@ pub struct IvfQuantizationStorage { /// open of it shares; see [`Self::with_index_file`]. `None` keeps the /// far gather pools to the plane access tracker. index_file: Option, + /// The file's plane-row layout, when its metadata declares one: built + /// and validated when the storage opens the file, or on first use when it + /// is reconstructed from a cache; see [`Self::plane_rows`]. + plane_rows: OnceLock>, } impl DeepSizeOf for IvfQuantizationStorage { @@ -1941,6 +1977,25 @@ impl IvfQuantizationStorage { if let Some(bytes) = quantizer_bytes { metadata.parse_buffer(bytes)?; } + // A file whose schema does not match its declared row layout is an + // invalid index; it fails here rather than on a later read. + let plane_rows = OnceLock::new(); + match metadata.row_layout() { + RQRowLayout::Columns => { + if let Some(name) = PACKED_COLUMNS + .iter() + .find(|name| schema.field(name).is_some()) + { + return Err(Error::index(format!( + "invalid IVF_RQ file: its metadata declares the column layout, but it has plane-row column {name}" + ))); + } + } + RQRowLayout::PlaneRows => { + let spec = Self::plane_rows_spec(&metadata, distance_type, &reader)?; + plane_rows.get_or_init(|| Arc::new(spec)); + } + } Ok(Self { reader, @@ -1958,6 +2013,53 @@ impl IvfQuantizationStorage { packed_metadata: OnceLock::new(), attach_schemas: AttachSchemas::default(), index_file: None, + plane_rows, + }) + } + + /// The plane-row layout of `reader`'s file, an IVF_RQ storage with + /// `metadata`, checked against the file's schema. + fn plane_rows_spec( + metadata: &Q::Metadata, + distance_type: DistanceType, + reader: &FileReader, + ) -> Result { + let Quantizer::Rabit(rq) = Q::from_metadata(metadata, distance_type)? else { + return Err(Error::index( + "invalid index file: only an IVF_RQ file can store plane rows", + )); + }; + let schema = arrow_schema::Schema::from(reader.schema().as_ref()); + PlaneRowsSpec::for_file(rq.metadata_ref(), &schema) + } + + /// How the file stores each row's fields, from its metadata. + pub fn row_layout(&self) -> RQRowLayout { + self.metadata.row_layout() + } + + /// The file's plane-row layout, `None` for the column layout. Reads of a + /// plane-row file go through it: they read the packed columns that hold + /// the fields they want and unpack them into those fields, so every read + /// returns what it returns from a column-layout file. + pub fn plane_rows(&self) -> Result>> { + if self.row_layout() == RQRowLayout::Columns { + return Ok(None); + } + if let Some(spec) = self.plane_rows.get() { + return Ok(Some(spec)); + } + let spec = Self::plane_rows_spec(&self.metadata, self.distance_type, &self.reader)?; + Ok(Some(self.plane_rows.get_or_init(|| Arc::new(spec)))) + } + + /// The fields reads return: the file's schema, or on a plane-row file + /// the column layout's fields its packed columns hold, then the columns + /// it carries. + pub fn logical_schema(&self) -> Result { + Ok(match self.plane_rows()? { + Some(spec) => spec.logical_schema().clone(), + None => self.schema(), }) } @@ -1999,6 +2101,7 @@ impl IvfQuantizationStorage { packed_metadata: OnceLock::new(), attach_schemas: AttachSchemas::default(), index_file: None, + plane_rows: OnceLock::new(), } } @@ -2078,9 +2181,10 @@ impl IvfQuantizationStorage { } /// Set whether reads take the file's small columns from the resident - /// store, resolved when the index opened. + /// store, resolved when the index opened. A plane-row file keeps no + /// small columns apart, so it never reads through the store. pub fn with_resident_columns_enabled(mut self, enabled: bool) -> Self { - self.resident_columns_enabled = enabled; + self.resident_columns_enabled = enabled && self.row_layout() == RQRowLayout::Columns; self } @@ -2093,8 +2197,12 @@ impl IvfQuantizationStorage { } /// Bytes the resident store holds for this storage's file, loaded or - /// not; see [`resident_columns_bytes`]. + /// not; see [`resident_columns_bytes`]. Zero for a plane-row file, whose + /// small columns are packed with the codes of their plane. pub fn resident_columns_bytes(&self) -> u64 { + if self.row_layout() == RQRowLayout::PlaneRows { + return 0; + } resident_columns_bytes(&self.schema(), self.num_rows()) } @@ -2249,6 +2357,8 @@ impl IvfQuantizationStorage { self.distance_type } + /// The file's schema: on a plane-row file, its packed columns; reads + /// return [`Self::logical_schema`]'s fields. pub fn schema(&self) -> SchemaRef { Arc::new(self.reader.schema().as_ref().into()) } @@ -2284,23 +2394,57 @@ impl IvfQuantizationStorage { /// uninstrumented reader is used. /// /// With [`Self::resident_columns_enabled`], only the code and bounds - /// columns are read from the file. + /// columns are read from the file. A plane-row file's partition is read + /// as its packed columns and unpacked. pub async fn load_partition( &self, part_id: usize, io_stats: Option, ) -> Result { + let batch = self + .read_partition_batch(part_id, io_stats.as_ref()) + .await?; + let (batch, remapper) = self.remap_partition_batch(batch).await?; + Q::Storage::try_from_batch_with_remapper( + batch, + self.metadata(), + self.distance_type, + remapper, + ) + } + + /// Every row of partition `part_id` as [`Self::logical_schema`]'s + /// columns, read through a reader that also records into `io_stats`: + /// what [`Self::load_partition`] builds the partition's storage from. + pub async fn read_partition_batch( + &self, + part_id: usize, + io_stats: Option<&IoStats>, + ) -> Result { let range = self.ivf.row_range(part_id); - let schema: SchemaRef = Arc::new(self.reader.schema().as_ref().into()); - let batch = if range.is_empty() { + let schema = self.logical_schema()?; + Ok(if range.is_empty() { RecordBatch::new_empty(schema) - } else if let Some(store) = self.resident_store(&schema, io_stats.as_ref()).await? { + } else if let Some(spec) = self.plane_rows()? { + let batches = self + .file_reader(io_stats, None) + .read_stream( + ReadBatchParams::Range(range), + u32::MAX, + 1, + FilterExpression::no_filter(), + ) + .await? + .try_collect::>() + .await?; + spec.unpack(&batches, &schema)? + } else if let Some(store) = self.resident_store(&schema, io_stats).await? { let params = ReadBatchParams::Range(range.clone()); - let reader = self.file_reader(io_stats.as_ref(), None); + let reader = self.file_reader(io_stats, None); self.read_with_resident_columns(store, schema, range, None, params, &reader) .await? } else { - let reader = self.file_reader(io_stats.as_ref(), None); + let reader = self.file_reader(io_stats, None); let batches = reader .read_stream( ReadBatchParams::Range(range), @@ -2311,32 +2455,35 @@ impl IvfQuantizationStorage { .await? .try_collect::>() .await?; - concat_batches(&schema, batches.iter())? - }; - let (batch, remapper) = self.remap_partition_batch(batch).await?; - Q::Storage::try_from_batch_with_remapper( - batch, - self.metadata(), - self.distance_type, - remapper, - ) + if Q::quantization_type() == QuantizationType::Rabit { + // In buffers of exactly its values, as a plane-row read + // unpacks them, so the partition entries of both layouts + // weigh the same. + exact_batch(&schema, &batches, false)? + } else { + concat_batches(&schema, batches.iter())? + } + }) } /// The file columns a native partition's code-only entry /// ([`PartitionCodes`]) holds, in the file's order: those the resident - /// store does not keep, the codes. - fn partition_code_columns(&self) -> Vec { - self.schema() + /// store does not keep, the codes. A plane-row file names the columns of + /// the column layout, which its reads unpack. + fn partition_code_columns(&self) -> Result> { + Ok(self + .logical_schema()? .fields() .iter() .filter(|field| !is_resident(field)) .map(|field| field.name().clone()) - .collect() + .collect()) } - /// [`Self::partition_code_columns`] as a projection of the file. + /// [`Self::partition_code_columns`] as a projection of a column-layout + /// file. pub fn partition_codes_projection(&self) -> Result { - let columns = self.partition_code_columns(); + let columns = self.partition_code_columns()?; let columns: Vec<&str> = columns.iter().map(String::as_str).collect(); lance_file::versions::reader_projection_from_column_names( self.reader.metadata().version(), @@ -2347,17 +2494,24 @@ impl IvfQuantizationStorage { /// Read the code-only entry of native partition `part_id` from the file, /// through a reader that also records into `io_stats`: its code columns - /// alone, which an object store serves in one request each. + /// alone, which an object store serves in one request each. A plane-row + /// file, which caches whole partitions, gives the same columns cut from + /// its unpacked partition. pub async fn read_partition_codes( &self, part_id: usize, io_stats: Option<&IoStats>, ) -> Result { - let columns = self.partition_code_columns(); + let columns = self.partition_code_columns()?; let columns: Vec<&str> = columns.iter().map(String::as_str).collect(); - let batch = self - .read_file_columns(self.ivf.row_range(part_id), &columns, io_stats) - .await?; + let batch = if self.plane_rows()?.is_some() { + // The code columns are packed with the partition's others. + let partition = self.read_partition_batch(part_id, io_stats).await?; + project_columns(&partition, &columns)? + } else { + self.read_file_columns(self.ivf.row_range(part_id), &columns, io_stats) + .await? + }; Ok(PartitionCodes(self.normalize_entry(batch)?)) } @@ -2497,7 +2651,7 @@ impl IvfQuantizationStorage { pub async fn prewarm_planes(&self, cache: &WeakLanceCache) -> Result<()> { use super::bq::layered::RQPrecision; let planes: Vec = std::iter::once(0) - .chain(self.planes_after_sign(RQPrecision::Full)) + .chain(self.planes_after_sign(RQPrecision::Full)?) .collect(); if !cache.plane_admission_gated() { for part_id in 0..self.num_partitions() { @@ -2546,9 +2700,15 @@ impl IvfQuantizationStorage { /// sign plane leaves them out, then the ex planes the precision scores. /// High precision prunes with the high bounds; full precision prunes /// with the error factors, or with the full bounds on a file without them. - fn planes_after_sign(&self, precision: super::bq::layered::RQPrecision) -> Vec { + fn planes_after_sign(&self, precision: super::bq::layered::RQPrecision) -> Result> { use super::bq::layered::RQPrecision; - let has_error_factors = self.reader.schema().field(ERROR_FACTORS_COLUMN).is_some(); + let has_error_factors = match self.plane_rows()? { + Some(spec) => spec + .logical_schema() + .field_with_name(ERROR_FACTORS_COLUMN) + .is_ok(), + None => self.reader.schema().field(ERROR_FACTORS_COLUMN).is_some(), + }; let (reads_bounds, last_ex_plane) = match precision { RQPrecision::Sign => (false, 0), RQPrecision::High => (true, 1), @@ -2559,7 +2719,7 @@ impl IvfQuantizationStorage { planes.push(SIGN_BOUNDS_PLANE); } planes.extend(1..=last_ex_plane); - planes + Ok(planes) } /// Read independent plane entries through the existing persistent cache. @@ -2579,7 +2739,7 @@ impl IvfQuantizationStorage { } return self.load_partition(part_id, io_stats).await; } - let after_sign = self.planes_after_sign(precision); + let after_sign = self.planes_after_sign(precision)?; let load_plane = |plane| self.load_plane(part_id, plane, cache, io_stats.clone()); let planes = if sequential_plane_loads()? || cache.plane_admission_gated() { // Admit the sign dependency first. The other reads can then @@ -2702,7 +2862,22 @@ impl IvfQuantizationStorage { plane: u8, io_stats: Option, ) -> Result { - let entry_columns = self.entry_columns(); + self.read_plane_entry_with(part_id, plane, self.entry_columns(), io_stats) + .await + } + + /// What a cache entry of plane `plane` of `part_id` holding + /// `entry_columns` holds, read from the file whatever this storage's own + /// entries hold ([`Self::entry_columns`]); see [`Self::read_plane_entry`]. + /// A plane-row file, which caches whole planes, gives the same columns + /// cut from its unpacked plane. + pub async fn read_plane_entry_with( + &self, + part_id: usize, + plane: u8, + entry_columns: EntryColumns, + io_stats: Option, + ) -> Result { if entry_columns == EntryColumns::All { return self.read_plane(part_id, plane, None, io_stats).await; } @@ -2713,10 +2888,15 @@ impl IvfQuantizationStorage { self.sign_bounds ))); } - let range = self.ivf.row_range(part_id); - let batch = self - .read_file_columns(range, &columns, io_stats.as_ref()) - .await?; + let batch = if self.plane_rows()?.is_some() { + // The entry's columns are packed with the plane's others. + let plane = self.read_plane(part_id, plane, None, io_stats).await?; + project_columns(&plane, &columns)? + } else { + let range = self.ivf.row_range(part_id); + self.read_file_columns(range, &columns, io_stats.as_ref()) + .await? + }; self.normalize_entry(batch) } @@ -3080,7 +3260,7 @@ impl IvfQuantizationStorage { let code_only = self.entry_columns() == EntryColumns::Codes; // The planes are independent. Overlap their cache/origin reads, // while preserving plane order when assembling the full row schema. - let planes = std::iter::once(0).chain(self.planes_after_sign(RQPrecision::Full)); + let planes = std::iter::once(0).chain(self.planes_after_sign(RQPrecision::Full)?); let batches = futures::future::try_join_all(planes.map(|plane| { let rows = &rows; let indices = &indices; @@ -3319,12 +3499,7 @@ impl IvfQuantizationStorage { self.sign_bounds ))); } - let projection = lance_file::versions::reader_projection_from_column_names( - self.reader.metadata().version(), - self.reader.schema(), - columns, - )?; - let schema = Arc::new(arrow_schema::Schema::from(projection.schema.as_ref())); + let (projection, schema) = self.read_projection(columns)?; let range = self.ivf.row_range(part_id); if range.is_empty() || rows.as_ref().is_some_and(Vec::is_empty) { return Ok(RecordBatch::new_empty(schema)); @@ -3356,6 +3531,44 @@ impl IvfQuantizationStorage { .read_with_resident_columns(store, schema, range, rows.as_deref(), params, &reader) .await; } + self.read_projected(&reader, params, projection, &schema, rows.is_none()) + .await + } + + /// How a read of the file's `columns` reads them, and the schema of the + /// batch it returns: the columns themselves, or on a plane-row file the + /// packed columns that hold them, which the read unpacks into them. + fn read_projection(&self, columns: &[&str]) -> Result<(ReaderProjection, SchemaRef)> { + let plane_rows = self.plane_rows()?; + let file_columns = match plane_rows { + Some(spec) => spec.packed_for(columns)?, + None => columns.to_vec(), + }; + let projection = lance_file::versions::reader_projection_from_column_names( + self.reader.metadata().version(), + self.reader.schema(), + &file_columns, + )?; + let schema = match plane_rows { + Some(spec) => spec.projection(columns)?, + None => Arc::new(arrow_schema::Schema::from(projection.schema.as_ref())), + }; + Ok((projection, schema)) + } + + /// Read `projection` ([`Self::read_projection`]) at `params` through + /// `reader` into one batch of `schema`, unpacking plane rows. A read of + /// every row of a partition (`whole`), which cache entries hold, returns + /// buffers of exactly its values, as an unpacked plane-row read does; a + /// read of selected rows keeps the decoder's buffers. + async fn read_projected( + &self, + reader: &FileReader, + params: ReadBatchParams, + projection: ReaderProjection, + schema: &SchemaRef, + whole: bool, + ) -> Result { let batches = reader .read_stream_projected( params, @@ -3367,7 +3580,11 @@ impl IvfQuantizationStorage { .await? .try_collect::>() .await?; - concat_batches(&schema, batches.iter()).map_err(Into::into) + match self.plane_rows()? { + Some(spec) => spec.unpack(&batches, schema), + None if whole => exact_batch(schema, &batches, false), + None => concat_batches(schema, batches.iter()).map_err(Into::into), + } } /// This storage's file reader, also recording its I/O into `io_stats`, @@ -3452,7 +3669,12 @@ impl IvfQuantizationStorage { .await? .try_collect::>() .await?; - Some(concat_batches(&file_schema, batches.iter())?) + Some(if rows.is_none() { + // A whole read, which entries hold: see `read_projected`. + exact_batch(&file_schema, &batches, false)? + } else { + concat_batches(&file_schema, batches.iter())? + }) }; store.attach( schema, @@ -3470,28 +3692,19 @@ impl IvfQuantizationStorage { columns: &[&str], io_stats: Option<&IoStats>, ) -> Result { - let projection = lance_file::versions::reader_projection_from_column_names( - self.reader.metadata().version(), - self.reader.schema(), - columns, - )?; - let schema = Arc::new(arrow_schema::Schema::from(projection.schema.as_ref())); + let (projection, schema) = self.read_projection(columns)?; if range.is_empty() { return Ok(RecordBatch::new_empty(schema)); } - let batches = self - .file_reader(io_stats, None) - .read_stream_projected( - ReadBatchParams::Range(range), - u32::MAX, - 1, - projection, - FilterExpression::no_filter(), - ) - .await? - .try_collect::>() - .await?; - Ok(concat_batches(&schema, batches.iter())?) + let reader = self.file_reader(io_stats, None); + self.read_projected( + &reader, + ReadBatchParams::Range(range), + projection, + &schema, + true, + ) + .await } /// Read only `columns` of the given rows; the batch keeps file schema order. @@ -3798,8 +4011,13 @@ impl IvfQuantizationStorage { } // Legacy synchronous remapping path. let frag_reuse_index = self.frag_reuse_index.clone(); + let plane_rows = self.plane_rows()?.cloned(); spawn_prewarm_materialization(move || { - let batch = compact_prewarm_batches(batches)?; + // Unpacking copies the rows into buffers of their own too. + let batch = match plane_rows { + Some(spec) => spec.unpack(&batches, spec.logical_schema())?, + None => compact_prewarm_batches(batches)?, + }; Q::Storage::try_from_batch_with_remapper( batch, &metadata, @@ -3818,13 +4036,13 @@ mod tests { HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES, IndexFileKey, LAZY_DENSE_BYTES_FRACTION_ENV, LAZY_DENSE_ENV, LAZY_DENSE_TO_EAGER_ENV, LAZY_EAGER_BEFORE_FULL_ENV, LAZY_FAR_INFLIGHT_ENV, LAZY_FAR_WINDOW_ENV, LAZY_FULL_ENV, LAZY_INLINE_ROWS_ENV, LAZY_MAX_RUNS_ENV, - LAZY_ORIGIN_GAP_BYTES_ENV, LAZY_ORIGIN_MAX_RUNS_ENV, LAZY_PROMOTE_ENV, - LAZY_PROMOTE_INFLIGHT_ENV, LAZY_WINDOW_ENV, LayeredLazyConfig, LazyFarPermits, - LazyOriginGap, LazyPromotion, ORIGIN_LATENCY_ENV, OriginLatencyClass, PlaneSource, - QueryScratchCapacity, QueryScratchPool, RESIDENT_COLUMNS_ENV, RESIDENT_LIFETIME_ENV, - ResidentColumnsSetting, ResidentLifetime, ResidentStoreSize, SEQUENTIAL_PLANE_LOADS_ENV, - SIGN_BOUNDS_ENV, SignBounds, compact_prewarm_batches, entry_columns_from, - entry_columns_setting, origin_latency_from, origin_latency_setting, + LAZY_ORIGIN_GAP_BYTES_ENV, LAZY_ORIGIN_MAX_RUNS_ENV, LAZY_PARTIAL_PUBLISH_ENV, + LAZY_PROMOTE_ENV, LAZY_PROMOTE_INFLIGHT_ENV, LAZY_WINDOW_ENV, LayeredLazyConfig, + LazyFarPermits, LazyOriginGap, LazyPromotion, ORIGIN_LATENCY_ENV, OriginLatencyClass, + PlaneSource, QueryScratchCapacity, QueryScratchPool, RESIDENT_COLUMNS_ENV, + RESIDENT_LIFETIME_ENV, ResidentColumnsSetting, ResidentLifetime, ResidentStoreSize, + SEQUENTIAL_PLANE_LOADS_ENV, SIGN_BOUNDS_ENV, SignBounds, compact_prewarm_batches, + entry_columns_from, entry_columns_setting, origin_latency_from, origin_latency_setting, origin_reads_whole_plane, plan_plane_gather, resident_columns_from, resident_columns_setting, resident_lifetime_from, resident_lifetime_setting, resident_store_fits, sequential_plane_loads, sequential_plane_loads_from, sign_bounds_from, @@ -3873,6 +4091,24 @@ mod tests { assert_ne!(compact_array.values().as_ptr(), parent_ptr); assert!(compact_array.get_array_memory_size() < parent_size); assert_eq!(compact_array.values(), &(10..20).collect::>()); + // In a buffer of exactly its values, as a plane-row read unpacks + // them, which the pieces of a read join into as well. + let exact_bytes = 10 * size_of::(); + assert_eq!(compact_array.values().inner().capacity(), exact_bytes); + let joined = + compact_prewarm_batches(vec![parent.slice(10, 3), parent.slice(13, 7)]).unwrap(); + assert_eq!(joined, compact); + assert_eq!( + joined.column(0).to_data().buffers()[0].capacity(), + exact_bytes + ); + // An exact batch is still copied, so an entry never shares a read. + let copied = compact_prewarm_batches(vec![compact.clone()]).unwrap(); + assert_eq!(copied, compact); + assert_ne!( + copied.column(0).to_data().buffers()[0].as_ptr(), + compact_array.values().as_ptr() as *const u8 + ); } #[test] @@ -3985,6 +4221,7 @@ mod tests { (LAZY_DENSE_TO_EAGER_ENV, "0"), (LAZY_FAR_WINDOW_ENV, "8"), (LAZY_FAR_INFLIGHT_ENV, " 2 "), + (LAZY_PARTIAL_PUBLISH_ENV, " on "), ]); let config = LayeredLazyConfig::from_lookup(|name| env.get(name).map(|value| value.to_string())) @@ -4006,6 +4243,7 @@ mod tests { dense_to_eager: DenseToEager::Off, far_window: 8, far_inflight: 2, + partial_publish: true, } ); assert_eq!( @@ -4020,6 +4258,16 @@ mod tests { LayeredLazyConfig::default().dense_to_eager, DenseToEager::Origin ); + assert!(!LayeredLazyConfig::default().partial_publish); + let env = HashMap::from([(LAZY_FULL_ENV, "1"), (LAZY_PARTIAL_PUBLISH_ENV, "off")]); + assert_eq!( + LayeredLazyConfig::from_lookup(|name| env.get(name).map(|value| value.to_string())) + .unwrap(), + LayeredLazyConfig { + enabled: true, + ..Default::default() + } + ); // The origin run cap does not depend on the eager-before-full switch, // and dense probes whose planes are on a slow origin go to the eager // scan unless set otherwise. @@ -4067,6 +4315,8 @@ mod tests { (LAZY_DENSE_TO_EAGER_ENV, "on"), (LAZY_FAR_WINDOW_ENV, "wide"), (LAZY_FAR_INFLIGHT_ENV, "0"), + (LAZY_PARTIAL_PUBLISH_ENV, "1"), + (LAZY_PARTIAL_PUBLISH_ENV, "partial"), ] { let enabled = |key: &str| -> Option { if key == LAZY_FULL_ENV { diff --git a/rust/lance/Cargo.toml b/rust/lance/Cargo.toml index 8d5690ec3a6..eddf3477943 100644 --- a/rust/lance/Cargo.toml +++ b/rust/lance/Cargo.toml @@ -87,6 +87,7 @@ snafu = { workspace = true } log = { workspace = true } serde_json = { workspace = true } serde = { workspace = true } +sha2.workspace = true permutation = { version = "0.4.0" } aws-sdk-dynamodb = { workspace = true, optional = true, default-features = false, features = ["default-https-client", "rt-tokio"] } tracing.workspace = true diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 13df58eb95d..0227f712766 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -30,8 +30,8 @@ use roaring::RoaringBitmap; use uuid::Uuid; use super::vector::ivf::{ - SteadyStateRebalance, VectorSegmentCompatibility, index_type_for_segmented_optimize, - optimize_vector_indices, select_steady_state_rebalance, vector_segment_compatibility, + SteadyStateRebalance, VectorSegmentCompatibility, optimize_vector_indices, + segmented_optimize_index_version, select_steady_state_rebalance, vector_segment_compatibility, }; use super::vector::{LogicalVectorIndex, VectorIndexParams, details, fresh_vector_segment_params}; use super::{CreateIndexBuilder, DatasetIndexExt, DatasetIndexInternalExt}; @@ -1245,8 +1245,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( removed_indices: Vec::new(), new_fragment_bitmap: base_unindexed_bitmap, new_dataset_version: dataset.manifest.version, - new_index_version: index_type_for_segmented_optimize(reference_index.as_ref())? - .version(), + new_index_version: segmented_optimize_index_version(reference_index.as_ref())?, new_index_details: reference_metadata .index_details .as_deref() @@ -1433,7 +1432,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( let index_version = if let Some(metadata) = removed_indices.first() { metadata.index_version as u32 } else { - index_type_for_segmented_optimize(reference_index.as_ref())?.version() as u32 + segmented_optimize_index_version(reference_index.as_ref())? as u32 }; Ok(( @@ -2874,10 +2873,11 @@ mod tests { assert_eq!(stats["num_indexed_fragments"], 2); assert_eq!(stats["num_unindexed_fragments"], 0); let appended_segments = dataset.load_indices_by_name("vector_idx").await.unwrap(); + let expected_version = crate::index::create::vector_index_version(&index_params) as i32; assert!( appended_segments .iter() - .all(|segment| segment.index_version == index_params.index_type().version()), + .all(|segment| segment.index_version == expected_version), "append must preserve the storage type's index version" ); let logical_index = dataset @@ -2964,8 +2964,7 @@ mod tests { "the reference-compatible append segment must merge with its source" ); assert_eq!( - merged[0].index_version, - index_params.index_type().version(), + merged[0].index_version, expected_version, "merge must preserve the storage type's index version" ); assert_eq!( diff --git a/rust/lance/src/index/create.rs b/rust/lance/src/index/create.rs index 46b67f38ed1..af6182d94a0 100644 --- a/rust/lance/src/index/create.rs +++ b/rust/lance/src/index/create.rs @@ -25,8 +25,10 @@ use crate::{ use futures::{FutureExt, future::BoxFuture}; use lance_core::datatypes::format_field_path; use lance_index::progress::{IndexBuildProgress, NoopIndexBuildProgress}; +use lance_index::vector::bq::RQRowLayout; use lance_index::{ - IndexParams, IndexType, registry::plugin_name_from_details_url, scalar::CreatedIndex, + IndexParams, IndexType, ivf_rq_index_version, registry::plugin_name_from_details_url, + scalar::CreatedIndex, }; use lance_index::{ metrics::NoOpMetricsCollector, @@ -73,6 +75,22 @@ fn scalar_params_from_inverted(params: &InvertedIndexParams) -> Result u32 { + match params.stages.last() { + Some(StageParams::RQ(rq)) + if params.index_type() == IndexType::IvfRq + || rq.row_layout == RQRowLayout::PlaneRows => + { + ivf_rq_index_version(rq.row_layout) as u32 + } + _ => params.index_type().version() as u32, + } +} + pub struct CreateIndexBuilder<'a> { dataset: &'a mut Dataset, columns: Vec, @@ -547,7 +565,7 @@ impl<'a> CreateIndexBuilder<'a> { .ok_or_else(|| { Error::index("Vector index type must take a VectorIndexParams".to_string()) })?; - let index_version = vec_params.index_type().version() as u32; + let index_version = vector_index_version(vec_params); let effective_fragments = effective_vector_fragments(self.dataset, self.fragments.as_deref()); @@ -4591,8 +4609,26 @@ mod tests { ); } + /// An IVF_RQ index is committed with the version of the row layout it + /// is written in: the column layout keeps version 2, which readers + /// without plane-row support accept, and plane rows take version 3, the + /// highest this build reads. A version above that is left out of the + /// usable indexes. + #[rstest] + #[case::columns(RQRowLayout::Columns, false)] + #[case::plane_rows_native(RQRowLayout::PlaneRows, false)] + #[case::plane_rows_layered(RQRowLayout::PlaneRows, true)] #[tokio::test] - async fn test_create_index_ivf_rq_preserves_index_version() { + async fn test_create_index_ivf_rq_version_follows_row_layout( + #[case] row_layout: RQRowLayout, + #[case] layered: bool, + ) { + use crate::dataset::transaction::Transaction; + use lance_index::pb::VectorIndexDetails; + use lance_index::pb::vector_index_details::{Compression, rabit_quantization}; + use lance_index::vector::bq::RQBuildParams; + use lance_index::{IVF_RQ_COLUMNS_INDEX_VERSION, IVF_RQ_INDEX_VERSION}; + let tmpdir = TempStrDir::default(); let dataset_uri = format!("file://{}", tmpdir.as_str()); @@ -4608,18 +4644,79 @@ mod tests { ); let mut dataset = Dataset::write(reader, &dataset_uri, None).await.unwrap(); - let params = VectorIndexParams::ivf_rq(4, 1, DistanceType::L2); + let num_bits = if layered { 7 } else { 1 }; + let params = VectorIndexParams::with_ivf_rq_params( + DistanceType::L2, + IvfBuildParams::new(4), + RQBuildParams::new(num_bits) + .with_layered(layered) + .with_row_layout(row_layout), + ); + let expected = match row_layout { + RQRowLayout::Columns => IVF_RQ_COLUMNS_INDEX_VERSION as i32, + RQRowLayout::PlaneRows => IVF_RQ_INDEX_VERSION as i32, + }; + assert_eq!(vector_index_version(¶ms) as i32, expected); + assert_eq!(ivf_rq_index_version(row_layout), expected); let committed = dataset .create_index(&["vector"], IndexType::IvfRq, None, ¶ms, false) .await .unwrap(); - - assert_eq!(committed.index_version, IndexType::IvfRq.version()); + assert_eq!(committed.index_version, expected); let loaded = dataset.load_indices_by_name(&committed.name).await.unwrap(); assert_eq!(loaded.len(), 1); - assert_eq!(loaded[0].index_version, IndexType::IvfRq.version()); + assert_eq!(loaded[0].index_version, expected); + let details = loaded[0] + .index_details + .as_ref() + .unwrap() + .to_msg::() + .unwrap(); + let Some(Compression::Rq(rq)) = details.compression else { + panic!("expected RQ details, got {details:?}"); + }; + let expected_layout = match row_layout { + RQRowLayout::Columns => rabit_quantization::RowLayout::Columns, + RQRowLayout::PlaneRows => rabit_quantization::RowLayout::PlaneRows, + }; + assert_eq!(rq.row_layout(), expected_layout); + assert_eq!(rq.layered, layered); + + // The index answers queries. + let query = arrow_array::Float32Array::from(vec![0.5f32; 16]); + let results = dataset + .scan() + .nearest("vector", &query, 5) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 5); + + // A version above what this build reads is left out. + let mut unreadable = loaded.clone(); + unreadable[0].index_version = IndexType::IvfRq.version() + 1; + let transaction = Transaction::new( + dataset.manifest.version, + Operation::CreateIndex { + new_indices: unreadable, + removed_indices: loaded, + }, + None, + ); + dataset + .apply_commit(transaction, &Default::default(), &Default::default()) + .await + .unwrap(); + assert!( + dataset + .load_indices_by_name(&committed.name) + .await + .unwrap() + .is_empty() + ); } #[tokio::test] diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index 09ad23cb3b5..5aab7835f82 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -2222,6 +2222,7 @@ fn derive_rabit_params(rabit_quantizer: &RabitQuantizer) -> RQBuildParams { RQBuildParams { num_bits: rabit_quantizer.num_bits(), layered: rabit_quantizer.metadata_ref().layered, + row_layout: rabit_quantizer.metadata_ref().row_layout, rotation_type: rabit_quantizer.rotation_type(), rotation: None, } diff --git a/rust/lance/src/index/vector/builder.rs b/rust/lance/src/index/vector/builder.rs index 88122af00a3..1974bbe963c 100644 --- a/rust/lance/src/index/vector/builder.rs +++ b/rust/lance/src/index/vector/builder.rs @@ -39,13 +39,14 @@ use lance_index::metrics::NoOpMetricsCollector; use lance_index::optimize::OptimizeOptions; use lance_index::progress::{IndexBuildProgress, NoopIndexBuildProgress}; use lance_index::scalar::RowIdRemapper; -use lance_index::vector::bq::storage::{RABIT_CODE_COLUMN, unpack_codes}; +use lance_index::vector::bq::plane_rows::{PlaneRowsSpec, requests_fullzip}; +use lance_index::vector::bq::storage::{RABIT_CODE_COLUMN, RQRowLayout, unpack_codes}; use lance_index::vector::hnsw::HNSW; use lance_index::vector::hnsw::remap::{remap_graph_batch, remap_graph_repair}; use lance_index::vector::kmeans::KMeansParams; use lance_index::vector::pq::storage::transpose; use lance_index::vector::quantizer::{ - QuantizationMetadata, QuantizationType, QuantizerBuildParams, + QuantizationMetadata, QuantizationType, Quantizer, QuantizerBuildParams, }; use lance_index::vector::quantizer::{QuantizerMetadata, QuantizerStorage}; use lance_index::vector::shared::{SupportedIvfIndexType, write_unified_ivf_and_index_metadata}; @@ -1716,12 +1717,28 @@ impl IvfIndexBuilder let storage_path = self.index_dir.clone().join(INDEX_AUXILIARY_FILE_NAME); let index_path = self.index_dir.clone().join(INDEX_FILE_NAME); + // An IVF_RQ quantizer whose metadata asks for plane rows (a new + // build's params, or the index an append, optimize or remap extends) + // writes its storage as one fixed-width row column per plane, every + // partition's rows packed after the sign codes take their stored form. + let plane_rows = match Into::::into(quantizer.clone()) { + Quantizer::Rabit(rq) if rq.metadata_ref().row_layout == RQRowLayout::PlaneRows => { + Some(PlaneRowsSpec::try_new(rq.metadata_ref())?) + } + _ => None, + }; let writer_options = FileWriterOptions::default(); let mut storage_writer = if is_flat { None } else { - let mut fields = vec![ROW_ID_FIELD.clone(), quantizer.field()]; - fields.extend(quantizer.extra_fields()); + let fields = match &plane_rows { + Some(spec) => spec.file_fields(requests_fullzip(self.format_version)), + None => { + let mut fields = vec![ROW_ID_FIELD.clone(), quantizer.field()]; + fields.extend(quantizer.extra_fields()); + fields + } + }; let storage_schema: Schema = (&arrow_schema::Schema::new(fields)).try_into()?; Some(file_versions::create_writer( self.format_version, @@ -1808,6 +1825,9 @@ impl IvfIndexBuilder let unpacked = Arc::new(unpack_codes(codes_fsl)); batch = batch.replace_column_by_name(RABIT_CODE_COLUMN, unpacked)?; } + if let Some(spec) = &plane_rows { + batch = spec.pack(&batch)?; + } if storage_writer.is_none() { let storage_schema: Schema = batch.schema_ref().as_ref().try_into()?; diff --git a/rust/lance/src/index/vector/details.rs b/rust/lance/src/index/vector/details.rs index 1bef3f39947..2ebc53ef383 100644 --- a/rust/lance/src/index/vector/details.rs +++ b/rust/lance/src/index/vector/details.rs @@ -26,7 +26,7 @@ use lance_linalg::distance::DistanceType; use lance_table::format::IndexMetadata; use serde::Serialize; -use lance_index::vector::bq::{RQBuildParams, RQRotationType}; +use lance_index::vector::bq::{RQBuildParams, RQRotationType, RQRowLayout}; use lance_index::vector::hnsw::builder::HnswBuildParams; use lance_index::vector::ivf::IvfBuildParams; use lance_index::vector::pq::PQBuildParams; @@ -295,7 +295,13 @@ pub fn vector_params_from_details(details: &prost_types::Any) -> Option RQRowLayout::Columns, + rabit_quantization::RowLayout::PlaneRows => RQRowLayout::PlaneRows, + }, + ), ) } (Some(hnsw), Some(Compression::Pq(pq))) => VectorIndexParams::with_ivf_hnsw_pq_params( @@ -688,9 +694,13 @@ async fn convert_v3_metadata_to_details( rabit_quantization::RotationType::Matrix } }; + let row_layout = match rq.row_layout { + RQRowLayout::Columns => rabit_quantization::RowLayout::Columns, + RQRowLayout::PlaneRows => rabit_quantization::RowLayout::PlaneRows, + }; Some(Compression::Rq(RabitQuantization { layered: rq.layered, - row_layout: rabit_quantization::RowLayout::Columns as i32, + row_layout: row_layout.into(), num_bits: rq.num_bits as u32, rotation_type: rotation_type.into(), })) @@ -1375,6 +1385,7 @@ mod tests { IvfSq, IvfRqMatrix, IvfRqFast, + IvfRqPlaneRows, IvfHnswFlat, IvfHnswPq, IvfHnswSq, @@ -1382,7 +1393,7 @@ mod tests { fn build_roundtrip_params(combo: Combo, metric: DistanceType) -> VectorIndexParams { use crate::index::vector::VectorIndexParams; - use lance_index::vector::bq::{RQBuildParams, RQRotationType}; + use lance_index::vector::bq::{RQBuildParams, RQRotationType, RQRowLayout}; use lance_index::vector::hnsw::builder::HnswBuildParams; use lance_index::vector::ivf::builder::IvfBuildParams; use lance_index::vector::pq::builder::PQBuildParams; @@ -1431,6 +1442,13 @@ mod tests { ivf, RQBuildParams::with_rotation_type(1, RQRotationType::Fast), ), + Combo::IvfRqPlaneRows => VectorIndexParams::with_ivf_rq_params( + metric, + ivf, + RQBuildParams::with_rotation_type(7, RQRotationType::Fast) + .with_layered(true) + .with_row_layout(RQRowLayout::PlaneRows), + ), Combo::IvfHnswFlat => VectorIndexParams::ivf_hnsw(metric, ivf, hnsw), Combo::IvfHnswPq => VectorIndexParams::with_ivf_hnsw_pq_params(metric, ivf, hnsw, pq), Combo::IvfHnswSq => VectorIndexParams::with_ivf_hnsw_sq_params(metric, ivf, hnsw, sq), @@ -1443,6 +1461,7 @@ mod tests { #[case::ivf_sq(Combo::IvfSq)] #[case::ivf_rq_matrix(Combo::IvfRqMatrix)] #[case::ivf_rq_fast(Combo::IvfRqFast)] + #[case::ivf_rq_plane_rows(Combo::IvfRqPlaneRows)] #[case::ivf_hnsw_flat(Combo::IvfHnswFlat)] #[case::ivf_hnsw_pq(Combo::IvfHnswPq)] #[case::ivf_hnsw_sq(Combo::IvfHnswSq)] @@ -1451,7 +1470,7 @@ mod tests { #[values(DistanceType::L2, DistanceType::Cosine)] metric: DistanceType, ) { use crate::index::vector::StageParams; - use lance_index::vector::bq::RQRotationType; + use lance_index::vector::bq::{RQRotationType, RQRowLayout}; let params = build_roundtrip_params(combo, metric); @@ -1503,6 +1522,16 @@ mod tests { _ => unreachable!(), }; assert_eq!(rq.rotation_type, expected); + assert_eq!(rq.row_layout, RQRowLayout::Columns); + assert!(!rq.layered); + } + Combo::IvfRqPlaneRows => { + let StageParams::RQ(rq) = &restored.stages[1] else { + panic!("expected RQ stage"); + }; + assert_eq!(rq.num_bits, 7); + assert!(rq.layered); + assert_eq!(rq.row_layout, RQRowLayout::PlaneRows); } Combo::IvfHnswFlat => { let StageParams::Hnsw(hnsw) = &restored.stages[1] else { diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 68e9e30ddc6..8be1688b7be 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -392,6 +392,22 @@ pub(crate) fn index_type_for_segmented_optimize(index: &dyn VectorIndex) -> Resu IndexType::try_from(index_type_string(sub_index_type, quantization_type).as_str()) } +/// The index version of a segment that optimizing `index` writes, which keeps +/// `index`'s row layout: an IVF_RQ segment's follows that layout (see +/// [`lance_index::ivf_rq_index_version`]), so a column-layout index keeps its +/// version; any other segment's is its type's. +pub(crate) fn segmented_optimize_index_version(index: &dyn VectorIndex) -> Result { + let index_type = index_type_for_segmented_optimize(index)?; + if index_type != IndexType::IvfRq { + return Ok(index_type.version()); + } + let row_layout = match index.quantizer() { + Quantizer::Rabit(rq) => rq.metadata_ref().row_layout, + _ => lance_index::vector::bq::RQRowLayout::Columns, + }; + Ok(lance_index::ivf_rq_index_version(row_layout)) +} + /// What a steady-state optimize (no new rows) should do to a logical IVF index. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum SteadyStateRebalance { @@ -484,13 +500,35 @@ pub(crate) enum VectorSegmentCompatibility { QueryCompatibleModelsDiffer, } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq)] enum VectorModelMismatch { StorageFormat, + /// IVF_RQ segments that store their rows in different layouts, the + /// first segment's and another's described. + RabitLayout { + expected: String, + found: String, + }, IvfCentroids, QuantizerMetadata, } +/// How an IVF_RQ quantizer lays out the rows its index stores, layered or +/// native and in columns or plane rows, as `native/columns`. `None` for +/// other quantizers. +fn rabit_layout(quantizer: &Quantizer) -> Option { + let Quantizer::Rabit(quantizer) = quantizer else { + return None; + }; + let metadata = quantizer.metadata(None); + let kind = if metadata.layered { + "layered" + } else { + "native" + }; + Some(format!("{kind}/{}", metadata.row_layout)) +} + fn vector_index_dimension(index: &dyn VectorIndex) -> usize { let ivf_dimension = index.ivf_model().dimension(); if ivf_dimension != 0 { @@ -576,7 +614,16 @@ fn vector_model_mismatch(indices: &[Arc]) -> Option {} } - if !shared_quantizer_model(&first_quantizer, &index.quantizer()) { + let quantizer = index.quantizer(); + // Segments laid out differently cannot be merged, whatever model + // they share: the merged segment has a single layout. + if let (Some(expected), Some(found)) = + (rabit_layout(&first_quantizer), rabit_layout(&quantizer)) + && expected != found + { + return Some(VectorModelMismatch::RabitLayout { expected, found }); + } + if !shared_quantizer_model(&first_quantizer, &quantizer) { return Some(VectorModelMismatch::QuantizerMetadata); } } @@ -603,6 +650,10 @@ fn validate_shared_vector_model(indices: &[Arc], operation: &st Some(VectorModelMismatch::StorageFormat) => Err(Error::index(format!( "{operation}: vector index segments do not share a storage format" ))), + Some(VectorModelMismatch::RabitLayout { expected, found }) => Err(Error::index(format!( + "{operation}: IVF_RQ segments of different row layouts cannot be merged: the \ + first segment is {expected} and another {found}" + ))), Some(VectorModelMismatch::IvfCentroids) => Err(Error::index(format!( "{operation}: vector index segments do not share IVF centroids" ))), @@ -638,6 +689,8 @@ fn shared_quantizer_model(left: &Quantizer, right: &Quantizer) -> bool { && left.code_dim == right.code_dim && left.num_bits == right.num_bits && left.packed == right.packed + && left.layered == right.layered + && left.row_layout == right.row_layout && left.query_estimator == right.query_estimator && left.fast_rotation_signs == right.fast_rotation_signs && match (&left.rotate_mat, &right.rotate_mat) { diff --git a/rust/lance/src/index/vector/ivf/partition_serde.rs b/rust/lance/src/index/vector/ivf/partition_serde.rs index e4dbf2c5268..a97fbe7a611 100644 --- a/rust/lance/src/index/vector/ivf/partition_serde.rs +++ b/rust/lance/src/index/vector/ivf/partition_serde.rs @@ -22,7 +22,9 @@ use lance_core::{Error, Result}; use lance_index::vector::bq::RQRotationType; use lance_index::vector::bq::builder::RabitQuantizer; use lance_index::vector::bq::raw_body::{read_raw_batch, write_raw_batch}; -use lance_index::vector::bq::storage::{RabitQuantizationMetadata, RabitQueryEstimator}; +use lance_index::vector::bq::storage::{ + RQRowLayout, RabitQuantizationMetadata, RabitQueryEstimator, +}; use lance_index::vector::flat::index::{FlatBinQuantizer, FlatMetadata, FlatQuantizer}; use lance_index::vector::pq::ProductQuantizer; use lance_index::vector::pq::storage::ProductQuantizationMetadata; @@ -502,6 +504,9 @@ impl CacheCodecImpl for PartitionEntry { // The storage batch already has packed codes; skip re-packing. packed: true, layered: header.layered, + // The storage batch holds one column per field whatever the + // index file's layout. + row_layout: RQRowLayout::Columns, query_estimator, }; let storage = ::Storage::try_from_batch( diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 52ca37339d0..f80f3d37d0f 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -54,7 +54,7 @@ use lance_index::vector::bq::layered::{EntryColumns, SignBounds}; use lance_index::vector::bq::layered_stats; use lance_index::vector::bq::partition_codes::PartitionCodesKey; use lance_index::vector::bq::rabit_ex_bits; -use lance_index::vector::bq::storage::{RabitQueryEstimator, SEGMENT_NUM_CODES}; +use lance_index::vector::bq::storage::{RQRowLayout, RabitQueryEstimator, SEGMENT_NUM_CODES}; use lance_index::vector::flat::index::{FlatBinQuantizer, FlatIndex, FlatQuantizer}; use lance_index::vector::graph::OrderedNode; use lance_index::vector::hnsw::HNSW; @@ -67,8 +67,9 @@ use lance_index::vector::sq::ScalarQuantizer; use lance_index::vector::storage::{ IndexFileKey, LayeredLazyConfig, OriginLatencyClass, PlaneAccessTracker, QueryResidual, QueryScratch, QueryScratchCapacity, QueryScratchPool, RabitRawQueryContext, ResidentColumns, - ResidentColumnsSetting, VectorStore, entry_columns_setting, origin_latency_setting, - resident_columns_setting, resident_lifetime_setting, resident_store_fits, sign_bounds_setting, + ResidentColumnsSetting, ResidentStoreSize, VectorStore, entry_columns_setting, + origin_latency_setting, resident_columns_setting, resident_lifetime_setting, + resident_store_fits, sign_bounds_setting, }; use lance_index::vector::v3::subindex::SubIndexType; use lance_index::{ @@ -100,6 +101,7 @@ use uuid::Uuid; use super::{IvfIndexPartitionStatistics, IvfIndexStatistics, maybe_centroids_for_stats}; mod lazy_full; +pub mod row_layout; pub(crate) type RabitSearchCacheCell = Arc>>>>; @@ -1978,10 +1980,26 @@ impl IVFIndex { /// Bytes the index's resident columns take, whether or not it keeps /// them; see [`lance_index::vector::storage::resident_columns_bytes`]. + /// Zero for a plane-row file, which never keeps a resident store. pub fn resident_columns_bytes(&self) -> u64 { self.storage.resident_columns_bytes() } + /// The size of the index's resident store, whether or not it keeps + /// one, and what an index cache charges for it once loaded, which the + /// open compared with the cache's pinned cap; see + /// [`ResidentColumnsSetting::admits`]. Zero for a plane-row file. + pub fn resident_store_size(&self) -> Result { + self.storage.resident_store_size() + } + + /// How the index's auxiliary file stores each row's fields. Only an + /// IVF_RQ index can store plane rows; reads unpack them, so results do + /// not depend on the layout. + pub fn row_layout(&self) -> RQRowLayout { + self.storage.row_layout() + } + /// Replace whether reads keep the small columns resident, so tests can /// open both in one process. #[cfg(test)] @@ -8765,6 +8783,9 @@ mod tests { /// settings route every predicted-dense probe (`all`), and the far /// window never applies. Promotions are off unless set, so some /// settings that gather sparse rows promote a plane after one read. + /// Thresholds are published once a probe is scored whole unless set, + /// so some settings publish one as soon as a lazy probe's survivors + /// fill the heap, partway through scoring it. /// The count of settings is prime, so over the parity test's cases /// the rotation pairs every setting with every `k`, filter, /// approximation mode and bounds choice. @@ -8840,6 +8861,21 @@ mod tests { dense_to_eager: DenseToEager::Off, ..enabled }, + // The first probe stays lazy, and its survivors filling the + // heap publish the threshold before it is scored whole. + LayeredLazyConfig { + dense_to_eager: DenseToEager::Off, + partial_publish: true, + ..enabled + }, + // Gathers wait for the heap to fill, which a lazy probe + // publishes partway through its scoring. + LayeredLazyConfig { + window: 4, + eager_before_full: false, + partial_publish: true, + ..enabled + }, ]; // On a high-latency origin a probe whose high or low plane no // cache tier holds is gathered within the far window. @@ -8869,6 +8905,21 @@ mod tests { promote: promote_once, ..enabled }, + // The defaults with mid-probe publishes. + LayeredLazyConfig { + partial_publish: true, + ..enabled + }, + // Far gathers released by a mid-probe publish once the heap + // fills, beyond no window. + LayeredLazyConfig { + window: 0, + eager_before_full: false, + dense_to_eager: DenseToEager::Off, + far_window: usize::MAX, + partial_publish: true, + ..enabled + }, ]; let low = low .into_iter() @@ -9183,9 +9234,47 @@ mod tests { } assert_eq!(stats.needed_not_fetched, 0, "{context}"); assert_lazy_chain_timing(&stats, context); + assert_partial_publish_counts(&stats, config, context); stats } + /// Check the mid-probe publish counters of `stats`, taken over lazy + /// scans under `config`. A scan's heap fills once, so it publishes a + /// mid-probe threshold at most once, and never with partial + /// publishes off; the gathers that take one are among the gathers, + /// never the first probe's, which is issued before its scoring. + fn assert_partial_publish_counts( + stats: &LayeredLazyStats, + config: LayeredLazyConfig, + context: &str, + ) { + let publishes: u64 = stats.mid_probe_full_publishes.iter().sum(); + let issues: u64 = stats.partial_threshold_issues.iter().sum(); + assert!(publishes <= stats.lazy_queries, "{stats:?} {context}"); + if !config.partial_publish { + assert_eq!( + ( + publishes, + issues, + stats.partial_threshold_rows.iter().sum::() + ), + (0, 0, 0), + "{context}" + ); + } + if publishes == 0 { + assert_eq!(issues, 0, "{stats:?} {context}"); + } + assert_eq!(stats.partial_threshold_issues[0], 0, "{context}"); + for bucket in 0..RANK_BUCKETS { + assert!( + stats.partial_threshold_issues[bucket] <= stats.lazy_probes[bucket] + && stats.partial_threshold_rows[bucket] <= stats.rows_fetched[bucket], + "bucket {bucket}: {stats:?} {context}" + ); + } + } + /// Check the release-chain timings of `stats`, taken over lazy scans /// that all completed. An eager scan times nothing; every lazy scan /// times its first probe's scoring once; a wait is timed only when it @@ -9449,6 +9538,150 @@ mod tests { ); } + /// A plane-row copy of the fixture answers every query, precision, + /// filter, cache and origin class as the column layout does, bit for + /// bit: eager, with a cascade, and lazily under every scan setting + /// (far window, one far permit, dense-to-eager routing, promotions). + #[rstest] + #[case::rq7_l2(7, DistanceType::L2, true)] + #[case::rq9_dot(9, DistanceType::Dot, true)] + #[case::native_rq7_l2(7, DistanceType::L2, false)] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_plane_rows_search_matches_columns( + #[case] bits: u8, + #[case] distance_type: DistanceType, + #[case] layered: bool, + ) { + use super::super::row_layout::rewrite_rq_row_layout; + use lance_index::vector::bq::storage::RQRowLayout; + let _serial = LAZY_TEST_LOCK.lock().await; + let columns_dir = TempStrDir::default(); + let (_, batch) = + write_rq_test_dataset(columns_dir.as_str(), bits, distance_type, layered).await; + let rows_dir = TempStrDir::default(); + let rows_path = format!("{}/plane_rows", rows_dir.as_str()); + copy_test_dir( + std::path::Path::new(columns_dir.as_str()), + std::path::Path::new(&rows_path), + ); + let mut converted = Dataset::open(&rows_path).await.unwrap(); + let name = converted.load_indices().await.unwrap()[0].name.clone(); + rewrite_rq_row_layout(&mut converted, &name, RQRowLayout::PlaneRows) + .await + .unwrap(); + + let vectors = batch["vector"].as_fixed_size_list(); + let keys = [vectors.value(0), vectors.value(777)]; + let filters = lazy_test_filters(); + let configs = lazy_test_configs(); + let precisions: &[RQPrecision] = if layered { + &[RQPrecision::Full, RQPrecision::High, RQPrecision::Sign] + } else { + &[RQPrecision::Full] + }; + let mut case = 0usize; + for cache in [ + LazyTestCache::Origin, + LazyTestCache::Small, + LazyTestCache::Ungated, + ] { + let (columns, _, _) = open_lazy_test_index(columns_dir.as_str(), cache).await; + let (rows, row_index, _) = open_lazy_test_index(&rows_path, cache).await; + assert_eq!(lazy_index(&row_index).row_layout(), RQRowLayout::PlaneRows); + for class in [OriginLatencyClass::Low, OriginLatencyClass::High] { + let open = move |dataset: Dataset| async move { + let (store, dir) = index_files(&dataset).await; + open_with_origin_latency(&dataset, store, dir, class).await + }; + let column = open(columns.clone()).await; + let row = open(rows.clone()).await; + assert!(!lazy_index(&row).resident_columns_enabled()); + for key in &keys { + for nprobes in [1, 4, LAZY_PARTITIONS] { + for k in [1, 10, 100, LAZY_ROWS + 1] { + for &precision in precisions { + case += 1; + let (filter_name, filter) = &filters[case % filters.len()]; + let mut query = lazy_test_query(key.clone(), k, nprobes); + query.rq_precision = precision; + // A cascade scores a partition's first stage at + // high precision only while its high plane is + // resident, so it compares the layouts where + // both caches hold the same planes: a + // high-latency column index caches code-only + // entries, which the plane-row index cannot. + if layered + && precision == RQPrecision::Full + && case.is_multiple_of(4) + && (class == OriginLatencyClass::Low + || cache == LazyTestCache::Origin) + { + query.rq_cascade_factor = Some(4); + } + let context = format!( + "cache={cache:?} class={class} nprobes={nprobes} k={k} precision={precision:?} cascade={:?} filter={filter_name}", + query.rq_cascade_factor + ); + assert_eq!( + result_bits( + &search_global(&column, &query, filter.clone()) + .await + .unwrap() + ), + result_bits( + &search_global(&row, &query, filter.clone()) + .await + .unwrap() + ), + "{context}" + ); + } + } + } + } + if !layered { + continue; + } + // Every lazy setting of the class on the widest probe set. + let query = lazy_test_query(keys[0].clone(), 100, LAZY_PARTITIONS); + let eager = result_bits( + &search_global(&column, &query, filters[0].1.clone()) + .await + .unwrap(), + ); + for &(config_class, config) in &configs { + if config_class != class { + continue; + } + let context = format!("cache={cache:?} class={class} config={config:?}"); + assert_lazy_matches_eager(&row, &query, &filters[0].1, config, &context) + .await; + lazy_index(&row).set_layered_lazy_config_for_test(config); + let lazy = search_global(&row, &query, filters[0].1.clone()) + .await + .unwrap(); + lazy_index(&row) + .set_layered_lazy_config_for_test(LayeredLazyConfig::default()); + assert_eq!(result_bits(&lazy), eager, "{context}"); + } + wait_for_promotions().await; + } + } + } + + fn copy_test_dir(from: &std::path::Path, to: &std::path::Path) { + std::fs::create_dir_all(to).unwrap(); + for entry in std::fs::read_dir(from).unwrap() { + let entry = entry.unwrap(); + let target = to.join(entry.file_name()); + if entry.file_type().unwrap().is_dir() { + copy_test_dir(&entry.path(), &target); + } else { + std::fs::copy(entry.path(), target).unwrap(); + } + } + } + /// An ungated backend is warmed partition by partition, so every plane /// is persisted even when RAM holds only part of the sign planes, and /// warm queries read nothing from the origin file, whether the first @@ -10991,6 +11224,81 @@ mod tests { assert_eq!(off.eager_before_full, 0, "{off:?}"); } + /// A lazy first probe whose survivors fill the heap publishes the + /// heap's top partway through its scoring, once per query at rank 0, + /// only with partial publishes on; the probe stays certain dense + /// either way. Routed to the eager scan, which publishes only once a + /// probe is scored whole, it publishes no mid-probe threshold. Every + /// setting matches the eager scan, and the gathers that took a + /// mid-probe threshold are counted within the lazy gathers (see + /// `assert_partial_publish_counts`). + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_layered_lazy_partial_publish_counts_mid_probe_thresholds() { + const PROBES: usize = 32; + const QUERIES: usize = 8; + const STAGING_STEPS: usize = 8; + const K: usize = 10; + let _serial = LAZY_TEST_LOCK.lock().await; + let dir = TempStrDir::default(); + let (_, batch) = write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; + let (_dataset, index, _) = + open_lazy_test_index(dir.as_str(), LazyTestCache::Origin).await; + let ivf = lazy_index(&index); + let vectors = batch["vector"].as_fixed_size_list(); + let filter: Arc = Arc::new(NoFilter); + let queries = queries_filled_by_first_probe(&index, vectors, K, PROBES, QUERIES); + let mut first_probe = [0; RANK_BUCKETS]; + first_probe[0] = QUERIES as u64; + ivf.set_lazy_prepare_parallelism_for_test(STAGING_STEPS); + for (dense_to_eager, partial_publish) in [ + (DenseToEager::Off, true), + (DenseToEager::Off, false), + (DenseToEager::All, true), + ] { + let config = LayeredLazyConfig { + partial_publish, + ..issue_policy(true, dense_to_eager) + }; + let mut publishes = [0; RANK_BUCKETS]; + let mut certain_dense = [0; RANK_BUCKETS]; + let mut routed = [0; RANK_BUCKETS]; + for (position, query) in queries.iter().enumerate() { + let context = format!( + "dense_to_eager={dense_to_eager} partial_publish={partial_publish} query={position}" + ); + let stats = + assert_lazy_matches_eager(&index, query, &filter, config, &context).await; + assert_eq!(stats.lazy_queries, 1, "{context}"); + for (totals, counts) in [ + (&mut publishes, stats.mid_probe_full_publishes), + (&mut certain_dense, stats.certain_dense), + (&mut routed, stats.dense_to_eager), + ] { + for (total, count) in totals.iter_mut().zip(counts) { + *total += count; + } + } + } + let context = + format!("dense_to_eager={dense_to_eager} partial_publish={partial_publish}"); + let lazy_first_probe = dense_to_eager == DenseToEager::Off; + let expected_publishes = if lazy_first_probe && partial_publish { + first_probe + } else { + [0; RANK_BUCKETS] + }; + assert_eq!(publishes, expected_publishes, "{context}"); + let (expected_dense, expected_routed) = if lazy_first_probe { + (first_probe, [0; RANK_BUCKETS]) + } else { + ([0; RANK_BUCKETS], first_probe) + }; + assert_eq!(certain_dense, expected_dense, "{context}"); + assert_eq!(routed, expected_routed, "{context}"); + } + ivf.set_lazy_prepare_parallelism_for_test(0); + } + /// A sparse gather the persistent tier cannot serve loads the whole /// plane once its origin row runs exceed the cap, whether gathers /// before the heap fills are deferred or not, with identical results. @@ -13197,16 +13505,16 @@ mod tests { assert_eq!(results[0], results[1]); } - /// Gaps that coalesced origin reads are checked at, in ascending - /// order so that neighbours compare a narrower gap with a wider one: - /// touching ranges only, V11's S3 block size, a narrower gap for a - /// high-latency origin, the gap `auto` gives one, and every range of - /// a column page. + /// Gaps that coalesced origin reads are checked at, in strictly + /// ascending order so that neighbours compare a narrower gap with a + /// wider one: touching ranges only, V11's S3 block size, the gap + /// `auto` gives a high-latency origin, the wider 1 MiB it gave + /// before, and every range of a column page. const COALESCE_TEST_GAPS: [u64; 5] = [ 0, 64 * 1024, - 256 * 1024, HIGH_LATENCY_LAZY_ORIGIN_GAP_BYTES, + 1024 * 1024, u64::MAX, ]; /// Spacing of the single rows of a sparse selection. @@ -13227,7 +13535,10 @@ mod tests { async fn test_layered_read_plane_matches_across_coalesce_gaps( #[case] version: LanceFileVersion, ) { - assert!(COALESCE_TEST_GAPS.is_sorted(), "{COALESCE_TEST_GAPS:?}"); + assert!( + COALESCE_TEST_GAPS.is_sorted_by(|narrower, wider| narrower < wider), + "{COALESCE_TEST_GAPS:?}" + ); let _serial = LAZY_TEST_LOCK.lock().await; let dir = TempStrDir::default(); write_lazy_test_dataset(dir.as_str(), 7, DistanceType::L2).await; @@ -13334,7 +13645,10 @@ mod tests { #[case] bits: u8, #[case] distance_type: DistanceType, ) { - assert!(COALESCE_TEST_GAPS.is_sorted(), "{COALESCE_TEST_GAPS:?}"); + assert!( + COALESCE_TEST_GAPS.is_sorted_by(|narrower, wider| narrower < wider), + "{COALESCE_TEST_GAPS:?}" + ); let _serial = LAZY_TEST_LOCK.lock().await; let dir = TempStrDir::default(); let (_, batch) = write_lazy_test_dataset(dir.as_str(), bits, distance_type).await; diff --git a/rust/lance/src/index/vector/ivf/v2/lazy_full.rs b/rust/lance/src/index/vector/ivf/v2/lazy_full.rs index eed33f3762e..544527a9aad 100644 --- a/rust/lance/src/index/vector/ivf/v2/lazy_full.rs +++ b/rust/lance/src/index/vector/ivf/v2/lazy_full.rs @@ -27,7 +27,12 @@ //! waiting would not save reads (`LANCE_RQ_LAZY_EAGER_BEFORE_FULL`): then //! it is gathered at once with `T` = +inf, selecting every accepted row as //! the eager load does; -//! - the query task scores ready probes on one heap and publishes progress. +//! - the query task scores ready probes on one heap and publishes progress +//! once each probe is scored whole. With `LANCE_RQ_LAZY_PARTIAL_PUBLISH=on`, +//! a lazy probe whose survivors first fill the heap publishes the heap's +//! top as soon as they do, partway through the probe; that threshold is +//! looser than the probe's final one, so the gathers it releases read more +//! rows. //! //! Every staging and gather step runs as its own task, so it makes progress //! whether or not the producer is polling the buffer that holds it; see @@ -121,6 +126,48 @@ struct LazyProgress { full: bool, /// The heap's top when `full`, otherwise +inf. threshold: f32, + /// Whether `threshold` was published partway through scoring the probe + /// after the `scored` ones, when its survivors first filled the heap: + /// the `k`-th best of the rows scored so far. Cleared once that probe is + /// scored whole; only set while `full`. + partial: bool, +} + +impl LazyProgress { + /// No probe scored and the heap empty. + fn start() -> Self { + Self { + scored: 0, + full: false, + threshold: f32::INFINITY, + partial: false, + } + } + + /// The `scored`-th probe was scored whole; `top` is the heap's top if + /// the heap holds `k` rows. + fn probe_scored(&mut self, scored: usize, top: Option) { + self.scored = scored; + self.partial = false; + if let Some(top) = top { + self.full = true; + self.threshold = top; + } + } + + /// The survivors of the probe being scored first filled the heap, whose + /// top is then `threshold`. + fn filled_mid_probe(&mut self, threshold: f32) { + self.full = true; + self.threshold = threshold; + self.partial = true; + } + + /// Whether a gather issued at this progress selects its survivors + /// against a threshold published partway through a probe. + fn partial_threshold(&self) -> bool { + self.full && self.partial + } } /// A probe after its sign stage. @@ -332,15 +379,20 @@ struct LazyScorer { scratch_pool: Arc, metrics: Arc, progress: Arc>, + /// Whether a lazy probe whose survivors fill the heap publishes the + /// threshold then, before it is scored whole; see + /// `LANCE_RQ_LAZY_PARTIAL_PUBLISH`. + partial_publish: bool, /// When the scan started, which its chain timings count from. started: Instant, } impl LazyScorer { /// Publish scoring progress to the gathers. The scan's first publish - /// comes from its first probe, once scored or once its rows fill the - /// heap, and is the earliest a gather waiting for the threshold or its - /// turn can be released; it is timed before any gather can see it. + /// comes from its first probe, once scored or, with partial publishes, + /// once its rows fill the heap, and is the earliest a gather waiting for + /// the threshold or its turn can be released; it is timed before any + /// gather can see it. fn publish(&self, update: impl FnOnce(&mut LazyProgress)) { self.progress.send_modify(|progress| { if progress.scored == 0 && !progress.full { @@ -535,11 +587,7 @@ impl IVFIndex { pre_filter.is_empty() && query.upper_bound.is_none(), &config, ); - let (progress_tx, progress_rx) = watch::channel(LazyProgress { - scored: 0, - full: false, - threshold: f32::INFINITY, - }); + let (progress_tx, progress_rx) = watch::channel(LazyProgress::start()); let scorer = LazyScorer { key: query.key.clone(), heap_capacity, @@ -554,6 +602,7 @@ impl IVFIndex { scratch_pool: self.scratch_pool.clone(), metrics: metrics.clone(), progress: Arc::new(progress_tx), + partial_publish: config.partial_publish, started, }; let far = match config.active_far_window(self.origin_latency) { @@ -1044,6 +1093,12 @@ impl IVFIndex { &mut survivors, &mut stage1, ); + if issue.partial_threshold() { + stats.partial_threshold_issues.incr(rank); + stats + .partial_threshold_rows + .add(rank, survivors.len() as u64); + } stats.stage1_candidates.add(rank, stage1.candidates as u64); stats .stage1_pruned_ub @@ -1142,7 +1197,8 @@ impl IVFIndex { } /// Score `batch` in probe order into `heap`, publishing progress after - /// every whole probe, matching the native scan threshold boundary. + /// every probe, and the threshold as soon as the heap fills when partial + /// publishes are on. fn score_lazy_batch( scorer: &LazyScorer, batch: Vec>, @@ -1171,13 +1227,7 @@ impl IVFIndex { .then(|| heap.peek().map(|node| node.dist.0)) .flatten(); let scored = *scored; - scorer.publish(|progress| { - progress.scored = scored; - if let Some(top) = top { - progress.full = true; - progress.threshold = top; - } - }); + scorer.publish(|progress| progress.probe_scored(scored, top)); Ok(()) }); stats.stage2_cpu_ns.add_elapsed(scoring); @@ -1283,7 +1333,12 @@ impl IVFIndex { &rows, stage1, heap, - |_threshold| {}, + |threshold| { + if scorer.partial_publish { + stats.mid_probe_full_publishes.incr(rank); + scorer.publish(|progress| progress.filled_mid_probe(threshold)); + } + }, )) })?; stats.stage2_exact.add(rank, counters.exact as u64); @@ -1472,9 +1527,9 @@ mod tests { permits: LazyFarPermits::try_new(1).unwrap(), }; let (progress_tx, mut progress) = watch::channel(LazyProgress { - scored: 0, full: true, threshold: 1.0, + ..LazyProgress::start() }); let held = far.permits.try_acquire().unwrap(); @@ -1505,6 +1560,41 @@ mod tests { assert_eq!(far.permits.in_flight(), 0); } + /// A gather parked on the threshold is released by a mid-probe publish + /// and, issued then, takes that threshold as partial. Once the probe is + /// scored whole the published threshold is its final one, no longer + /// partial, and the heap stays full whatever later probes publish. + #[tokio::test] + async fn mid_probe_threshold_stays_partial_until_its_probe_is_scored() { + let start = LazyProgress::start(); + assert!(!start.full && !start.partial_threshold()); + assert_eq!(start.threshold, f32::INFINITY); + let (progress_tx, mut progress) = watch::channel(start); + // Probe 0 is scored without filling the heap. + progress_tx.send_modify(|progress| progress.probe_scored(1, None)); + assert!(!progress.borrow().full); + + { + let released = |progress: &LazyProgress| progress.full || progress.scored >= 3; + let waiting = progress.wait_for(released); + tokio::pin!(waiting); + assert!(futures::poll!(&mut waiting).is_pending()); + progress_tx.send_modify(|progress| progress.filled_mid_probe(2.0)); + let issue = *waiting.await.unwrap(); + assert_eq!((issue.scored, issue.full, issue.threshold), (1, true, 2.0)); + assert!(issue.partial_threshold()); + } + + progress_tx.send_modify(|progress| progress.probe_scored(2, Some(1.5))); + let issue = *progress.borrow(); + assert_eq!((issue.scored, issue.full, issue.threshold), (2, true, 1.5)); + assert!(!issue.partial_threshold()); + progress_tx.send_modify(|progress| progress.probe_scored(3, Some(1.0))); + let issue = *progress.borrow(); + assert_eq!((issue.scored, issue.full, issue.threshold), (3, true, 1.0)); + assert!(!issue.partial_threshold()); + } + #[test] fn count_missing_compares_sorted_offsets() { assert_eq!(count_missing(&[], &[1, 2]), 0); diff --git a/rust/lance/src/index/vector/ivf/v2/row_layout.rs b/rust/lance/src/index/vector/ivf/v2/row_layout.rs new file mode 100644 index 00000000000..0190a5a61b5 --- /dev/null +++ b/rust/lance/src/index/vector/ivf/v2/row_layout.rs @@ -0,0 +1,2635 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Converting an IVF_RQ index between the column and plane-row layouts. +//! +//! [`rewrite_rq_row_layout`] rewrites the auxiliary file of a dataset's +//! single-segment IVF_RQ index in the other layout and commits it as a new +//! segment. Both layouts hold the same values, so nothing is re-encoded: the +//! index file, the partition offsets and lengths, the global buffers and their +//! positions, and every other schema metadata value stay as they are. Only the +//! internal columns change, and the storage metadata gains or loses its +//! `row_layout` member, as text, so every other byte of it is kept. +//! [`verify_rq_row_layout`] checks a converted index against its source. + +use std::collections::{BTreeMap, HashMap}; +use std::sync::Arc; + +use arrow::compute::concat_batches; +use arrow_array::RecordBatch; +use arrow_array::cast::AsArray; +use arrow_array::types::UInt8Type; +use arrow_schema::{Field, Schema as ArrowSchema}; +use futures::TryStreamExt; +use lance_core::cache::LanceCache; +use lance_core::deepsize::{Context, DeepSizeOf}; +use lance_core::{Error, Result}; +use lance_encoding::decoder::{DecoderPlugins, FilterExpression}; +use lance_file::reader::{FileReader, FileReaderOptions}; +use lance_file::versions as file_versions; +use lance_file::writer::{FileWriteSummary, FileWriterOptions}; +use lance_index::pb::VectorIndexDetails; +use lance_index::pb::vector_index_details::{Compression, rabit_quantization}; +use lance_index::vector::bq::builder::RabitQuantizer; +use lance_index::vector::bq::layered::{EntryColumns, SIGN_BOUNDS_PLANE, SignBounds}; +use lance_index::vector::bq::plane_rows::{PlaneRowsSpec, requests_fullzip}; +use lance_index::vector::bq::storage::{ + RABIT_METADATA_KEY, RQRowLayout, RabitQuantizationMetadata, RabitQuantizationStorage, +}; +use lance_index::vector::storage::{IvfQuantizationStorage, STORAGE_METADATA_KEY, VectorStore}; +use lance_index::{INDEX_AUXILIARY_FILE_NAME, INDEX_FILE_NAME, ivf_rq_index_version}; +use lance_io::ReadBatchParams; +use lance_io::object_store::ObjectStore; +use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; +use lance_io::utils::CachedFileSize; +use lance_table::format::{IndexFile, IndexMetadata}; +use object_store::path::Path; +use serde::Serialize; +use sha2::{Digest, Sha256}; +use uuid::Uuid; + +use super::read_partition_window_batches; +use crate::Dataset; +use crate::dataset::transaction::{Operation, Transaction}; +use crate::index::DatasetIndexExt; + +/// The storage metadata key of the row layout. +const ROW_LAYOUT_KEY: &str = "row_layout"; +/// The storage metadata member the plane-row layout adds, as JSON text. +const PLANE_ROWS_MEMBER: &str = r#""row_layout":"plane_rows""#; + +/// Partition bytes above which one partition's read of a plane-row file is a +/// large GET, reported by [`verify_rq_row_layout`]. +const LARGE_PARTITION_BYTES: u64 = 16 * 1024 * 1024; + +/// One schema metadata value a conversion rewrote. +#[derive(Debug, Clone, Serialize)] +pub struct MetadataPatch { + pub key: String, + pub source: String, + pub converted: String, +} + +/// What [`rewrite_rq_row_layout`] did. +#[derive(Debug, Clone, Serialize)] +pub struct RowLayoutRewrite { + pub index_name: String, + pub old_uuid: String, + pub new_uuid: String, + pub source_layout: String, + pub target_layout: String, + /// The index version the new segment was committed with. + pub index_version: i32, + /// The Lance file version of both auxiliary files. + pub file_version: String, + pub num_partitions: usize, + pub num_rows: u64, + pub source_aux_bytes: u64, + pub aux_bytes: u64, + pub index_file_bytes: u64, + /// User global buffers copied, in order, at their positions. + pub global_buffers: usize, + pub metadata: Vec, + /// The dataset version the conversion committed. + pub dataset_version: u64, +} + +/// Rewrite the auxiliary file of `dataset`'s IVF_RQ index `index_name` in +/// `target` and commit it as a new segment with the matching index version, +/// in place of the old one. The index must be one segment, in the other +/// layout, whose `_indices//` directory `dataset` holds. The old +/// directory is left for the caller to remove once the conversion is +/// verified, so earlier dataset versions stay readable until then. +pub async fn rewrite_rq_row_layout( + dataset: &mut Dataset, + index_name: &str, + target: RQRowLayout, +) -> Result { + rewrite_rq_row_layout_with_options(dataset, index_name, target, FileWriterOptions::default()) + .await +} + +/// [`rewrite_rq_row_layout`], writing the new auxiliary file with +/// `writer_options`, which only change where its pages end. +pub async fn rewrite_rq_row_layout_with_options( + dataset: &mut Dataset, + index_name: &str, + target: RQRowLayout, + writer_options: FileWriterOptions, +) -> Result { + rewrite_rq_row_layout_impl(dataset, index_name, target, writer_options, None).await +} + +/// The conversion, writing each partition in batches of at most +/// `batch_rows` rows when set (one batch otherwise), so tests can end pages +/// inside partitions. +async fn rewrite_rq_row_layout_impl( + dataset: &mut Dataset, + index_name: &str, + target: RQRowLayout, + writer_options: FileWriterOptions, + batch_rows: Option, +) -> Result { + let old = single_rq_segment(dataset, index_name).await?; + let store = dataset.object_store.clone(); + let old_dir = dataset.indices_dir().join(old.uuid.to_string()); + let old_aux = old_dir.clone().join(INDEX_AUXILIARY_FILE_NAME); + let old_index = old_dir.clone().join(INDEX_FILE_NAME); + for path in [&old_aux, &old_index] { + if !store.exists(path).await? { + return Err(Error::invalid_input(format!( + "converting index {index_name}: {path} is missing; convert a dataset that holds the index's directory" + ))); + } + } + let source = open_rq_storage(&store, &old_aux).await?; + let metadata = source.metadata(); + let source_layout = metadata.row_layout; + if source_layout == target { + return Err(Error::invalid_input(format!( + "index {index_name} is already in the {target} layout" + ))); + } + let reader = source.reader(); + let file_schema = ArrowSchema::from(reader.schema().as_ref()); + let spec = match target { + RQRowLayout::PlaneRows => PlaneRowsSpec::for_column_file(metadata, &file_schema)?, + RQRowLayout::Columns => PlaneRowsSpec::for_file(metadata, &file_schema)?, + }; + let patches = patch_schema_metadata(file_schema.metadata(), target)?; + let mut schema_metadata = file_schema.metadata().clone(); + for patch in &patches { + schema_metadata.insert(patch.key.clone(), patch.converted.clone()); + } + let version = reader.metadata().version(); + let fields: Vec = match target { + RQRowLayout::PlaneRows => spec.file_fields(requests_fullzip(version)), + RQRowLayout::Columns => spec + .logical_schema() + .fields() + .iter() + .map(|field| field.as_ref().clone()) + .collect(), + }; + let schema = ArrowSchema::new_with_metadata(fields, schema_metadata); + + let new_uuid = Uuid::new_v4(); + let new_dir = dataset.indices_dir().join(new_uuid.to_string()); + let new_aux = new_dir.clone().join(INDEX_AUXILIARY_FILE_NAME); + let mut writer = file_versions::create_writer( + version, + store.create(&new_aux).await?, + lance_core::datatypes::Schema::try_from(&schema)?, + writer_options, + )?; + // The IVF buffer and the rotation matrix are referenced by position, so + // they are copied in order to the same positions. + let global_buffers = reader.metadata().file_buffers.len().saturating_sub(1); + for position in 1..=global_buffers as u32 { + let bytes = reader.read_global_buffer(position).await?; + let copied = writer.add_global_buffer(bytes).await?; + if copied != position { + return Err(Error::internal(format!( + "converting index {index_name}: global buffer {position} was written at {copied}" + ))); + } + } + let ivf = source.ivf(); + let file_schema_ref = Arc::new(file_schema.clone()); + for partition in 0..ivf.num_partitions() { + let range = ivf.row_range(partition); + if range.is_empty() { + continue; + } + let batches = reader + .read_stream( + ReadBatchParams::Range(range), + u32::MAX, + 1, + FilterExpression::no_filter(), + ) + .await? + .try_collect::>() + .await?; + let batch = match target { + RQRowLayout::PlaneRows => spec.pack(&concat_batches(&file_schema_ref, &batches)?)?, + RQRowLayout::Columns => spec.unpack(&batches, spec.logical_schema())?, + }; + let step = batch_rows.unwrap_or(batch.num_rows()).max(1); + for start in (0..batch.num_rows()).step_by(step) { + let rows = step.min(batch.num_rows() - start); + writer.write_batch(&batch.slice(start, rows)).await?; + } + } + let FileWriteSummary { + num_rows, + size_bytes: aux_bytes, + } = writer.finish().await?; + if num_rows != reader.num_rows() { + return Err(Error::internal(format!( + "converting index {index_name}: wrote {num_rows} of the source's {} rows", + reader.num_rows() + ))); + } + let new_index = new_dir.join(INDEX_FILE_NAME); + store.copy(&old_index, &new_index).await?; + let index_file_bytes = store.size(&new_index).await?; + let source_aux_bytes = store.size(&old_aux).await?; + + let index_version = ivf_rq_index_version(target); + let new = IndexMetadata { + uuid: new_uuid, + index_details: Some(Arc::new(details_with_row_layout(&old, target)?)), + index_version, + files: Some(vec![ + IndexFile { + path: INDEX_AUXILIARY_FILE_NAME.to_string(), + size_bytes: aux_bytes, + }, + IndexFile { + path: INDEX_FILE_NAME.to_string(), + size_bytes: index_file_bytes, + }, + ]), + ..old.clone() + }; + let transaction = Transaction::new( + dataset.manifest.version, + Operation::CreateIndex { + new_indices: vec![new], + removed_indices: vec![old.clone()], + }, + None, + ); + dataset + .apply_commit(transaction, &Default::default(), &Default::default()) + .await?; + Ok(RowLayoutRewrite { + index_name: index_name.to_string(), + old_uuid: old.uuid.to_string(), + new_uuid: new_uuid.to_string(), + source_layout: source_layout.to_string(), + target_layout: target.to_string(), + index_version, + file_version: version.to_manifest_string().to_string(), + num_partitions: ivf.num_partitions(), + num_rows, + source_aux_bytes, + aux_bytes, + index_file_bytes, + global_buffers, + metadata: patches, + dataset_version: dataset.manifest.version, + }) +} + +/// The one segment of IVF_RQ index `index_name`. +async fn single_rq_segment(dataset: &Dataset, index_name: &str) -> Result { + let mut segments = dataset.load_indices_by_name(index_name).await?; + if segments.len() != 1 { + return Err(Error::invalid_input(format!( + "converting index {index_name}: it has {} segments; only a single-segment index converts", + segments.len() + ))); + } + let segment = segments.remove(0); + rq_details(&segment)?; + Ok(segment) +} + +/// The segment's vector details, which must describe IVF_RQ. +fn rq_details(segment: &IndexMetadata) -> Result { + let details = segment + .index_details + .as_ref() + .filter(|details| !details.value.is_empty()) + .ok_or_else(|| { + Error::invalid_input(format!( + "converting index {}: segment {} has no vector index details", + segment.name, segment.uuid + )) + })? + .to_msg::() + .map_err(|error| { + Error::invalid_input(format!( + "converting index {}: segment {} is not a vector index: {error}", + segment.name, segment.uuid + )) + })?; + if !matches!(details.compression, Some(Compression::Rq(_))) + || details.hnsw_index_config.is_some() + { + return Err(Error::invalid_input(format!( + "converting index {}: segment {} is not an IVF_RQ index", + segment.name, segment.uuid + ))); + } + Ok(details) +} + +/// `segment`'s details with the row layout `target`. +fn details_with_row_layout( + segment: &IndexMetadata, + target: RQRowLayout, +) -> Result { + let mut details = rq_details(segment)?; + if let Some(Compression::Rq(rq)) = details.compression.as_mut() { + rq.set_row_layout(match target { + RQRowLayout::Columns => rabit_quantization::RowLayout::Columns, + RQRowLayout::PlaneRows => rabit_quantization::RowLayout::PlaneRows, + }); + } + prost_types::Any::from_msg(&details) + .map_err(|error| Error::internal(format!("encoding vector index details: {error}"))) +} + +/// The IVF_RQ storage of the auxiliary file at `path`. +async fn open_rq_storage( + store: &Arc, + path: &Path, +) -> Result> { + let scheduler = ScanScheduler::new(store.clone(), SchedulerConfig::max_bandwidth(store)); + let reader = FileReader::try_open( + scheduler + .open_file(path, &CachedFileSize::unknown()) + .await?, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await?; + IvfQuantizationStorage::try_new(reader, None).await +} + +/// The schema metadata values a conversion to `target` rewrites: the +/// storage metadata, and the quantizer metadata files written by the +/// distributed merger also carry, each with the plane-row member added as +/// its last member or removed, as text. +fn patch_schema_metadata( + metadata: &HashMap, + target: RQRowLayout, +) -> Result> { + let insert = target == RQRowLayout::PlaneRows; + let storage = metadata.get(STORAGE_METADATA_KEY).ok_or_else(|| { + Error::invalid_input(format!( + "the auxiliary file has no {STORAGE_METADATA_KEY} metadata" + )) + })?; + // The storage metadata is a JSON list of one JSON string, so the member + // goes in escaped, before the closing brace of that string's object. + let escaped = serde_json::to_string(PLANE_ROWS_MEMBER)?; + let escaped = &escaped[1..escaped.len() - 1]; + let converted = patch_object_text(storage, escaped, insert)?; + check_storage_metadata(storage, &converted, target)?; + let mut patches = vec![MetadataPatch { + key: STORAGE_METADATA_KEY.to_string(), + source: storage.clone(), + converted, + }]; + if let Some(rabit) = metadata.get(RABIT_METADATA_KEY) { + let converted = patch_object_text(rabit, PLANE_ROWS_MEMBER, insert)?; + check_metadata_object(rabit, &converted, target)?; + patches.push(MetadataPatch { + key: RABIT_METADATA_KEY.to_string(), + source: rabit.clone(), + converted, + }); + } + Ok(patches) +} + +/// `text`, which ends with a JSON object's closing brace and possibly +/// closing list and string delimiters, with `member` added as the object's +/// last member or, when `!insert`, removed with its separator; every other +/// byte is kept. +fn patch_object_text(text: &str, member: &str, insert: bool) -> Result { + if insert { + let (open, close) = text + .find('{') + .zip(text.rfind('}')) + .ok_or_else(|| Error::invalid_input(format!("{text:?} holds no JSON object")))?; + let separator = if text[open + 1..close].trim().is_empty() { + "" + } else { + "," + }; + return Ok(format!( + "{}{separator}{member}{}", + &text[..close], + &text[close..] + )); + } + for candidate in [ + format!(",{member}"), + format!("{member},"), + member.to_string(), + ] { + if let Some(start) = text.find(&candidate) { + return Ok(format!( + "{}{}", + &text[..start], + &text[start + candidate.len()..] + )); + } + } + Err(Error::invalid_input(format!( + "{text:?} has no {member} member to remove" + ))) +} + +/// Check that `converted`, the storage metadata after a conversion to +/// `target`, is `source`'s single entry with the row layout member alone +/// added or removed. +fn check_storage_metadata(source: &str, converted: &str, target: RQRowLayout) -> Result<()> { + let entry = |text: &str| -> Result { + let mut entries: Vec = serde_json::from_str(text)?; + if entries.len() != 1 { + return Err(Error::invalid_input(format!( + "the storage metadata holds {} entries, expected one", + entries.len() + ))); + } + Ok(entries.remove(0)) + }; + check_metadata_object(&entry(source)?, &entry(converted)?, target) +} + +/// Check that `converted`, an IVF_RQ metadata object after a conversion to +/// `target`, is `source` with the row layout member alone added or removed, +/// and reads back as `target`. +fn check_metadata_object(source: &str, converted: &str, target: RQRowLayout) -> Result<()> { + let object = |text: &str| -> Result> { + match serde_json::from_str(text)? { + serde_json::Value::Object(object) => Ok(object), + other => Err(Error::invalid_input(format!( + "the IVF_RQ metadata {other} is not a JSON object" + ))), + } + }; + let mut expected = object(source)?; + match target { + RQRowLayout::PlaneRows => { + if expected.contains_key(ROW_LAYOUT_KEY) { + return Err(Error::invalid_input(format!( + "the IVF_RQ metadata {source} already declares a row layout" + ))); + } + expected.insert( + ROW_LAYOUT_KEY.to_string(), + serde_json::Value::String(RQRowLayout::PlaneRows.to_string()), + ); + } + RQRowLayout::Columns => { + expected.remove(ROW_LAYOUT_KEY); + } + } + let read: RabitQuantizationMetadata = serde_json::from_str(converted)?; + if object(converted)? != expected || read.row_layout != target { + return Err(Error::internal(format!( + "rewriting the IVF_RQ metadata {source} for {target} gave {converted}" + ))); + } + Ok(()) +} + +/// How a plane-row file's pages fall on its partitions. +#[derive(Debug, Clone, Default, Serialize)] +pub struct PageStats { + /// Pages of the packed columns, summed over them. + pub pages: usize, + /// Partition reads of a packed column whose rows span more than one of + /// its pages, summed over the packed columns. + pub straddling_partitions: usize, + /// The largest number of bytes one partition takes, over all its + /// packed columns. + pub max_partition_bytes: u64, + /// Partitions taking more than 16 MiB. + pub partitions_over_16mib: usize, +} + +/// Shapes of reads and cache entries whose memory [`verify_rq_row_layout`] +/// compares, as [`RowLayoutVerification::deep_sizes`] names them. +pub mod deep_size_shapes { + /// Every row of a partition, which partition entries are built from. + pub const PARTITION_BATCH: &str = "partition_batch"; + /// Every row of a plane of a layered index, which a plane entry holding + /// every column holds. + pub const PLANE: &str = "plane"; + /// Every third row of a plane of a layered index: a sparse read, which + /// no cache entry holds. + pub const PLANE_SPARSE: &str = "plane_sparse"; + /// The storage a partition entry holds. + pub const PARTITION_STORAGE: &str = "partition_storage"; + /// The storage a partition entry holds when a prewarm built it from a + /// read of several partitions. + pub const PREWARM_PARTITION_STORAGE: &str = "prewarm_partition_storage"; + /// A code-only entry: a native partition's or a layered plane's code + /// (and bounds) columns. + pub const CODE_ONLY_ENTRY: &str = "code_only_entry"; + /// A native partition's code-only entry as a prewarm cuts it from a read + /// of several partitions. + pub const PREWARM_CODE_ONLY_ENTRY: &str = "prewarm_code_only_entry"; +} + +/// The memory one shape of read or cache entry takes on both sides. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)] +pub struct DeepSizes { + /// Reads or entries compared. + pub count: usize, + /// Those whose two sides take different memory. + pub mismatches: usize, + pub source_bytes: u64, + pub converted_bytes: u64, +} + +impl DeepSizes { + fn record(&mut self, source: u64, converted: u64) { + self.count += 1; + self.mismatches += usize::from(source != converted); + self.source_bytes += source; + self.converted_bytes += converted; + } +} + +/// What [`verify_rq_row_layout`] found. +#[derive(Debug, Clone, Default, Serialize)] +pub struct RowLayoutVerification { + pub source_uuid: String, + pub converted_uuid: String, + pub source_layout: String, + pub converted_layout: String, + /// The Lance file version of each side's auxiliary file. + pub source_file_version: String, + pub converted_file_version: String, + pub partitions: usize, + pub rows: u64, + /// Reads compared: every partition whole, and for a layered index every + /// plane, whole and at every third row. + pub reads: usize, + /// Cache entries compared: every partition's storage, read alone and + /// cut from a prewarm's read, and its code-only entries. + pub entries: usize, + /// Reads and entries whose batches differ in a field or a value. + pub batch_mismatches: usize, + /// Whole reads and cache entries whose two sides take different memory, + /// so would be charged differently in a cache, and their memory. + pub deep_size_mismatches: usize, + pub source_deep_bytes: u64, + pub converted_deep_bytes: u64, + /// Sparse reads whose two sides take different memory, and their + /// memory. No cache entry holds a sparse read, so this is reported + /// apart: a column-layout read of selected rows of a Lance 2.0 file can + /// keep more buffer than its values. + pub sparse_deep_size_mismatches: usize, + pub sparse_source_deep_bytes: u64, + pub sparse_converted_deep_bytes: u64, + /// The memory of each shape of read and entry ([`deep_size_shapes`]). + pub deep_sizes: BTreeMap, + /// SHA-256 over every partition's values, rows packed as the plane-row + /// layout packs them, for each side. + pub source_digest: String, + pub converted_digest: String, + /// Whether the column layout's fields read from both sides are the same. + pub schema_matches: bool, + /// The storage metadata of both sides, and whether the converted one is + /// the source's with the row layout alone added or removed, as text. + pub source_storage_metadata: String, + pub converted_storage_metadata: String, + pub metadata_patch_ok: bool, + /// Whether every other schema metadata value is byte for byte the same. + pub other_metadata_equal: bool, + pub rotate_mat_position: Option, + pub rotate_mat_position_equal: bool, + /// SHA-256 of each user global buffer of the source, and whether the + /// converted file's are the same, in the same positions. + pub global_buffer_digests: Vec, + pub global_buffers_equal: bool, + /// Whether both index files hold the same bytes. + pub index_file_equal: bool, + /// Page stats of whichever side stores plane rows. + pub pages: PageStats, +} + +impl RowLayoutVerification { + /// Whether the converted index holds the source's values, files and + /// metadata, but for the layout, and every whole read and cache entry of + /// it takes the memory the source's does, so a cache charges both the + /// same. Sparse reads, which no entry holds, are reported apart. + pub fn is_ok(&self) -> bool { + self.batch_mismatches == 0 + && self.deep_size_mismatches == 0 + && self.source_digest == self.converted_digest + && self.schema_matches + && self.metadata_patch_ok + && self.other_metadata_equal + && self.rotate_mat_position_equal + && self.global_buffers_equal + && self.index_file_equal + } + + /// Count a read or entry of `shape` whose values are `same` and whose + /// sides take `source` and `converted` bytes. + fn record(&mut self, shape: &str, same: bool, source: u64, converted: u64) { + self.batch_mismatches += usize::from(!same); + self.deep_sizes + .entry(shape.to_string()) + .or_default() + .record(source, converted); + let mismatch = usize::from(source != converted); + if shape == deep_size_shapes::PLANE_SPARSE { + self.sparse_deep_size_mismatches += mismatch; + self.sparse_source_deep_bytes += source; + self.sparse_converted_deep_bytes += converted; + } else { + self.deep_size_mismatches += mismatch; + self.source_deep_bytes += source; + self.converted_deep_bytes += converted; + } + } + + /// Count read `a` of the source against read `b` of the converted index. + fn compare(&mut self, shape: &str, a: &RecordBatch, b: &RecordBatch) { + self.reads += usize::from(!is_entry_shape(shape)); + self.entries += usize::from(is_entry_shape(shape)); + let same = a.schema().fields() == b.schema().fields() && a.columns() == b.columns(); + self.record(shape, same, deep_bytes(a), deep_bytes(b)); + } + + /// Count the storage of a partition entry of each side. + fn compare_storage( + &mut self, + shape: &str, + a: &RabitQuantizationStorage, + b: &RabitQuantizationStorage, + ) { + self.entries += 1; + let same = a.len() == b.len() && a.row_ids().eq(b.row_ids()); + self.record( + shape, + same, + a.deep_size_of() as u64, + b.deep_size_of() as u64, + ); + } +} + +/// `storage`'s partitions as a prewarm cuts them from one read of every +/// partition, which it builds their entries from. +async fn prewarm_window( + storage: &IvfQuantizationStorage, +) -> Result>> { + let schema = Arc::new(ArrowSchema::from(storage.reader().schema().as_ref())); + read_partition_window_batches( + storage.reader(), + None, + &schema, + storage.ivf(), + 0..storage.num_partitions(), + None, + ) + .await +} + +/// Whether `shape` is a cache entry rather than a read. +fn is_entry_shape(shape: &str) -> bool { + [ + deep_size_shapes::CODE_ONLY_ENTRY, + deep_size_shapes::PREWARM_CODE_ONLY_ENTRY, + ] + .contains(&shape) +} + +/// The memory `batch`'s arrays take, as a cache charges it. +fn deep_bytes(batch: &RecordBatch) -> u64 { + batch.deep_size_of_children(&mut Context::new()) as u64 +} + +/// Check that index `index_name` of `converted`, a single segment converted +/// from `source`'s, holds the same values: every partition and, for a +/// layered index, every plane under both bounds placements, whole and at +/// every third row; the same global buffers in the same positions, the same +/// index file, and the same schema metadata but for the row layout member. +pub async fn verify_rq_row_layout( + source: &Dataset, + converted: &Dataset, + index_name: &str, +) -> Result { + let source_segment = single_rq_segment(source, index_name).await?; + let converted_segment = single_rq_segment(converted, index_name).await?; + let segment_file = |dataset: &Dataset, segment: &IndexMetadata, file: &str| { + dataset + .indices_dir() + .join(segment.uuid.to_string()) + .join(file) + }; + let source_aux = segment_file(source, &source_segment, INDEX_AUXILIARY_FILE_NAME); + let converted_aux = segment_file(converted, &converted_segment, INDEX_AUXILIARY_FILE_NAME); + let open = |dataset: &Dataset, path: &Path, sign_bounds: SignBounds| { + let store = dataset.object_store.clone(); + let path = path.clone(); + async move { + Ok::<_, Error>( + open_rq_storage(&store, &path) + .await? + .with_sign_bounds(sign_bounds), + ) + } + }; + let lazy = [ + open(source, &source_aux, SignBounds::Lazy).await?, + open(converted, &converted_aux, SignBounds::Lazy).await?, + ]; + let eager = [ + open(source, &source_aux, SignBounds::Eager).await?, + open(converted, &converted_aux, SignBounds::Eager).await?, + ]; + let [src, conv] = &lazy; + let file_version = |storage: &IvfQuantizationStorage| -> String { + storage + .reader() + .metadata() + .version() + .to_manifest_string() + .to_string() + }; + let mut report = RowLayoutVerification { + source_uuid: source_segment.uuid.to_string(), + converted_uuid: converted_segment.uuid.to_string(), + source_layout: src.row_layout().to_string(), + converted_layout: conv.row_layout().to_string(), + source_file_version: file_version(src), + converted_file_version: file_version(conv), + partitions: src.num_partitions(), + rows: src.num_rows(), + schema_matches: src.logical_schema()?.fields() == conv.logical_schema()?.fields(), + ..Default::default() + }; + if src.row_layout() == conv.row_layout() { + return Err(Error::invalid_input(format!( + "index {index_name} is in the {} layout on both sides", + src.row_layout() + ))); + } + if src.ivf().lengths != conv.ivf().lengths || src.num_rows() != conv.num_rows() { + return Err(Error::invalid_input(format!( + "index {index_name}: the converted partitions differ from the source's" + ))); + } + + // The packed rows of both sides digest the values the same way. + let plane_rows = |storage: &IvfQuantizationStorage| -> Result { + let schema = ArrowSchema::from(storage.reader().schema().as_ref()); + match storage.row_layout() { + RQRowLayout::PlaneRows => PlaneRowsSpec::for_file(storage.metadata(), &schema), + RQRowLayout::Columns => PlaneRowsSpec::for_column_file(storage.metadata(), &schema), + } + }; + let specs = [plane_rows(src)?, plane_rows(conv)?]; + let mut digests = [Sha256::new(), Sha256::new()]; + let layered = src.is_layered_rq(); + let [mut src_prewarm, mut conv_prewarm] = + [prewarm_window(src).await?, prewarm_window(conv).await?].map(Vec::into_iter); + // A native partition's code-only entry as a prewarm of the column side + // cuts it; the plane-row layout caches none, so its twin is the read of + // the same columns. + let column_side = if src.row_layout() == RQRowLayout::Columns { + src + } else { + conv + }; + let mut prewarm_codes = if layered { + Vec::new() + } else { + let projection = column_side.partition_codes_projection()?; + let schema = Arc::new(ArrowSchema::from(projection.schema.as_ref())); + read_partition_window_batches( + column_side.reader(), + Some(&projection), + &schema, + column_side.ivf(), + 0..column_side.num_partitions(), + None, + ) + .await? + } + .into_iter(); + for partition in 0..report.partitions { + let batches = [ + src.read_partition_batch(partition, None).await?, + conv.read_partition_batch(partition, None).await?, + ]; + report.compare(deep_size_shapes::PARTITION_BATCH, &batches[0], &batches[1]); + for ((digest, spec), batch) in digests.iter_mut().zip(&specs).zip(&batches) { + let packed = spec.pack(batch)?; + for column in spec.packed_columns() { + let rows = packed + .column_by_name(column.name()) + .and_then(|rows| rows.as_fixed_size_list_opt()) + .and_then(|rows| rows.values().as_primitive_opt::()) + .ok_or_else(|| { + Error::internal(format!("packed column {} is missing", column.name())) + })?; + digest.update(rows.values()); + } + } + report.compare_storage( + deep_size_shapes::PARTITION_STORAGE, + &src.load_partition(partition, None).await?, + &conv.load_partition(partition, None).await?, + ); + let prewarmed = [src_prewarm.next(), conv_prewarm.next()].map(|batches| { + batches.ok_or_else(|| Error::internal("a prewarm read missed a partition")) + }); + let [src_prewarmed, conv_prewarmed] = prewarmed; + report.compare_storage( + deep_size_shapes::PREWARM_PARTITION_STORAGE, + &src.materialize_partition_for_prewarm(src_prewarmed?) + .await?, + &conv + .materialize_partition_for_prewarm(conv_prewarmed?) + .await?, + ); + if !layered { + let codes = [ + src.read_partition_codes(partition, None).await?, + conv.read_partition_codes(partition, None).await?, + ]; + report.compare(deep_size_shapes::CODE_ONLY_ENTRY, &codes[0].0, &codes[1].0); + let prewarmed = column_side + .partition_codes_from_batches( + prewarm_codes + .next() + .ok_or_else(|| Error::internal("a prewarm read missed a partition"))?, + )? + .0; + let [a, b] = if column_side.row_layout() == src.row_layout() { + [&prewarmed, &codes[1].0] + } else { + [&codes[0].0, &prewarmed] + }; + report.compare(deep_size_shapes::PREWARM_CODE_ONLY_ENTRY, a, b); + continue; + } + let rows = src.partition_size(partition); + let every_third: Vec = (0..rows as u32).step_by(3).collect(); + for (storages, planes) in [ + (&lazy, &[0, 1, 2, SIGN_BOUNDS_PLANE][..]), + (&eager, &[0][..]), + ] { + let [a, b] = storages; + for &plane in planes { + for selected in [None, Some(every_third.clone())] { + let shape = if selected.is_some() { + deep_size_shapes::PLANE_SPARSE + } else { + deep_size_shapes::PLANE + }; + report.compare( + shape, + &a.read_plane(partition, plane, selected.clone(), None) + .await?, + &b.read_plane(partition, plane, selected.clone(), None) + .await?, + ); + } + report.compare( + deep_size_shapes::CODE_ONLY_ENTRY, + &a.read_plane_entry_with(partition, plane, EntryColumns::Codes, None) + .await?, + &b.read_plane_entry_with(partition, plane, EntryColumns::Codes, None) + .await?, + ); + } + } + } + let [source_digest, converted_digest] = digests.map(|digest| hex(&digest.finalize())); + report.source_digest = source_digest; + report.converted_digest = converted_digest; + + // Metadata: the storage metadata differs by the row layout member alone, + // and every other value is the same. + let metadata = |storage: &IvfQuantizationStorage| { + storage.reader().schema().metadata.clone() + }; + let [source_metadata, converted_metadata] = [metadata(src), metadata(conv)]; + let patched = patch_schema_metadata(&source_metadata, conv.row_layout()); + report.source_storage_metadata = source_metadata + .get(STORAGE_METADATA_KEY) + .cloned() + .unwrap_or_default(); + report.converted_storage_metadata = converted_metadata + .get(STORAGE_METADATA_KEY) + .cloned() + .unwrap_or_default(); + report.metadata_patch_ok = patched.as_ref().is_ok_and(|patches| { + patches + .iter() + .all(|patch| converted_metadata.get(&patch.key) == Some(&patch.converted)) + }); + let patched_keys = [STORAGE_METADATA_KEY, RABIT_METADATA_KEY]; + let others = |metadata: &HashMap| { + metadata + .iter() + .filter(|(key, _)| !patched_keys.contains(&key.as_str())) + .map(|(key, value)| (key.clone(), value.clone())) + .collect::>() + }; + report.other_metadata_equal = others(&source_metadata) == others(&converted_metadata) + && source_metadata.contains_key(RABIT_METADATA_KEY) + == converted_metadata.contains_key(RABIT_METADATA_KEY); + report.rotate_mat_position = src.metadata().rotate_mat_position; + report.rotate_mat_position_equal = + src.metadata().rotate_mat_position == conv.metadata().rotate_mat_position; + + // Global buffers, in their positions, and the index file. + let buffers = |storage: &IvfQuantizationStorage| { + storage + .reader() + .metadata() + .file_buffers + .len() + .saturating_sub(1) + }; + report.global_buffers_equal = buffers(src) == buffers(conv); + for position in 1..=buffers(src) as u32 { + let source_bytes = src.reader().read_global_buffer(position).await?; + report + .global_buffer_digests + .push(hex(&Sha256::digest(&source_bytes))); + if report.global_buffers_equal { + let converted_bytes = conv.reader().read_global_buffer(position).await?; + report.global_buffers_equal &= source_bytes == converted_bytes; + } + } + let index_bytes = |dataset: &Dataset, segment: &IndexMetadata| { + let store = dataset.object_store.clone(); + let path = segment_file(dataset, segment, INDEX_FILE_NAME); + async move { store.read_one_all(&path).await } + }; + report.index_file_equal = index_bytes(source, &source_segment).await? + == index_bytes(converted, &converted_segment).await?; + + let plane_row_side = if src.row_layout() == RQRowLayout::PlaneRows { + src + } else { + conv + }; + report.pages = page_stats(plane_row_side)?; + Ok(report) +} + +/// How the pages of `storage`'s plane-row file fall on its partitions. +fn page_stats(storage: &IvfQuantizationStorage) -> Result { + let spec = storage + .plane_rows()? + .ok_or_else(|| Error::internal("page stats need a plane-row file"))?; + let reader = storage.reader(); + let mut stats = PageStats::default(); + let ivf = storage.ivf(); + let row_bytes: usize = spec + .packed_columns() + .iter() + .map(|column| column.stride()) + .sum(); + for partition in 0..ivf.num_partitions() { + let bytes = (ivf.partition_size(partition) * row_bytes) as u64; + stats.max_partition_bytes = stats.max_partition_bytes.max(bytes); + stats.partitions_over_16mib += usize::from(bytes > LARGE_PARTITION_BYTES); + } + for column in spec.packed_columns() { + let projection = lance_file::versions::reader_projection_from_column_names( + reader.metadata().version(), + reader.schema(), + &[column.name()], + )?; + let index = projection.column_indices[0] as usize; + let info = + reader.metadata().column_infos.get(index).ok_or_else(|| { + Error::internal(format!("column {} has no page info", column.name())) + })?; + stats.pages += info.page_infos.len(); + // The first row of every page after the first. + let mut starts = Vec::with_capacity(info.page_infos.len()); + let mut next = 0u64; + for page in info.page_infos.iter() { + next += page.num_rows; + starts.push(next); + } + starts.pop(); + for partition in 0..ivf.num_partitions() { + let range = ivf.row_range(partition); + if range.is_empty() { + continue; + } + let (start, end) = (range.start as u64, range.end as u64); + stats.straddling_partitions += usize::from( + starts + .iter() + .any(|&page_start| page_start > start && page_start < end), + ); + } + } + Ok(stats) +} + +fn hex(bytes: &[u8]) -> String { + bytes.iter().map(|byte| format!("{byte:02x}")).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow_array::types::{Float32Type, UInt64Type}; + use arrow_array::{ + Array, ArrayRef, FixedSizeListArray, Float32Array, RecordBatchIterator, UInt64Array, + }; + use arrow_schema::DataType; + use bytes::Bytes; + use lance_arrow::{FixedSizeListArrayExt, RecordBatchExt}; + use lance_core::ROW_ID; + use lance_core::cache::CacheCodec; + use lance_core::utils::tempfile::TempStrDir; + use lance_encoding::decoder::PageEncoding; + use lance_encoding::format::pb21::page_layout::Layout; + use lance_file::version::LanceFileVersion; + use lance_index::IndexType; + use lance_index::optimize::OptimizeOptions; + use lance_index::vector::DIST_COL; + use lance_index::vector::bq::layered::{PlaneBatch, RQPrecision, SignBounds}; + use lance_index::vector::bq::storage::{RABIT_CODE_COLUMN, unpack_codes}; + use lance_index::vector::bq::{RQBuildParams, RQRotationType}; + use lance_index::vector::ivf::IvfBuildParams; + use lance_index::vector::storage::VectorStore; + use lance_io::scheduler::IoStats; + use lance_linalg::distance::DistanceType; + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + use rstest::rstest; + + use crate::dataset::optimize::{CompactionOptions, compact_files}; + use crate::dataset::{WriteMode, WriteParams}; + use crate::index::vector::VectorIndexParams; + use crate::index::vector::ivf::v2::IVFIndex; + use crate::index::{DatasetIndexInternalExt, vector::VectorIndex}; + use lance_index::metrics::NoOpMetricsCollector; + use lance_index::vector::flat::index::FlatIndex; + + const INDEX_NAME: &str = "vector_idx"; + const DIM: usize = 64; + /// Partitions of the test index; the data clusters around all but the + /// last centroid, which is left with an empty partition. + const PARTITIONS: usize = 6; + const ROWS: usize = 3_000; + /// A page budget well below a partition's bytes, and batches of a number + /// of rows that divides no partition, so pages end inside partitions. + const SMALL_PAGE_BYTES: u64 = 4096; + const SMALL_BATCH_ROWS: usize = 97; + + type IvfRq = IVFIndex; + + fn centroids() -> Vec { + let mut rng = StdRng::seed_from_u64(0x9470); + (0..PARTITIONS * DIM) + .map(|_| rng.random_range(-1.0f32..1.0)) + .collect() + } + + /// `rows` vectors clustered around the first `PARTITIONS - 1` centroids, + /// with ids from `first_id`. + fn test_batch(rows: usize, first_id: u64, seed: u64) -> RecordBatch { + let centroids = centroids(); + let mut rng = StdRng::seed_from_u64(seed); + let mut values = Vec::with_capacity(rows * DIM); + for _ in 0..rows { + let cluster = rng.random_range(0..PARTITIONS - 1); + for dim in 0..DIM { + values.push(centroids[cluster * DIM + dim] + rng.random_range(-0.05f32..0.05)); + } + } + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), DIM as i32) + .unwrap(); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new("vector", vectors.data_type().clone(), true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(UInt64Array::from_iter_values( + first_id..first_id + rows as u64, + )), + Arc::new(vectors), + ], + ) + .unwrap() + } + + async fn write_dataset( + uri: &str, + version: LanceFileVersion, + max_rows_per_file: usize, + ) -> Dataset { + let batch = test_batch(ROWS, 0, 1); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let params = WriteParams { + data_storage_version: Some(version), + max_rows_per_file, + mode: WriteMode::Create, + ..Default::default() + }; + Dataset::write(reader, uri, Some(params)).await.unwrap() + } + + fn ivf_params() -> IvfBuildParams { + let centroids = + FixedSizeListArray::try_new_from_values(Float32Array::from(centroids()), DIM as i32) + .unwrap(); + IvfBuildParams::try_with_centroids(PARTITIONS, Arc::new(centroids)).unwrap() + } + + fn rq_params(layered: bool, rotation_type: RQRotationType) -> RQBuildParams { + RQBuildParams::with_rotation_type(7, rotation_type).with_layered(layered) + } + + async fn create_index(dataset: &mut Dataset, rq: RQBuildParams) { + let params = VectorIndexParams::with_ivf_rq_params(DistanceType::L2, ivf_params(), rq); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + } + + /// A column-layout dataset at `uri` with a native or layered index. + async fn column_dataset( + uri: &str, + layered: bool, + rotation_type: RQRotationType, + version: LanceFileVersion, + ) -> Dataset { + let mut dataset = write_dataset(uri, version, ROWS).await; + create_index(&mut dataset, rq_params(layered, rotation_type)).await; + dataset + } + + /// Every file under `dir` with its size, sorted. + fn list_files(dir: &std::path::Path) -> Vec<(std::path::PathBuf, u64)> { + let mut files = Vec::new(); + for entry in std::fs::read_dir(dir).unwrap() { + let entry = entry.unwrap(); + if entry.file_type().unwrap().is_dir() { + files.extend(list_files(&entry.path())); + } else { + files.push((entry.path(), entry.metadata().unwrap().len())); + } + } + files.sort(); + files + } + + fn copy_dir(from: &std::path::Path, to: &std::path::Path) { + std::fs::create_dir_all(to).unwrap(); + for entry in std::fs::read_dir(from).unwrap() { + let entry = entry.unwrap(); + let target = to.join(entry.file_name()); + if entry.file_type().unwrap().is_dir() { + copy_dir(&entry.path(), &target); + } else { + std::fs::copy(entry.path(), target).unwrap(); + } + } + } + + /// A full copy of the dataset at `from`, index directories included, at + /// `to`, converted in place to plane rows. + async fn converted_copy( + from: &str, + to: &str, + options: FileWriterOptions, + ) -> (Dataset, RowLayoutRewrite) { + converted_copy_in_batches(from, to, options, None).await + } + + /// [`converted_copy`], writing at most `batch_rows` rows per batch. + async fn converted_copy_in_batches( + from: &str, + to: &str, + options: FileWriterOptions, + batch_rows: Option, + ) -> (Dataset, RowLayoutRewrite) { + copy_dir(std::path::Path::new(from), std::path::Path::new(to)); + let mut dataset = Dataset::open(to).await.unwrap(); + let rewrite = rewrite_rq_row_layout_impl( + &mut dataset, + INDEX_NAME, + RQRowLayout::PlaneRows, + options, + batch_rows, + ) + .await + .unwrap(); + (dataset, rewrite) + } + + async fn index_storage(dataset: &Dataset) -> IvfQuantizationStorage { + let segment = single_rq_segment(dataset, INDEX_NAME).await.unwrap(); + let path = dataset + .indices_dir() + .join(segment.uuid.to_string()) + .join(INDEX_AUXILIARY_FILE_NAME); + open_rq_storage(&dataset.object_store, &path).await.unwrap() + } + + async fn open_index(dataset: &Dataset) -> Arc { + let segment = single_rq_segment(dataset, INDEX_NAME).await.unwrap(); + dataset + .open_vector_index("vector", &segment.uuid, &NoOpMetricsCollector) + .await + .unwrap() + } + + fn query(seed: u64) -> ArrayRef { + test_batch(1, 0, seed)["vector"] + .as_fixed_size_list() + .value(0) + } + + /// Row ids and distance bits of a query, in result order. + async fn search( + dataset: &Dataset, + query: &dyn Array, + k: usize, + nprobes: usize, + precision: RQPrecision, + ) -> (Vec, Vec) { + let result = dataset + .scan() + .nearest("vector", query, k) + .unwrap() + .nprobes(nprobes) + .rq_precision(precision) + .with_row_id() + .try_into_batch() + .await + .unwrap(); + ( + result[ROW_ID] + .as_primitive::() + .values() + .to_vec(), + result[DIST_COL] + .as_primitive::() + .values() + .iter() + .map(|dist| dist.to_bits()) + .collect(), + ) + } + + async fn assert_same_results(a: &Dataset, b: &Dataset, layered: bool) { + let precisions: &[RQPrecision] = if layered { + &[RQPrecision::Full, RQPrecision::High, RQPrecision::Sign] + } else { + &[RQPrecision::Full] + }; + for seed in [11, 12, 13] { + let query = query(seed); + for (k, nprobes) in [(1, 1), (10, 2), (100, PARTITIONS), (ROWS + 1, PARTITIONS)] { + for &precision in precisions { + assert_eq!( + search(a, query.as_ref(), k, nprobes, precision).await, + search(b, query.as_ref(), k, nprobes, precision).await, + "seed={seed} k={k} nprobes={nprobes} precision={precision:?}" + ); + } + } + } + } + + fn assert_same_batch(a: &RecordBatch, b: &RecordBatch, context: &str) { + assert_eq!(a.schema().fields(), b.schema().fields(), "{context}"); + assert_eq!(a.columns(), b.columns(), "{context}"); + } + + /// The conversion runs in place on a copy that still holds the old index + /// directory: it commits a new segment at version 3 whose files it + /// records, every query answers as before, earlier versions stay readable + /// until the caller removes the old directory, and the source is not + /// written. Only `row_layout` is added to the metadata, as text. + #[rstest] + #[case::native_fast(false, RQRotationType::Fast)] + #[case::native_matrix(false, RQRotationType::Matrix)] + #[case::layered_fast(true, RQRotationType::Fast)] + #[case::layered_matrix(true, RQRotationType::Matrix)] + #[tokio::test] + async fn test_rewrite_rq_row_layout_in_place( + #[case] layered: bool, + #[case] rotation_type: RQRotationType, + ) { + let source_dir = TempStrDir::default(); + let copy_dir_ = TempStrDir::default(); + let source = column_dataset( + source_dir.as_str(), + layered, + rotation_type, + LanceFileVersion::V2_2, + ) + .await; + let source_segment = single_rq_segment(&source, INDEX_NAME).await.unwrap(); + let source_files = list_files(std::path::Path::new(source_dir.as_str())); + let copy_path = format!("{}/copy", copy_dir_.as_str()); + let (converted, rewrite) = converted_copy( + source_dir.as_str(), + ©_path, + FileWriterOptions::default(), + ) + .await; + + assert_eq!(rewrite.old_uuid, source_segment.uuid.to_string()); + assert_ne!(rewrite.new_uuid, rewrite.old_uuid); + assert_eq!(rewrite.source_layout, "columns"); + assert_eq!(rewrite.target_layout, "plane_rows"); + assert_eq!(rewrite.index_version, 3); + assert_eq!(rewrite.num_rows, ROWS as u64); + assert_eq!(rewrite.num_partitions, PARTITIONS); + let expected_buffers = match rotation_type { + RQRotationType::Fast => 1, + RQRotationType::Matrix => 2, + }; + assert_eq!(rewrite.global_buffers, expected_buffers); + assert_eq!(rewrite.metadata.len(), 1); + let patch = &rewrite.metadata[0]; + assert_eq!(patch.key, STORAGE_METADATA_KEY); + let close = patch.source.rfind('}').unwrap(); + assert_eq!( + patch.converted, + format!( + r#"{},\"row_layout\":\"plane_rows\"{}"#, + &patch.source[..close], + &patch.source[close..] + ) + ); + + // The committed segment: the new uuid alone, at version 3, with the + // layout in its details and its files' sizes recorded. + let segments = converted.load_indices_by_name(INDEX_NAME).await.unwrap(); + assert_eq!(segments.len(), 1); + let segment = &segments[0]; + assert_eq!(segment.uuid.to_string(), rewrite.new_uuid); + assert_eq!(segment.index_version, 3); + assert_eq!(segment.fragment_bitmap, source_segment.fragment_bitmap); + assert_eq!(segment.dataset_version, source_segment.dataset_version); + let details = rq_details(segment).unwrap(); + let Some(Compression::Rq(rq)) = details.compression else { + unreachable!() + }; + assert_eq!(rq.row_layout(), rabit_quantization::RowLayout::PlaneRows); + let new_dir = converted.indices_dir().join(rewrite.new_uuid.as_str()); + for file in segment.files.as_ref().unwrap() { + let size = converted + .object_store + .size(&new_dir.clone().join(file.path.as_str())) + .await + .unwrap(); + assert_eq!(size, file.size_bytes, "{}", file.path); + } + let index = open_index(&converted).await; + let index = index.as_any().downcast_ref::().unwrap(); + assert_eq!(index.row_layout(), RQRowLayout::PlaneRows); + assert!(!index.resident_columns_enabled()); + assert_eq!(index.resident_columns_bytes(), 0); + + assert_same_results(&source, &converted, layered).await; + let verification = verify_rq_row_layout(&source, &converted, INDEX_NAME) + .await + .unwrap(); + assert!(verification.is_ok(), "{verification:#?}"); + assert_eq!(verification.deep_size_mismatches, 0, "{verification:#?}"); + assert!(verification.reads >= PARTITIONS); + assert_eq!(verification.source_layout, "columns"); + assert_eq!(verification.converted_layout, "plane_rows"); + assert_eq!(verification.global_buffer_digests.len(), expected_buffers); + assert_eq!( + verification.rotate_mat_position, + match rotation_type { + RQRotationType::Fast => None, + RQRotationType::Matrix => Some(2), + } + ); + assert_eq!(verification.pages.straddling_partitions, 0); + assert!(verification.pages.max_partition_bytes > 0); + assert_eq!(verification.pages.partitions_over_16mib, 0); + + // The version before the conversion still reads the old directory. + let previous = converted + .checkout_version(rewrite.dataset_version - 1) + .await + .unwrap(); + assert_eq!( + single_rq_segment(&previous, INDEX_NAME).await.unwrap().uuid, + source_segment.uuid + ); + assert_same_results(&source, &previous, layered).await; + // Removing it leaves the converted index whole. + converted + .object_store + .remove_dir_all(converted.indices_dir().join(rewrite.old_uuid.as_str())) + .await + .unwrap(); + let reopened = Dataset::open(©_path).await.unwrap(); + assert_same_results(&source, &reopened, layered).await; + + // The source was not written. + assert_eq!( + list_files(std::path::Path::new(source_dir.as_str())), + source_files + ); + assert_eq!( + single_rq_segment(&source, INDEX_NAME).await.unwrap().uuid, + source_segment.uuid + ); + } + + /// Converting back to columns restores the source's auxiliary file: + /// its fields, every value, its metadata byte for byte, its global + /// buffers, and version 2. + #[rstest] + #[tokio::test] + async fn test_rewrite_rq_row_layout_round_trips(#[values(false, true)] layered: bool) { + let source_dir = TempStrDir::default(); + let copy_dir_ = TempStrDir::default(); + let source = column_dataset( + source_dir.as_str(), + layered, + RQRotationType::Matrix, + LanceFileVersion::V2_2, + ) + .await; + let copy_path = format!("{}/copy", copy_dir_.as_str()); + let (mut converted, _) = converted_copy( + source_dir.as_str(), + ©_path, + FileWriterOptions::default(), + ) + .await; + let back = rewrite_rq_row_layout(&mut converted, INDEX_NAME, RQRowLayout::Columns) + .await + .unwrap(); + assert_eq!(back.index_version, 2); + assert_eq!(back.target_layout, "columns"); + assert_eq!( + converted.load_indices_by_name(INDEX_NAME).await.unwrap()[0].index_version, + 2 + ); + + let original = index_storage(&source).await; + let round_trip = index_storage(&converted).await; + assert_eq!(round_trip.row_layout(), RQRowLayout::Columns); + let schema = |storage: &IvfQuantizationStorage| { + ArrowSchema::from(storage.reader().schema().as_ref()) + }; + assert_eq!(schema(&original), schema(&round_trip)); + assert_eq!( + original.reader().metadata().file_buffers.len(), + round_trip.reader().metadata().file_buffers.len() + ); + for position in 1..original.reader().metadata().file_buffers.len() as u32 { + assert_eq!( + original + .reader() + .read_global_buffer(position) + .await + .unwrap(), + round_trip + .reader() + .read_global_buffer(position) + .await + .unwrap() + ); + } + for partition in 0..PARTITIONS { + assert_same_batch( + &original + .read_partition_batch(partition, None) + .await + .unwrap(), + &round_trip + .read_partition_batch(partition, None) + .await + .unwrap(), + &format!("partition {partition}"), + ); + } + assert_same_results(&source, &converted, layered).await; + } + + #[tokio::test] + async fn test_rewrite_rq_row_layout_rejects_what_it_cannot_convert() { + let source_dir = TempStrDir::default(); + let copy_dir_ = TempStrDir::default(); + let mut source = column_dataset( + source_dir.as_str(), + true, + RQRotationType::Fast, + LanceFileVersion::V2_2, + ) + .await; + let copy_path = format!("{}/copy", copy_dir_.as_str()); + let (mut converted, rewrite) = converted_copy( + source_dir.as_str(), + ©_path, + FileWriterOptions::default(), + ) + .await; + + // Already in the target layout. + let err = rewrite_rq_row_layout(&mut converted, INDEX_NAME, RQRowLayout::PlaneRows) + .await + .unwrap_err(); + assert!( + err.to_string().contains("already in the plane_rows layout"), + "{err}" + ); + let err = rewrite_rq_row_layout(&mut source, INDEX_NAME, RQRowLayout::Columns) + .await + .unwrap_err(); + assert!( + err.to_string().contains("already in the columns layout"), + "{err}" + ); + + // A dataset copied without the index's directory. + converted + .object_store + .remove_dir_all(converted.indices_dir().join(rewrite.new_uuid.as_str())) + .await + .unwrap(); + let err = rewrite_rq_row_layout(&mut converted, INDEX_NAME, RQRowLayout::Columns) + .await + .unwrap_err(); + assert!(err.to_string().contains("is missing"), "{err}"); + + // An index of several segments. + let batch = test_batch(200, ROWS as u64, 2); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + source.append(reader, None).await.unwrap(); + source + .optimize_indices(&OptimizeOptions::append()) + .await + .unwrap(); + assert_eq!( + source.load_indices_by_name(INDEX_NAME).await.unwrap().len(), + 2 + ); + let err = rewrite_rq_row_layout(&mut source, INDEX_NAME, RQRowLayout::PlaneRows) + .await + .unwrap_err(); + assert!(err.to_string().contains("2 segments"), "{err}"); + + // Not an IVF_RQ index. + let other_dir = TempStrDir::default(); + let mut other = write_dataset(other_dir.as_str(), LanceFileVersion::V2_2, ROWS).await; + other + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + &VectorIndexParams::with_ivf_flat_params(DistanceType::L2, ivf_params()), + true, + ) + .await + .unwrap(); + let err = rewrite_rq_row_layout(&mut other, INDEX_NAME, RQRowLayout::PlaneRows) + .await + .unwrap_err(); + assert!(err.to_string().contains("not an IVF_RQ index"), "{err}"); + } + + /// Text patches add the member last, or remove it with its separator, + /// and keep every other byte. + #[test] + fn test_patch_object_text() { + assert_eq!( + patch_object_text(r#"{"a":1}"#, PLANE_ROWS_MEMBER, true).unwrap(), + r#"{"a":1,"row_layout":"plane_rows"}"# + ); + assert_eq!( + patch_object_text("{ }", PLANE_ROWS_MEMBER, true).unwrap(), + r#"{ "row_layout":"plane_rows"}"# + ); + for (text, expected) in [ + (r#"{"a":1,"row_layout":"plane_rows"}"#, r#"{"a":1}"#), + (r#"{"row_layout":"plane_rows","a":1}"#, r#"{"a":1}"#), + ( + r#"{"a":1,"row_layout":"plane_rows","b":2}"#, + r#"{"a":1,"b":2}"#, + ), + (r#"{"row_layout":"plane_rows"}"#, "{}"), + ] { + assert_eq!( + patch_object_text(text, PLANE_ROWS_MEMBER, false).unwrap(), + expected + ); + } + assert!(patch_object_text(r#"{"a":1}"#, PLANE_ROWS_MEMBER, false).is_err()); + assert!(patch_object_text("[1]", PLANE_ROWS_MEMBER, true).is_err()); + + let storage = serde_json::to_string(&vec![ + r#"{"code_dim":64,"num_bits":7,"packed":true}"#.to_string(), + ]) + .unwrap(); + let metadata = HashMap::from([(STORAGE_METADATA_KEY.to_string(), storage.clone())]); + let patches = patch_schema_metadata(&metadata, RQRowLayout::PlaneRows).unwrap(); + let entries: Vec = serde_json::from_str(&patches[0].converted).unwrap(); + let read: RabitQuantizationMetadata = serde_json::from_str(&entries[0]).unwrap(); + assert_eq!(read.row_layout, RQRowLayout::PlaneRows); + // Metadata that already declares a layout is not patched again. + let declared = HashMap::from([( + STORAGE_METADATA_KEY.to_string(), + patches[0].converted.clone(), + )]); + assert!(patch_schema_metadata(&declared, RQRowLayout::PlaneRows).is_err()); + let back = patch_schema_metadata(&declared, RQRowLayout::Columns).unwrap(); + assert_eq!(back[0].converted, storage); + } + + /// Every read of a plane-row file returns what the same read of its + /// column-layout source returns: whole partitions, whole planes under + /// both bounds placements, selected rows at every coalesce gap, and no + /// rows; on 2.0 and 2.2 files, with pages that hold whole partitions or + /// end inside them. + #[rstest] + #[tokio::test] + async fn test_plane_rows_reads_match_column_reads( + #[values(false, true)] layered: bool, + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_2)] version: LanceFileVersion, + #[values(false, true)] small_pages: bool, + ) { + let source_dir = TempStrDir::default(); + let copy_dir_ = TempStrDir::default(); + let source = + column_dataset(source_dir.as_str(), layered, RQRotationType::Fast, version).await; + let options = if small_pages { + FileWriterOptions { + data_cache_bytes: Some(SMALL_PAGE_BYTES), + max_page_bytes: Some(SMALL_PAGE_BYTES), + ..Default::default() + } + } else { + FileWriterOptions::default() + }; + let copy_path = format!("{}/copy", copy_dir_.as_str()); + let (converted, _) = converted_copy_in_batches( + source_dir.as_str(), + ©_path, + options, + small_pages.then_some(SMALL_BATCH_ROWS), + ) + .await; + let verification = verify_rq_row_layout(&source, &converted, INDEX_NAME) + .await + .unwrap(); + // Whole reads and cache entries weigh the same in a cache; a sparse + // column read of a 2.0 file can keep more buffer than its values. + assert!(verification.is_ok(), "{verification:#?}"); + assert_eq!(verification.deep_size_mismatches, 0, "{verification:#?}"); + for file_version in [ + &verification.source_file_version, + &verification.converted_file_version, + ] { + assert_eq!(file_version, version.resolve().to_manifest_string()); + } + assert!( + verification.sparse_converted_deep_bytes <= verification.sparse_source_deep_bytes, + "{verification:#?}" + ); + if version != LanceFileVersion::V2_0 { + assert_eq!( + verification.sparse_deep_size_mismatches, 0, + "{verification:#?}" + ); + } + if small_pages { + assert!( + verification.pages.straddling_partitions > 0, + "{verification:#?}" + ); + } + + let mut rng = StdRng::seed_from_u64(u64::from(layered) + 2 * u64::from(small_pages)); + for sign_bounds in [SignBounds::Lazy, SignBounds::Eager] { + let columns = index_storage(&source).await.with_sign_bounds(sign_bounds); + let rows = index_storage(&converted) + .await + .with_sign_bounds(sign_bounds); + assert_eq!(rows.row_layout(), RQRowLayout::PlaneRows); + assert_eq!( + rows.logical_schema().unwrap().fields(), + columns.logical_schema().unwrap().fields() + ); + for partition in 0..PARTITIONS { + let context = format!("{sign_bounds:?} partition {partition}"); + assert_same_batch( + &columns.read_partition_batch(partition, None).await.unwrap(), + &rows.read_partition_batch(partition, None).await.unwrap(), + &context, + ); + let a = columns.load_partition(partition, None).await.unwrap(); + let b = rows.load_partition(partition, None).await.unwrap(); + assert_eq!(a.len(), b.len(), "{context}"); + assert_eq!( + a.row_ids().collect::>(), + b.row_ids().collect::>(), + "{context}" + ); + if !layered { + continue; + } + let size = columns.partition_size(partition) as u32; + let some: Vec = (0..size).filter(|_| rng.random_bool(0.2)).collect(); + let planes: &[u8] = match sign_bounds { + SignBounds::Lazy => &[0, 1, 2, SIGN_BOUNDS_PLANE], + SignBounds::Eager => &[0, 1, 2], + }; + for &plane in planes { + for selected in [None, Some(Vec::new()), Some(some.clone())] { + let context = format!( + "{context} plane {plane} rows {:?}", + selected.as_ref().map(Vec::len) + ); + assert_same_batch( + &columns + .read_plane(partition, plane, selected.clone(), None) + .await + .unwrap(), + &rows + .read_plane(partition, plane, selected.clone(), None) + .await + .unwrap(), + &context, + ); + } + for gap in [0, 64 * 1024, 1024 * 1024, u64::MAX] { + assert_same_batch( + &columns + .read_plane_with_coalesce_gap( + partition, + plane, + Some(some.clone()), + None, + Some(gap), + ) + .await + .unwrap(), + &rows + .read_plane_with_coalesce_gap( + partition, + plane, + Some(some.clone()), + None, + Some(gap), + ) + .await + .unwrap(), + &format!("{context} plane {plane} gap {gap}"), + ); + } + } + } + } + + // On 2.2 every packed page is full-zip: rows times the stride, flat. + if version == LanceFileVersion::V2_2 { + let storage = index_storage(&converted).await; + let spec = storage.plane_rows().unwrap().unwrap().clone(); + let reader = storage.reader(); + for column in spec.packed_columns() { + let projection = lance_file::versions::reader_projection_from_column_names( + reader.metadata().version(), + reader.schema(), + &[column.name()], + ) + .unwrap(); + let info = &reader.metadata().column_infos[projection.column_indices[0] as usize]; + assert!(!info.page_infos.is_empty()); + for page in info.page_infos.iter() { + let PageEncoding::Structural(layout) = &page.encoding else { + panic!("{} page is not structural", column.name()); + }; + assert!( + matches!(layout.layout, Some(Layout::FullZipLayout(_))), + "{} page is {layout:?}", + column.name() + ); + let bytes: u64 = page + .buffer_offsets_and_sizes + .iter() + .map(|(_, size)| size) + .sum(); + assert_eq!( + bytes, + page.num_rows * column.stride() as u64, + "{}", + column.name() + ); + } + } + } + } + + /// A copy of the plane-row dataset at `from` at `to`, converted in place + /// back to the column layout with `options`, at most `batch_rows` rows + /// per batch. + async fn column_copy_in_batches( + from: &str, + to: &str, + options: FileWriterOptions, + batch_rows: Option, + ) -> Dataset { + copy_dir(std::path::Path::new(from), std::path::Path::new(to)); + let mut dataset = Dataset::open(to).await.unwrap(); + rewrite_rq_row_layout_impl( + &mut dataset, + INDEX_NAME, + RQRowLayout::Columns, + options, + batch_rows, + ) + .await + .unwrap(); + dataset + } + + /// Partitions whose rows of `column` span more than one of its pages. + fn straddling_partitions( + storage: &IvfQuantizationStorage, + column: &str, + ) -> usize { + let reader = storage.reader(); + let projection = lance_file::versions::reader_projection_from_column_names( + reader.metadata().version(), + reader.schema(), + &[column], + ) + .unwrap(); + let info = &reader.metadata().column_infos[projection.column_indices[0] as usize]; + let mut page_starts = Vec::new(); + let mut next = 0; + for page in info.page_infos.iter() { + next += page.num_rows; + page_starts.push(next); + } + page_starts.pop(); + (0..storage.num_partitions()) + .filter(|&partition| { + let range = storage.ivf().row_range(partition); + page_starts + .iter() + .any(|&start| start > range.start as u64 && start < range.end as u64) + }) + .count() + } + + /// Twins weigh the same in a cache even where the column layout's reads + /// span pages: every whole read, partition storage (read alone or cut + /// from a prewarm's read) and code-only entry of a column file whose + /// pages end inside partitions takes the memory the plane-row twin's + /// does, on 2.0 and 2.2 files and in either direction of a conversion. + #[rstest] + #[tokio::test] + async fn test_column_entries_weigh_as_plane_rows( + #[values(false, true)] layered: bool, + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_2)] version: LanceFileVersion, + ) { + let source_dir = TempStrDir::default(); + let rows_dir = TempStrDir::default(); + let columns_dir = TempStrDir::default(); + column_dataset(source_dir.as_str(), layered, RQRotationType::Fast, version).await; + let rows_path = format!("{}/rows", rows_dir.as_str()); + let (rows, _) = converted_copy( + source_dir.as_str(), + &rows_path, + FileWriterOptions::default(), + ) + .await; + let columns = column_copy_in_batches( + &rows_path, + &format!("{}/columns", columns_dir.as_str()), + FileWriterOptions { + data_cache_bytes: Some(SMALL_PAGE_BYTES), + max_page_bytes: Some(SMALL_PAGE_BYTES), + ..Default::default() + }, + Some(SMALL_BATCH_ROWS), + ) + .await; + let storage = index_storage(&columns).await; + assert_eq!(storage.row_layout(), RQRowLayout::Columns); + assert!(straddling_partitions(&storage, RABIT_CODE_COLUMN) > 0); + + let entry_shapes: &[&str] = if layered { + &[ + deep_size_shapes::PARTITION_BATCH, + deep_size_shapes::PLANE, + deep_size_shapes::PARTITION_STORAGE, + deep_size_shapes::PREWARM_PARTITION_STORAGE, + deep_size_shapes::CODE_ONLY_ENTRY, + ] + } else { + &[ + deep_size_shapes::PARTITION_BATCH, + deep_size_shapes::PARTITION_STORAGE, + deep_size_shapes::PREWARM_PARTITION_STORAGE, + deep_size_shapes::CODE_ONLY_ENTRY, + deep_size_shapes::PREWARM_CODE_ONLY_ENTRY, + ] + }; + for (source, converted) in [(&columns, &rows), (&rows, &columns)] { + let verification = verify_rq_row_layout(source, converted, INDEX_NAME) + .await + .unwrap(); + assert!(verification.is_ok(), "{verification:#?}"); + assert_eq!(verification.deep_size_mismatches, 0, "{verification:#?}"); + assert_eq!( + verification.source_deep_bytes, verification.converted_deep_bytes, + "{verification:#?}" + ); + for shape in entry_shapes { + let sizes = &verification.deep_sizes[*shape]; + assert!(sizes.count >= PARTITIONS, "{shape}: {sizes:?}"); + assert_eq!(sizes.mismatches, 0, "{shape}: {sizes:?}"); + assert_eq!( + sizes.source_bytes, sizes.converted_bytes, + "{shape}: {sizes:?}" + ); + } + assert_eq!( + verification + .deep_sizes + .contains_key(deep_size_shapes::PLANE_SPARSE), + layered + ); + } + } + + /// A plane-row file reads a native partition, or a whole plane of a + /// layered one, in one request, an eager sign plane (sign and bounds) in + /// two, and selected rows in one request per run of adjacent rows; no + /// read loads a resident store, whatever residency asks. + #[rstest] + #[tokio::test] + async fn test_plane_rows_cut_origin_requests(#[values(false, true)] layered: bool) { + let source_dir = TempStrDir::default(); + let copy_dir_ = TempStrDir::default(); + column_dataset( + source_dir.as_str(), + layered, + RQRotationType::Fast, + LanceFileVersion::V2_2, + ) + .await; + let copy_path = format!("{}/copy", copy_dir_.as_str()); + let (converted, _) = converted_copy( + source_dir.as_str(), + ©_path, + FileWriterOptions::default(), + ) + .await; + let requests = |stats: &IoStats| stats.snapshot().iops; + + let lazy = index_storage(&converted) + .await + .with_resident_columns_enabled(true); + assert!(!lazy.resident_columns_enabled()); + assert_eq!(lazy.resident_columns_bytes(), 0); + let eager = index_storage(&converted) + .await + .with_sign_bounds(SignBounds::Eager); + let lazy = lazy.with_sign_bounds(SignBounds::Lazy); + let mut rng = StdRng::seed_from_u64(7); + for partition in 0..PARTITIONS { + let size = lazy.partition_size(partition) as u32; + let stats = IoStats::new(); + lazy.read_partition_batch(partition, Some(&stats)) + .await + .unwrap(); + // One request per packed column: one for a native partition, + // at most four for a layered one, whose plane pages can be close + // enough in the file to share a request. + let columns = if layered { 4 } else { 1 }; + if size == 0 { + assert_eq!(requests(&stats), 0, "partition {partition}"); + } else if layered { + assert!( + (1..=columns).contains(&requests(&stats)), + "partition {partition}: {} requests", + requests(&stats) + ); + } else { + assert_eq!(requests(&stats), 1, "partition {partition}"); + } + if !layered { + continue; + } + for (storage, plane, expected) in [ + (&lazy, 0, 1), + (&eager, 0, 2), + (&lazy, 1, 1), + (&lazy, 2, 1), + (&lazy, SIGN_BOUNDS_PLANE, 1), + ] { + let stats = IoStats::new(); + storage + .read_plane(partition, plane, None, Some(stats.clone())) + .await + .unwrap(); + let expected = if size == 0 { 0 } else { expected }; + assert_eq!( + requests(&stats), + expected, + "partition {partition} plane {plane}" + ); + } + let selected: Vec = (0..size).filter(|_| rng.random_bool(0.3)).collect(); + let runs = selected + .iter() + .enumerate() + .filter(|(index, row)| *index == 0 || selected[index - 1] + 1 != **row) + .count() as u64; + for plane in [1, 2] { + let stats = IoStats::new(); + lazy.read_plane_with_coalesce_gap( + partition, + plane, + Some(selected.clone()), + Some(stats.clone()), + Some(0), + ) + .await + .unwrap(); + assert_eq!( + requests(&stats), + runs, + "partition {partition} plane {plane}" + ); + } + } + } + + /// Building with plane-row params writes the values a column build + /// writes, packed: the same rows of every partition, at version 3. + #[rstest] + #[tokio::test] + async fn test_plane_rows_build_matches_conversion(#[values(false, true)] layered: bool) { + let columns_dir = TempStrDir::default(); + let rows_dir = TempStrDir::default(); + let columns = column_dataset( + columns_dir.as_str(), + layered, + RQRotationType::Fast, + LanceFileVersion::V2_2, + ) + .await; + let column_storage = index_storage(&columns).await; + let mut rows = write_dataset(rows_dir.as_str(), LanceFileVersion::V2_2, ROWS).await; + let mut rq = + rq_params(layered, RQRotationType::Fast).with_row_layout(RQRowLayout::PlaneRows); + rq.rotation = Some(column_storage.metadata().clone()); + create_index(&mut rows, rq).await; + assert_eq!( + rows.load_indices_by_name(INDEX_NAME).await.unwrap()[0].index_version, + 3 + ); + let row_storage = index_storage(&rows).await; + assert_eq!(row_storage.row_layout(), RQRowLayout::PlaneRows); + assert!(row_storage.metadata().packed); + assert_eq!(row_storage.ivf().lengths, column_storage.ivf().lengths); + + // Rows in id order, sign codes as each row's own: the partitions' + // row order can differ between builds. + let spec = PlaneRowsSpec::try_new(row_storage.metadata()).unwrap(); + let rows_by_id = |batch: RecordBatch| -> Vec<(u64, Vec)> { + let codes = Arc::new(unpack_codes(batch[RABIT_CODE_COLUMN].as_fixed_size_list())); + let batch = batch + .replace_column_by_name(RABIT_CODE_COLUMN, codes) + .unwrap(); + let packed = spec.pack(&batch).unwrap(); + let mut rows: Vec<(u64, Vec)> = Vec::new(); + let ids = batch[ROW_ID].as_primitive::(); + for row in 0..batch.num_rows() { + let mut bytes = Vec::new(); + for column in spec.packed_columns() { + let list = packed[column.name()].as_fixed_size_list(); + bytes.extend_from_slice(list.value(row).as_primitive::().values()); + } + rows.push((ids.value(row), bytes)); + } + rows.sort(); + rows + }; + for partition in 0..PARTITIONS { + assert_eq!( + rows_by_id( + column_storage + .read_partition_batch(partition, None) + .await + .unwrap() + ), + rows_by_id( + row_storage + .read_partition_batch(partition, None) + .await + .unwrap() + ), + "partition {partition}" + ); + } + } + + /// The ex planes a plane-row file reads serialize to the same cache + /// entries as the column layout's, so both share one NVMe format. + #[tokio::test] + async fn test_plane_rows_ex_plane_entries_match() { + let source_dir = TempStrDir::default(); + let copy_dir_ = TempStrDir::default(); + let source = column_dataset( + source_dir.as_str(), + true, + RQRotationType::Fast, + LanceFileVersion::V2_2, + ) + .await; + let copy_path = format!("{}/copy", copy_dir_.as_str()); + let (converted, _) = converted_copy( + source_dir.as_str(), + ©_path, + FileWriterOptions::default(), + ) + .await; + let codec = CacheCodec::from_impl::(); + let columns = index_storage(&source).await; + let rows = index_storage(&converted).await; + let entry = |batch: RecordBatch| { + let mut bytes = Vec::new(); + codec + .serialize( + &(Arc::new(PlaneBatch(batch)) as Arc), + &mut bytes, + ) + .unwrap(); + Bytes::from(bytes) + }; + for partition in 0..PARTITIONS { + for plane in [1, 2] { + assert_eq!( + entry( + columns + .read_plane(partition, plane, None, None) + .await + .unwrap() + ), + entry(rows.read_plane(partition, plane, None, None).await.unwrap()), + "partition {partition} plane {plane}" + ); + } + } + } + + /// A file whose schema does not match the layout its metadata declares + /// is an invalid index when it opens. + #[tokio::test] + async fn test_plane_rows_open_rejects_layout_mismatch() { + let source_dir = TempStrDir::default(); + let copy_dir_ = TempStrDir::default(); + let source = column_dataset( + source_dir.as_str(), + true, + RQRotationType::Fast, + LanceFileVersion::V2_2, + ) + .await; + let copy_path = format!("{}/copy", copy_dir_.as_str()); + let (converted, _) = converted_copy( + source_dir.as_str(), + ©_path, + FileWriterOptions::default(), + ) + .await; + let columns = index_storage(&source).await; + let rows = index_storage(&converted).await; + let store = converted.object_store.clone(); + let storage_metadata = |storage: &IvfQuantizationStorage| { + storage + .reader() + .schema() + .metadata + .get(STORAGE_METADATA_KEY) + .unwrap() + .clone() + }; + // Each file's columns under the other's storage metadata. + for (data, metadata, case) in [ + ( + &rows, + storage_metadata(&columns), + "plane rows declared as columns", + ), + ( + &columns, + storage_metadata(&rows), + "columns declared as plane rows", + ), + ] { + let reader = data.reader(); + let mut schema_metadata = reader.schema().metadata.clone(); + schema_metadata.insert(STORAGE_METADATA_KEY.to_string(), metadata); + let schema = ArrowSchema::new_with_metadata( + ArrowSchema::from(reader.schema().as_ref()).fields().clone(), + schema_metadata, + ); + let path = Path::from(format!("{copy_path}/mismatch/{}", case.replace(' ', "_"))); + let mut writer = file_versions::create_writer( + reader.metadata().version(), + store.create(&path).await.unwrap(), + lance_core::datatypes::Schema::try_from(&schema).unwrap(), + FileWriterOptions::default(), + ) + .unwrap(); + for position in 1..reader.metadata().file_buffers.len() as u32 { + writer + .add_global_buffer(reader.read_global_buffer(position).await.unwrap()) + .await + .unwrap(); + } + let batches = reader + .read_stream( + ReadBatchParams::RangeFull, + u32::MAX, + 1, + FilterExpression::no_filter(), + ) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + for batch in &batches { + writer.write_batch(batch).await.unwrap(); + } + writer.finish().await.unwrap(); + let err = open_rq_storage(&store, &path).await.unwrap_err(); + assert!(matches!(err, Error::Index { .. }), "{case}: {err:?}"); + } + } + + /// Remapping a plane-row index for compaction, at once or through a + /// fragment reuse index, keeps plane rows and version 3, and answers as + /// its column-layout twin remapped the same way. + #[rstest] + #[tokio::test] + async fn test_plane_rows_survive_remap(#[values(false, true)] defer_index_remap: bool) { + let columns_dir = TempStrDir::default(); + let copy_dir_ = TempStrDir::default(); + let mut columns = + write_dataset(columns_dir.as_str(), LanceFileVersion::V2_2, ROWS / 4).await; + create_index(&mut columns, rq_params(true, RQRotationType::Fast)).await; + let copy_path = format!("{}/copy", copy_dir_.as_str()); + let (mut rows, _) = converted_copy( + columns_dir.as_str(), + ©_path, + FileWriterOptions::default(), + ) + .await; + let options = CompactionOptions { + defer_index_remap, + ..Default::default() + }; + for dataset in [&mut columns, &mut rows] { + dataset.delete("id % 3 = 0").await.unwrap(); + compact_files(dataset, options.clone(), None).await.unwrap(); + } + let segment = single_rq_segment(&rows, INDEX_NAME).await.unwrap(); + assert_eq!(segment.index_version, 3); + let index = open_index(&rows).await; + let index = index.as_any().downcast_ref::().unwrap(); + assert_eq!(index.row_layout(), RQRowLayout::PlaneRows); + assert_eq!( + single_rq_segment(&columns, INDEX_NAME) + .await + .unwrap() + .index_version, + 2 + ); + assert_same_results(&columns, &rows, true).await; + } + + /// Appending to an index keeps its layout: a delta segment of a + /// plane-row index is plane rows at version 3, of a column-layout index + /// columns at version 2, and a merge through optimize writes the + /// reference segment's layout and version. A plane-row index answers as + /// its column-layout twin, built from the same model. + #[rstest] + #[tokio::test] + async fn test_optimize_append_keeps_row_layout_and_version( + #[values(false, true)] layered: bool, + ) { + let columns_dir = TempStrDir::default(); + let rows_dir = TempStrDir::default(); + let mut columns = column_dataset( + columns_dir.as_str(), + layered, + RQRotationType::Fast, + LanceFileVersion::V2_2, + ) + .await; + let mut rq = + rq_params(layered, RQRotationType::Fast).with_row_layout(RQRowLayout::PlaneRows); + rq.rotation = Some(index_storage(&columns).await.metadata().clone()); + let mut rows = write_dataset(rows_dir.as_str(), LanceFileVersion::V2_2, ROWS).await; + create_index(&mut rows, rq).await; + + let batch = test_batch(400, ROWS as u64, 3); + for (dataset, layout, version) in [ + (&mut columns, RQRowLayout::Columns, 2), + (&mut rows, RQRowLayout::PlaneRows, 3), + ] { + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + dataset.append(reader, None).await.unwrap(); + dataset + .optimize_indices(&OptimizeOptions::append()) + .await + .unwrap(); + let segments = dataset.load_indices_by_name(INDEX_NAME).await.unwrap(); + assert_eq!(segments.len(), 2, "{layout}"); + for segment in &segments { + assert_eq!(segment.index_version, version, "{layout}"); + let index = dataset + .open_vector_index("vector", &segment.uuid, &NoOpMetricsCollector) + .await + .unwrap(); + let index = index.as_any().downcast_ref::().unwrap(); + assert_eq!(index.row_layout(), layout); + } + } + assert_same_results(&columns, &rows, layered).await; + + for (dataset, layout, version) in [ + (&mut columns, RQRowLayout::Columns, 2), + (&mut rows, RQRowLayout::PlaneRows, 3), + ] { + dataset + .optimize_indices(&OptimizeOptions::merge(2)) + .await + .unwrap(); + let segment = single_rq_segment(dataset, INDEX_NAME).await.unwrap(); + assert_eq!(segment.index_version, version, "{layout}"); + assert_eq!(index_storage(dataset).await.row_layout(), layout); + } + assert_same_results(&columns, &rows, layered).await; + } + + /// A steady-state optimize that rewrites one oversized segment keeps its + /// layout and version. + #[tokio::test] + async fn test_segment_rebalance_keeps_layout_and_version() { + let dir = TempStrDir::default(); + let mut dataset = write_dataset(dir.as_str(), LanceFileVersion::V2_2, ROWS).await; + let mut ivf = ivf_params(); + ivf.target_partition_size = Some(256); + let params = VectorIndexParams::with_ivf_rq_params( + DistanceType::L2, + ivf, + rq_params(true, RQRotationType::Fast).with_row_layout(RQRowLayout::PlaneRows), + ); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some(INDEX_NAME.to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + // A delta segment with more than 4 x 256 rows in one partition. + let centroids = centroids(); + let mut values = Vec::with_capacity(1_500 * DIM); + let mut rng = StdRng::seed_from_u64(5); + for _ in 0..1_500 { + values.extend( + centroids[..DIM] + .iter() + .map(|value| value + rng.random_range(-0.05f32..0.05)), + ); + } + let vectors = + FixedSizeListArray::try_new_from_values(Float32Array::from(values), DIM as i32) + .unwrap(); + let schema = test_batch(1, 0, 0).schema(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(UInt64Array::from_iter_values( + ROWS as u64..ROWS as u64 + 1_500, + )), + Arc::new(vectors), + ], + ) + .unwrap(); + dataset + .append(RecordBatchIterator::new(vec![Ok(batch)], schema), None) + .await + .unwrap(); + dataset + .optimize_indices(&OptimizeOptions::append()) + .await + .unwrap(); + let before = dataset.load_indices_by_name(INDEX_NAME).await.unwrap(); + assert_eq!(before.len(), 2); + dataset + .optimize_indices(&OptimizeOptions::default()) + .await + .unwrap(); + let after = dataset.load_indices_by_name(INDEX_NAME).await.unwrap(); + assert_eq!(after.len(), 2, "one segment was rewritten in place"); + assert_ne!( + after.iter().map(|segment| segment.uuid).collect::>(), + before + .iter() + .map(|segment| segment.uuid) + .collect::>(), + "the steady-state optimize rewrote a segment" + ); + for segment in &after { + assert_eq!(segment.index_version, 3); + let index = dataset + .open_vector_index("vector", &segment.uuid, &NoOpMetricsCollector) + .await + .unwrap(); + let index = index.as_any().downcast_ref::().unwrap(); + assert_eq!(index.row_layout(), RQRowLayout::PlaneRows); + } + } + + /// Distributed segments, built from one model on fragment groups, merge + /// into one segment at their version in the column layout; the merger + /// does not support plane-row segments. + #[rstest] + #[tokio::test] + async fn test_merge_segments_by_row_layout( + #[values(RQRowLayout::Columns, RQRowLayout::PlaneRows)] row_layout: RQRowLayout, + ) { + let model_dir = TempStrDir::default(); + let dir = TempStrDir::default(); + let model = column_dataset( + model_dir.as_str(), + false, + RQRotationType::Fast, + LanceFileVersion::V2_2, + ) + .await; + let mut rq = rq_params(false, RQRotationType::Fast).with_row_layout(row_layout); + rq.rotation = Some(index_storage(&model).await.metadata().clone()); + let params = VectorIndexParams::with_ivf_rq_params(DistanceType::L2, ivf_params(), rq); + let mut dataset = write_dataset(dir.as_str(), LanceFileVersion::V2_2, ROWS / 2).await; + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 2); + let mut segments = Vec::new(); + for fragment in &fragments { + let segment = dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name(INDEX_NAME.to_string()) + .fragments(vec![fragment.id() as u32]) + .execute_uncommitted() + .await + .unwrap(); + assert_eq!(segment.index_version, ivf_rq_index_version(row_layout)); + segments.push(segment); + } + let merged = dataset.merge_existing_index_segments(segments).await; + match row_layout { + RQRowLayout::Columns => { + let merged = merged.unwrap(); + assert_eq!(merged.index_version, 2); + } + RQRowLayout::PlaneRows => { + let err = merged.unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. }), "{err:?}"); + } + } + } + + /// Segments of one index whose IVF_RQ rows are laid out differently, in + /// columns or plane rows, layered or native, are never merged into one: + /// an explicit merge fails naming both layouts, and a default optimize + /// indexes new rows in a segment of the last segment's layout and keeps + /// the others as they are. + #[rstest] + #[case::row_layout((false, RQRowLayout::Columns), (false, RQRowLayout::PlaneRows))] + #[case::layered((false, RQRowLayout::Columns), (true, RQRowLayout::Columns))] + #[case::layered_plane_rows((true, RQRowLayout::PlaneRows), (false, RQRowLayout::PlaneRows))] + #[tokio::test] + async fn test_optimize_does_not_merge_mixed_rq_layouts( + #[case] first: (bool, RQRowLayout), + #[case] second: (bool, RQRowLayout), + ) { + let model_dir = TempStrDir::default(); + let dir = TempStrDir::default(); + let model = column_dataset( + model_dir.as_str(), + false, + RQRotationType::Fast, + LanceFileVersion::V2_2, + ) + .await; + let rotation = index_storage(&model).await.metadata().clone(); + let mut dataset = write_dataset(dir.as_str(), LanceFileVersion::V2_2, ROWS / 2).await; + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 2); + let mut segments = Vec::new(); + for (fragment, (layered, row_layout)) in fragments.iter().zip([first, second]) { + let mut rq = rq_params(layered, RQRotationType::Fast).with_row_layout(row_layout); + rq.rotation = Some(rotation.clone()); + let params = VectorIndexParams::with_ivf_rq_params(DistanceType::L2, ivf_params(), rq); + segments.push( + dataset + .create_index_builder(&["vector"], IndexType::Vector, ¶ms) + .name(INDEX_NAME.to_string()) + .fragments(vec![fragment.id() as u32]) + .execute_uncommitted() + .await + .unwrap(), + ); + } + let uuids: Vec = segments.iter().map(|segment| segment.uuid).collect(); + dataset + .commit_existing_index_segments(INDEX_NAME, "vector", segments) + .await + .unwrap(); + let layout = |(layered, row_layout): (bool, RQRowLayout)| { + format!( + "{}/{row_layout}", + if layered { "layered" } else { "native" } + ) + }; + + let err = dataset + .optimize_indices(&OptimizeOptions::merge(2)) + .await + .unwrap_err(); + let message = err.to_string(); + assert!(message.contains("different row layouts"), "{message}"); + assert!(message.contains(&layout(first)), "{message}"); + assert!(message.contains(&layout(second)), "{message}"); + + let batch = test_batch(200, ROWS as u64, 2); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + dataset.append(reader, None).await.unwrap(); + dataset + .optimize_indices(&OptimizeOptions::default()) + .await + .unwrap(); + let segments = dataset.load_indices_by_name(INDEX_NAME).await.unwrap(); + assert_eq!(segments.len(), 3); + for uuid in &uuids { + assert!(segments.iter().any(|segment| segment.uuid == *uuid)); + } + let appended = segments + .iter() + .find(|segment| !uuids.contains(&segment.uuid)) + .unwrap(); + let Some(Compression::Rq(rq)) = rq_details(appended).unwrap().compression else { + unreachable!() + }; + let row_layout = match rq.row_layout() { + rabit_quantization::RowLayout::Columns => RQRowLayout::Columns, + rabit_quantization::RowLayout::PlaneRows => RQRowLayout::PlaneRows, + }; + assert_eq!(row_layout, second.1); + assert_eq!(appended.index_version, ivf_rq_index_version(second.1)); + let index = dataset + .open_vector_index("vector", &appended.uuid, &NoOpMetricsCollector) + .await + .unwrap(); + let index = index.as_any().downcast_ref::().unwrap(); + assert_eq!(index.row_layout(), second.1); + // Every row is found across the segments of both layouts. + let (ids, _) = search( + &dataset, + query(11).as_ref(), + ROWS + 200, + PARTITIONS, + RQPrecision::Full, + ) + .await; + assert_eq!(ids.len(), ROWS + 200); + } +}