diff --git a/rust/lance-encoding/benches/common/mod.rs b/rust/lance-encoding/benches/common/mod.rs new file mode 100644 index 00000000000..0ca46e250d7 --- /dev/null +++ b/rust/lance-encoding/benches/common/mod.rs @@ -0,0 +1,275 @@ +use std::sync::Arc; + +use arrow_schema::DataType; +use lance_core::{Error, Result, datatypes::Field}; +use lance_encoding::{ + compression::{ + BlockCompressor, CompressionStrategy, field_metadata_params, finalize_miniblock_compressor, + reject_packed_struct_per_value, try_bitpacking_block, try_bitpacking_miniblock, + try_byte_stream_split_miniblock, try_fixed_packed_struct_miniblock, try_fixed_u8_rle_block, + try_fixed_u8_rle_miniblock, try_general_block, try_raw_block, + try_raw_fixed_size_list_miniblock, try_raw_fixed_width_miniblock, try_raw_per_value, + try_uncompressed_fixed_width_miniblock, try_variable_packed_struct_per_value, + try_variable_width_miniblock, try_variable_width_per_value, + }, + compression_config::{CompressionFieldParams, CompressionParams}, + data::DataBlock, + encoder::{ + ColumnIndexSequence, FieldEncoder, FieldEncodingContext, FieldEncodingStrategy, + structural::{ + PrimitiveFieldEncoding, PrimitivePageEncoding, try_create_binary_blob, try_create_list, + try_create_map, try_create_struct, try_create_structural_blob, + try_create_structural_fixed_size_list, + }, + }, + encodings::logical::primitive::{fullzip::PerValueCompressor, miniblock::MiniBlockCompressor}, + format::pb21::CompressiveEncoding, +}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BenchEncoding { + Array, + StructuralU16, + StructuralU32, +} + +impl std::fmt::Display for BenchEncoding { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::Array => "array", + Self::StructuralU16 => "structural-u16", + Self::StructuralU32 => "structural-u32", + }) + } +} + +#[derive(Debug, Clone)] +struct BenchCompressionStrategy { + encoding: BenchEncoding, + params: CompressionParams, +} + +impl BenchCompressionStrategy { + fn field_params(&self, field: &Field) -> CompressionFieldParams { + let mut params = self + .params + .get_field_params(&field.name, &field.data_type()); + let mut metadata = field_metadata_params(field); + if self.encoding == BenchEncoding::StructuralU16 + && metadata + .minichunk_size + .is_some_and(|size| size >= 32 * 1024) + { + metadata.minichunk_size = None; + } + params.merge(&metadata); + params + } +} + +impl CompressionStrategy for BenchCompressionStrategy { + fn create_miniblock_compressor( + &self, + field: &Field, + data: &DataBlock, + ) -> Result> { + let params = self.field_params(field); + let compressor = + if let Some(compressor) = try_uncompressed_fixed_width_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = try_byte_stream_split_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = try_fixed_u8_rle_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = try_bitpacking_miniblock(data) { + compressor + } else if let Some(compressor) = try_raw_fixed_width_miniblock(data) { + compressor + } else if let Some(compressor) = try_variable_width_miniblock(field, data, ¶ms)? { + compressor + } else if let Some(compressor) = try_fixed_packed_struct_miniblock(data)? { + compressor + } else if let Some(compressor) = try_raw_fixed_size_list_miniblock(data) { + compressor + } else { + return Err(Error::not_supported_source( + format!( + "Mini-block compression not yet supported for block type {}", + data.name() + ) + .into(), + )); + }; + finalize_miniblock_compressor(data, compressor, ¶ms) + } + + fn create_per_value( + &self, + field: &Field, + data: &DataBlock, + ) -> Result> { + let params = self.field_params(field); + if let Some(compressor) = try_raw_per_value(data) { + return Ok(compressor); + } + let packed = match self.encoding { + BenchEncoding::StructuralU16 => reject_packed_struct_per_value(field, data)?, + BenchEncoding::StructuralU32 => { + try_variable_packed_struct_per_value(Arc::new(self.clone()), field, data)? + } + BenchEncoding::Array => unreachable!(), + }; + if let Some(compressor) = packed { + return Ok(compressor); + } + if let Some(compressor) = try_variable_width_per_value(field, data, ¶ms)? { + return Ok(compressor); + } + Err(Error::not_supported_source( + format!( + "Per-value compression not yet supported for block type {}", + data.name() + ) + .into(), + )) + } + + fn create_block_compressor( + &self, + field: &Field, + data: &DataBlock, + ) -> Result<(Box, CompressiveEncoding)> { + let params = self.field_params(field); + if self.encoding == BenchEncoding::StructuralU32 + && let Some(compressor) = try_fixed_u8_rle_block(data, ¶ms)? + { + return Ok(compressor); + } + if let Some(compressor) = try_bitpacking_block(data) { + return Ok(compressor); + } + if self.encoding == BenchEncoding::StructuralU32 + && let Some(compressor) = try_general_block(data, ¶ms)? + { + return Ok(compressor); + } + if let Some(compressor) = try_raw_block(data) { + return Ok(compressor); + } + Err(Error::not_supported_source( + format!( + "Block compression not yet supported for block type {}", + data.name() + ) + .into(), + )) + } +} + +#[derive(Debug)] +struct BenchFieldEncodingStrategy { + encoding: BenchEncoding, + primitive: PrimitiveFieldEncoding, +} + +impl FieldEncodingStrategy for BenchFieldEncodingStrategy { + fn create_field_encoder( + &self, + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, + ) -> Result> { + if let Some(encoder) = + try_create_binary_blob(&self.primitive, field, column_index, context)? + { + return Ok(encoder); + } + if self.encoding == BenchEncoding::StructuralU32 + && let Some(encoder) = + try_create_structural_blob(&self.primitive, field, column_index, context)? + { + return Ok(encoder); + } + if field.is_blob() { + return Err(Error::invalid_input_source( + format!( + "Blob encoding is not available for field '{}' with data type {}", + field.name, + field.data_type() + ) + .into(), + )); + } + if self.encoding == BenchEncoding::StructuralU32 { + if let Some(encoder) = try_create_map(field, column_index, context)? { + return Ok(encoder); + } + if let Some(encoder) = + try_create_structural_fixed_size_list(field, column_index, context)? + { + return Ok(encoder); + } + } + if let Some(encoder) = self.primitive.try_create(field, column_index, context)? { + return Ok(encoder); + } + if self.encoding == BenchEncoding::StructuralU16 { + if matches!( + field.data_type(), + DataType::FixedSizeList(item, _) + if matches!(item.data_type(), DataType::Struct(_)) + ) { + return Err(Error::not_supported_source( + "FixedSizeList is not enabled by the selected file format".into(), + )); + } + if matches!(field.data_type(), DataType::Map(_, _)) { + return Err(Error::not_supported_source( + "Map data type is not enabled by the selected file format".into(), + )); + } + } + if let Some(encoder) = try_create_list(field, column_index, context)? { + return Ok(encoder); + } + if let Some(encoder) = try_create_struct(field, column_index, context)? { + return Ok(encoder); + } + Err(Error::not_supported_source( + format!( + "{} has no field encoding for '{}' with data type {}", + self.encoding, + field.name, + field.data_type() + ) + .into(), + )) + } +} + +pub fn encoding_strategy(encoding: BenchEncoding) -> Box { + if encoding == BenchEncoding::Array { + return Box::new(lance_encoding::array_encoding::ArrayFieldEncodingStrategy::new()); + } + + let compression = Arc::new(BenchCompressionStrategy { + encoding, + params: CompressionParams::default(), + }); + let page_encodings = match encoding { + BenchEncoding::StructuralU16 => vec![ + PrimitivePageEncoding::reject_sparse(), + PrimitivePageEncoding::dense_u16(compression), + ], + BenchEncoding::StructuralU32 => vec![ + PrimitivePageEncoding::reject_sparse(), + PrimitivePageEncoding::constant(), + PrimitivePageEncoding::dense_u32(compression), + ], + BenchEncoding::Array => unreachable!(), + }; + Box::new(BenchFieldEncodingStrategy { + encoding, + primitive: PrimitiveFieldEncoding::new(page_encodings), + }) +} diff --git a/rust/lance-encoding/benches/decoder.rs b/rust/lance-encoding/benches/decoder.rs index cc0404e1bb3..21cc2334ece 100644 --- a/rust/lance-encoding/benches/decoder.rs +++ b/rust/lance-encoding/benches/decoder.rs @@ -11,15 +11,18 @@ use lance_core::cache::LanceCache; use lance_datagen::ArrayGeneratorExt; use lance_encoding::{ decoder::{ - DecodeBatchScheduler, DecoderConfig, DecoderPlugins, FilterExpression, create_decode_stream, + DecodeBatchScheduler, DecoderConfig, DecoderPlugins, EncodedBatchLayout, FilterExpression, + create_decode_stream, }, - encoder::{EncodingOptions, default_encoding_strategy, encode_batch}, - version::LanceFileVersion, + encoder::{EncodingOptions, encode_batch}, }; use tokio::sync::mpsc::unbounded_channel; use rand::Rng; +pub mod common; +use common::{BenchEncoding, encoding_strategy}; + const PRIMITIVE_TYPES: &[DataType] = &[ DataType::Date32, DataType::Date64, @@ -49,6 +52,15 @@ const PRIMITIVE_TYPES: &[DataType] = &[ // schema doesn't yet parse them in the context of a fixed size list. const PRIMITIVE_TYPES_FOR_FSL: &[DataType] = &[DataType::Int8, DataType::Float32]; +fn encoded_batch_layout(encoding: BenchEncoding) -> EncodedBatchLayout { + match encoding { + BenchEncoding::Array => EncodedBatchLayout::Array, + BenchEncoding::StructuralU16 | BenchEncoding::StructuralU32 => { + EncodedBatchLayout::Structural + } + } +} + fn bench_decode(c: &mut Criterion) { let rt = tokio::runtime::Runtime::new().unwrap(); let mut group = c.benchmark_group("decode_primitive"); @@ -64,7 +76,7 @@ fn bench_decode(c: &mut Criterion) { .unwrap(); let lance_schema = Arc::new(lance_core::datatypes::Schema::try_from(data.schema().as_ref()).unwrap()); - let encoding_strategy = default_encoding_strategy(LanceFileVersion::default()); + let encoding_strategy = encoding_strategy(BenchEncoding::StructuralU16); let encoded = rt .block_on(encode_batch( &data, @@ -81,7 +93,7 @@ fn bench_decode(c: &mut Criterion) { &FilterExpression::no_filter(), Arc::::default(), false, - LanceFileVersion::default(), + EncodedBatchLayout::Structural, Some(Arc::new(LanceCache::no_cache())), )) .unwrap(); @@ -95,14 +107,14 @@ fn bench_decode_fsl(c: &mut Criterion) { let rt = tokio::runtime::Runtime::new().unwrap(); let mut group = c.benchmark_group("decode_fsl"); const NUM_BYTES: u64 = 1024 * 1024 * 128; - for version in [ - LanceFileVersion::V2_0, - LanceFileVersion::V2_1, - LanceFileVersion::V2_2, + for encoding in [ + BenchEncoding::Array, + BenchEncoding::StructuralU16, + BenchEncoding::StructuralU32, ] { for data_type in PRIMITIVE_TYPES_FOR_FSL { for dimension in [4, 16, 32, 64, 128] { - let nullable_choices: &[bool] = if version == LanceFileVersion::V2_0 { + let nullable_choices: &[bool] = if encoding == BenchEncoding::Array { &[false] } else { &[false, true] @@ -110,7 +122,7 @@ fn bench_decode_fsl(c: &mut Criterion) { for nullable in nullable_choices { let func_name = format!( "{:?}_{}_v{}_null{}", - data_type, dimension, version, nullable + data_type, dimension, encoding, nullable ) .to_lowercase(); group.throughput(criterion::Throughput::Bytes(NUM_BYTES)); @@ -133,7 +145,7 @@ fn bench_decode_fsl(c: &mut Criterion) { lance_core::datatypes::Schema::try_from(data.schema().as_ref()) .unwrap(), ); - let encoding_strategy = default_encoding_strategy(version); + let encoding_strategy = encoding_strategy(encoding); let encoded = rt .block_on(encode_batch( &data, @@ -149,7 +161,7 @@ fn bench_decode_fsl(c: &mut Criterion) { &FilterExpression::no_filter(), Arc::::default(), false, - version, + encoded_batch_layout(encoding), Some(Arc::new(LanceCache::no_cache())), )) .unwrap(); @@ -199,7 +211,7 @@ fn bench_decode_str_with_dict_encoding(c: &mut Criterion) { let lance_schema = Arc::new(lance_core::datatypes::Schema::try_from(data.schema().as_ref()).unwrap()); - let encoding_strategy = default_encoding_strategy(LanceFileVersion::default()); + let encoding_strategy = encoding_strategy(BenchEncoding::StructuralU16); let encoded = rt .block_on(encode_batch( &data, @@ -215,7 +227,7 @@ fn bench_decode_str_with_dict_encoding(c: &mut Criterion) { &FilterExpression::no_filter(), Arc::::default(), false, - LanceFileVersion::default(), + EncodedBatchLayout::Structural, Some(Arc::new(LanceCache::no_cache())), )) .unwrap(); @@ -274,7 +286,7 @@ fn bench_decode_packed_struct(c: &mut Criterion) { RecordBatch::try_new(Arc::new(new_schema.clone()), data.columns().to_vec()).unwrap(); let lance_schema = Arc::new(lance_core::datatypes::Schema::try_from(&new_schema).unwrap()); - let encoding_strategy = default_encoding_strategy(LanceFileVersion::V2_2); + let encoding_strategy = encoding_strategy(BenchEncoding::StructuralU32); let encoded = rt .block_on(encode_batch( &data, @@ -291,7 +303,7 @@ fn bench_decode_packed_struct(c: &mut Criterion) { &FilterExpression::no_filter(), Arc::::default(), false, - LanceFileVersion::V2_2, + EncodedBatchLayout::Structural, Some(Arc::new(LanceCache::no_cache())), )) .unwrap(); @@ -331,7 +343,7 @@ fn bench_decode_str_with_fixed_size_binary_encoding(c: &mut Criterion) { let lance_schema = Arc::new(lance_core::datatypes::Schema::try_from(data.schema().as_ref()).unwrap()); - let encoding_strategy = default_encoding_strategy(LanceFileVersion::default()); + let encoding_strategy = encoding_strategy(BenchEncoding::StructuralU16); let encoded = rt .block_on(encode_batch( &data, @@ -347,7 +359,7 @@ fn bench_decode_str_with_fixed_size_binary_encoding(c: &mut Criterion) { &FilterExpression::no_filter(), Arc::::default(), false, - LanceFileVersion::default(), + EncodedBatchLayout::Structural, Some(Arc::new(LanceCache::no_cache())), )) .unwrap(); @@ -398,7 +410,7 @@ fn bench_decode_compressed(c: &mut Criterion) { let lance_schema = Arc::new(lance_core::datatypes::Schema::try_from(schema.as_ref()).unwrap()); // V2_2+ required for general compression - let encoding_strategy = default_encoding_strategy(LanceFileVersion::V2_2); + let encoding_strategy = encoding_strategy(BenchEncoding::StructuralU32); // Encode once during setup let encoded = rt @@ -423,7 +435,7 @@ fn bench_decode_compressed(c: &mut Criterion) { &FilterExpression::no_filter(), Arc::::default(), false, - LanceFileVersion::V2_2, + EncodedBatchLayout::Structural, Some(Arc::new(LanceCache::no_cache())), )) .unwrap(); @@ -476,7 +488,7 @@ fn bench_decode_compressed_parallel(c: &mut Criterion) { let lance_schema = Arc::new(lance_core::datatypes::Schema::try_from(schema.as_ref()).unwrap()); - let encoding_strategy = default_encoding_strategy(LanceFileVersion::V2_2); + let encoding_strategy = encoding_strategy(BenchEncoding::StructuralU32); let encoded = rt .block_on(encode_batch( diff --git a/rust/lance-encoding/benches/encoder.rs b/rust/lance-encoding/benches/encoder.rs index 08eb89d32fe..02ffd92083c 100644 --- a/rust/lance-encoding/benches/encoder.rs +++ b/rust/lance-encoding/benches/encoder.rs @@ -7,15 +7,15 @@ use arrow_array::{ArrayRef, BooleanArray, ListArray, RecordBatch}; use arrow_buffer::{OffsetBuffer, ScalarBuffer}; use arrow_schema::{DataType, Field, Schema}; use criterion::{Criterion, criterion_group, criterion_main}; -use lance_encoding::{ - encoder::{EncodingOptions, default_encoding_strategy, encode_batch}, - version::LanceFileVersion, -}; +use lance_encoding::encoder::{EncodingOptions, encode_batch}; + +pub mod common; +use common::{BenchEncoding, encoding_strategy}; fn encode_batch_sync(rt: &tokio::runtime::Runtime, data: &RecordBatch) { let lance_schema = Arc::new(lance_core::datatypes::Schema::try_from(data.schema().as_ref()).unwrap()); - let encoding_strategy = default_encoding_strategy(LanceFileVersion::V2_2); + let encoding_strategy = encoding_strategy(BenchEncoding::StructuralU32); rt.block_on(encode_batch( data, @@ -68,7 +68,7 @@ fn bench_encode_compressed(c: &mut Criterion) { let lance_schema = Arc::new(lance_core::datatypes::Schema::try_from(schema.as_ref()).unwrap()); // V2_2+ required for general compression - let encoding_strategy = default_encoding_strategy(LanceFileVersion::V2_2); + let encoding_strategy = encoding_strategy(BenchEncoding::StructuralU32); group.throughput(criterion::Throughput::Elements( (NUM_ROWS * NUM_COLUMNS) as u64, diff --git a/rust/lance-encoding/src/array_encoding.rs b/rust/lance-encoding/src/array_encoding.rs new file mode 100644 index 00000000000..a76c6a186ec --- /dev/null +++ b/rust/lance-encoding/src/array_encoding.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Encoding and decoding mechanisms described by [`crate::format::pb::ArrayEncoding`]. +//! +//! File versions decide which mechanisms to compose and accept. This module +//! contains only the reusable implementation of that persisted grammar. + +pub mod logical; +pub mod physical; +mod strategy; + +pub use strategy::ArrayFieldEncodingStrategy; diff --git a/rust/lance-encoding/src/previous/encodings/logical.rs b/rust/lance-encoding/src/array_encoding/logical.rs similarity index 100% rename from rust/lance-encoding/src/previous/encodings/logical.rs rename to rust/lance-encoding/src/array_encoding/logical.rs diff --git a/rust/lance-encoding/src/previous/encodings/logical/binary.rs b/rust/lance-encoding/src/array_encoding/logical/binary.rs similarity index 97% rename from rust/lance-encoding/src/previous/encodings/logical/binary.rs rename to rust/lance-encoding/src/array_encoding/logical/binary.rs index 00715f64511..697c1503e0a 100644 --- a/rust/lance-encoding/src/previous/encodings/logical/binary.rs +++ b/rust/lance-encoding/src/array_encoding/logical/binary.rs @@ -19,7 +19,7 @@ use crate::{ DecodeArrayTask, FilterExpression, MessageType, NextDecodeTask, PriorityRange, ScheduledScanLine, SchedulerContext, }, - previous::decoder::{DecoderReady, FieldScheduler, LogicalPageDecoder, SchedulingJob}, + decoder::{DecoderReady, FieldScheduler, LogicalPageDecoder, SchedulingJob}, }; /// Wraps a varbin scheduler and uses a BinaryPageDecoder to cast @@ -41,7 +41,7 @@ impl SchedulingJob for BinarySchedulingJob<'_> { .decoders .into_iter() .map(|message| { - let decoder = message.into_legacy(); + let decoder = message.into_array(); MessageType::DecoderReady(DecoderReady { path: decoder.path, decoder: Box::new(BinaryPageDecoder { @@ -179,7 +179,7 @@ impl DecodeArrayTask for BinaryArrayDecoder { DataType::LargeUtf8 => Self::from_list_array::(arr.as_list::()), _ => panic!("Binary decoder does not support this data type"), }; - // data_size is only tracked in the v2.1 structural decode path; the legacy + // data_size is only tracked in the v2.1 structural decode path; the v2.0 array // v2.0 path does not need it so we return 0. Ok((result, 0)) } diff --git a/rust/lance-encoding/src/previous/encodings/logical/blob.rs b/rust/lance-encoding/src/array_encoding/logical/blob.rs similarity index 96% rename from rust/lance-encoding/src/previous/encodings/logical/blob.rs rename to rust/lance-encoding/src/array_encoding/logical/blob.rs index 20f6afc5b36..0bf3a39fca1 100644 --- a/rust/lance-encoding/src/previous/encodings/logical/blob.rs +++ b/rust/lance-encoding/src/array_encoding/logical/blob.rs @@ -23,9 +23,9 @@ use crate::{ DecodeArrayTask, FilterExpression, MessageType, NextDecodeTask, PriorityRange, ScheduledScanLine, SchedulerContext, }, + decoder::{DecoderReady, FieldScheduler, LogicalPageDecoder, SchedulingJob}, encoder::{EncodeTask, FieldEncoder, OutOfLineBuffers}, format::pb::{Blob, ColumnEncoding, column_encoding}, - previous::decoder::{DecoderReady, FieldScheduler, LogicalPageDecoder, SchedulingJob}, repdef::RepDefBuilder, }; @@ -67,7 +67,7 @@ impl SchedulingJob for BlobFieldSchedulingJob<'_> { let next_descriptions = self.descriptions_job.schedule_next(context, priority)?; let mut priority = priority.current_priority(); let decoders = next_descriptions.decoders.into_iter().map(|decoder| { - let decoder = decoder.into_legacy(); + let decoder = decoder.into_array(); let path = decoder.path; let mut decoder = decoder.decoder; let num_rows = decoder.num_rows(); @@ -285,7 +285,7 @@ impl DecodeArrayTask for BlobArrayDecodeTask { buffer.extend_from_slice(&bytes); } let data_buf = Buffer::from_vec(buffer); - // data_size is only tracked in the v2.1 structural decode path; the legacy + // data_size is only tracked in the v2.1 structural decode path; the v2.0 array // v2.0 path does not need it so we return 0. Ok(( Arc::new(LargeBinaryArray::new(offsets, data_buf, self.validity)), @@ -417,10 +417,10 @@ mod tests { use super::BlobFieldDecoder; use crate::{ EncodingsIo, + decoder::LogicalPageDecoder, format::pb::column_encoding, - previous::decoder::LogicalPageDecoder, + testing::TestEncoding, testing::{TestCases, check_round_trip_encoding_of_data, check_specific_random}, - version::LanceFileVersion, }; static BLOB_META: LazyLock> = LazyLock::new(|| { @@ -433,11 +433,7 @@ mod tests { #[test_log::test(tokio::test)] async fn test_basic_blob() { let field = Field::new("", DataType::LargeBinary, false).with_metadata(BLOB_META.clone()); - check_specific_random( - field, - TestCases::basic().with_max_file_version(LanceFileVersion::V2_1), - ) - .await; + check_specific_random(field, TestCases::basic().with_array_and_u16_encodings()).await; } #[test_log::test(tokio::test)] @@ -446,10 +442,10 @@ mod tests { let val2: &[u8] = &[7, 8, 9]; let array = Arc::new(LargeBinaryArray::from(vec![Some(val1), None, Some(val2)])); let test_cases = TestCases::default() - .with_max_file_version(LanceFileVersion::V2_1) + .with_array_and_u16_encodings() .with_expected_encoding("packed_struct") - .with_verify_encoding(Arc::new(|cols, version| { - if version < &LanceFileVersion::V2_1 { + .with_verify_encoding(Arc::new(|cols, encoding| { + if *encoding == TestEncoding::Array { // In 2.0 we used a special "column encoding" to mark blob fields. In 2.1 we // don't do this and just rely on the regular page encoding. assert_eq!(cols.len(), 1); @@ -465,9 +461,9 @@ mod tests { .await; let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_1) - .with_verify_encoding(Arc::new(|cols, version| { - if version < &LanceFileVersion::V2_1 { + .with_structural_encodings() + .with_verify_encoding(Arc::new(|cols, encoding| { + if *encoding == TestEncoding::Array { assert_eq!(cols.len(), 1); let col = &cols[0]; assert!(!matches!( diff --git a/rust/lance-encoding/src/previous/encodings/logical/list.rs b/rust/lance-encoding/src/array_encoding/logical/list.rs similarity index 98% rename from rust/lance-encoding/src/previous/encodings/logical/list.rs rename to rust/lance-encoding/src/array_encoding/logical/list.rs index 3de886a21db..10ffb215798 100644 --- a/rust/lance-encoding/src/previous/encodings/logical/list.rs +++ b/rust/lance-encoding/src/array_encoding/logical/list.rs @@ -22,19 +22,19 @@ use tokio::task::JoinHandle; use crate::{ EncodingsIo, + array_encoding::logical::r#struct::{SimpleStructDecoder, SimpleStructScheduler}, buffer::LanceBuffer, data::{BlockInfo, DataBlock, FixedWidthDataBlock}, decoder::{ - DecodeArrayTask, DecodeBatchScheduler, FilterExpression, ListPriorityRange, MessageType, - NextDecodeTask, PageEncoding, PriorityRange, ScheduledScanLine, SchedulerContext, + DecodeArrayTask, DecodeBatchScheduler, FieldScheduler, FilterExpression, ListPriorityRange, + LogicalPageDecoder, MessageType, NextDecodeTask, PageEncoding, PriorityRange, + ScheduledScanLine, SchedulerContext, SchedulingJob, }, - encoder::{EncodeTask, EncodedColumn, EncodedPage, FieldEncoder, OutOfLineBuffers}, - format::pb, - previous::{ - decoder::{FieldScheduler, LogicalPageDecoder, SchedulingJob}, - encoder::{ArrayEncoder, EncodedArray}, - encodings::logical::r#struct::{SimpleStructDecoder, SimpleStructScheduler}, + encoder::{ + ArrayEncoder, EncodeTask, EncodedArray, EncodedColumn, EncodedPage, FieldEncoder, + OutOfLineBuffers, }, + format::pb, repdef::RepDefBuilder, utils::accumulation::AccumulationQueue, }; @@ -397,7 +397,7 @@ async fn indirect_schedule_task( for message in indirect_messages { for decoder in message.decoders { - let decoder = decoder.into_legacy(); + let decoder = decoder.into_array(); if !decoder.path.is_empty() { root_decoder.accept_child(decoder)?; } @@ -467,7 +467,7 @@ impl SchedulingJob for ListFieldSchedulingJob<'_> { .into_iter() .next() .unwrap() - .into_legacy() + .into_array() .decoder; let items_scheduler = self.scheduler.items_scheduler.clone(); @@ -690,7 +690,7 @@ impl DecodeArrayTask for ListDecodeTask { } _ => panic!("ListDecodeTask with data type that is not i32 or i64"), }; - // data_size is only tracked in the v2.1 structural decode path; the legacy + // data_size is only tracked in the v2.1 structural decode path; the v2.0 array // v2.0 path does not need it so we return 0. Ok((array, 0)) } @@ -958,10 +958,10 @@ impl ListOffsetsEncoder { let (data, description) = array.into_buffers(); Ok(EncodedPage { data, - description: PageEncoding::Legacy(description), + description: PageEncoding::Array(description), num_rows, column_idx, - row_number: 0, // Legacy encoders do not use + row_number: 0, // V2.0 encoders do not use }) }) .map(|res_res| res_res.unwrap()) diff --git a/rust/lance-encoding/src/previous/encodings/logical/primitive.rs b/rust/lance-encoding/src/array_encoding/logical/primitive.rs similarity index 97% rename from rust/lance-encoding/src/previous/encodings/logical/primitive.rs rename to rust/lance-encoding/src/array_encoding/logical/primitive.rs index d1debf3ef33..a35bc316a1a 100644 --- a/rust/lance-encoding/src/previous/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/array_encoding/logical/primitive.rs @@ -10,10 +10,10 @@ use futures::{FutureExt, future::BoxFuture}; use log::trace; use crate::decoder::{ColumnBuffers, PageBuffers}; -use crate::previous::decoder::{FieldScheduler, LogicalPageDecoder, SchedulingJob}; -use crate::previous::encoder::ArrayEncodingStrategy; +use crate::decoder::{FieldScheduler, LogicalPageDecoder, SchedulingJob}; +use crate::encoder::ArrayEncodingStrategy; use crate::utils::accumulation::AccumulationQueue; -use crate::{data::DataBlock, previous::encodings::physical::decoder_from_array_encoding}; +use crate::{array_encoding::physical::decoder_from_array_encoding, data::DataBlock}; use lance_core::{Error, Result, datatypes::Field}; use crate::{ @@ -74,7 +74,7 @@ impl PrimitiveFieldScheduler { positions_and_sizes: &page.buffer_offsets_and_sizes, }; let scheduler = decoder_from_array_encoding( - page.encoding.as_legacy(), + page.encoding.as_array(), &page_buffers, &data_type, ); @@ -316,7 +316,7 @@ impl DecodeArrayTask for PrimitiveFieldDecodeTask { return Ok((new_array, 0)); } } - // data_size is only tracked in the v2.1 structural decode path; the legacy + // data_size is only tracked in the v2.1 structural decode path; the v2.0 array // v2.0 path does not need it so we return 0. Ok((array, 0)) } @@ -436,10 +436,10 @@ impl PrimitiveFieldEncoder { let (data, description) = array.into_buffers(); Ok(EncodedPage { data, - description: PageEncoding::Legacy(description), + description: PageEncoding::Array(description), num_rows: num_values, column_idx, - row_number: 0, // legacy encoders do not use + row_number: 0, // v2.0 encoders do not use }) }) .map(|res_res| { diff --git a/rust/lance-encoding/src/previous/encodings/logical/struct.rs b/rust/lance-encoding/src/array_encoding/logical/struct.rs similarity index 99% rename from rust/lance-encoding/src/previous/encodings/logical/struct.rs rename to rust/lance-encoding/src/array_encoding/logical/struct.rs index b117f74ce2b..045b3bca71d 100644 --- a/rust/lance-encoding/src/previous/encodings/logical/struct.rs +++ b/rust/lance-encoding/src/array_encoding/logical/struct.rs @@ -12,7 +12,7 @@ use crate::{ DecodeArrayTask, FilterExpression, MessageType, NextDecodeTask, PriorityRange, ScheduledScanLine, SchedulerContext, }, - previous::decoder::{DecoderReady, FieldScheduler, LogicalPageDecoder, SchedulingJob}, + decoder::{DecoderReady, FieldScheduler, LogicalPageDecoder, SchedulingJob}, }; use arrow_array::{ArrayRef, StructArray}; use arrow_schema::{DataType, Field, Fields}; @@ -57,7 +57,7 @@ struct EmptyStructDecodeTask { impl DecodeArrayTask for EmptyStructDecodeTask { fn decode(self: Box) -> Result<(ArrayRef, u64)> { - // data_size is only tracked in the v2.1 structural decode path; the legacy + // data_size is only tracked in the v2.1 structural decode path; the v2.0 array // v2.0 path does not need it so we return 0. Ok(( Arc::new(StructArray::new_empty_fields(self.num_rows as usize, None)), @@ -607,7 +607,7 @@ impl DecodeArrayTask for SimpleStructDecodeTask { .into_iter() .map(|child| child.decode()) .collect::>>()?; - // data_size is only tracked in the v2.1 structural decode path; the legacy + // data_size is only tracked in the v2.1 structural decode path; the v2.0 array // v2.0 path does not need it so we return 0. Ok(( Arc::new(StructArray::try_new(self.child_fields, child_arrays, None)?), diff --git a/rust/lance-encoding/src/previous/encodings/physical.rs b/rust/lance-encoding/src/array_encoding/physical.rs similarity index 99% rename from rust/lance-encoding/src/previous/encodings/physical.rs rename to rust/lance-encoding/src/array_encoding/physical.rs index a3cb0adb4a7..66af9e806df 100644 --- a/rust/lance-encoding/src/previous/encodings/physical.rs +++ b/rust/lance-encoding/src/array_encoding/physical.rs @@ -5,16 +5,16 @@ use arrow_schema::DataType; use lance_arrow::DataTypeExt; use crate::{ - buffer::LanceBuffer, - decoder::{PageBuffers, PageScheduler}, - encodings::physical::block::{CompressionConfig, CompressionScheme}, - format::pb::{self, PackedStruct}, - previous::encodings::physical::{ + array_encoding::physical::{ basic::BasicPageScheduler, binary::BinaryPageScheduler, bitmap::DenseBitmapScheduler, dictionary::DictionaryPageScheduler, fixed_size_list::FixedListScheduler, fsst::FsstPageScheduler, packed_struct::PackedStructPageScheduler, value::ValuePageScheduler, }, + buffer::LanceBuffer, + decoder::{PageBuffers, PageScheduler}, + encodings::physical::block::{CompressionConfig, CompressionScheme}, + format::pb::{self, PackedStruct}, }; pub mod basic; @@ -293,9 +293,9 @@ pub fn decoder_from_array_encoding( #[cfg(test)] mod tests { + use crate::array_encoding::physical::get_buffer_decoder; use crate::decoder::{ColumnBuffers, FileBuffers, PageBuffers}; use crate::format::pb; - use crate::previous::encodings::physical::get_buffer_decoder; #[test] fn test_get_buffer_decoder_for_compressed_buffer() { diff --git a/rust/lance-encoding/src/previous/encodings/physical/basic.rs b/rust/lance-encoding/src/array_encoding/physical/basic.rs similarity index 99% rename from rust/lance-encoding/src/previous/encodings/physical/basic.rs rename to rust/lance-encoding/src/array_encoding/physical/basic.rs index ec098c3fff4..6dd7326a97f 100644 --- a/rust/lance-encoding/src/previous/encodings/physical/basic.rs +++ b/rust/lance-encoding/src/array_encoding/physical/basic.rs @@ -11,8 +11,8 @@ use crate::{ EncodingsIo, data::{AllNullDataBlock, BlockInfo, DataBlock, NullableDataBlock}, decoder::{PageScheduler, PrimitivePageDecoder}, + encoder::{ArrayEncoder, EncodedArray}, format::ProtobufUtils, - previous::encoder::{ArrayEncoder, EncodedArray}, }; use lance_core::Result; diff --git a/rust/lance-encoding/src/previous/encodings/physical/binary.rs b/rust/lance-encoding/src/array_encoding/physical/binary.rs similarity index 99% rename from rust/lance-encoding/src/previous/encodings/physical/binary.rs rename to rust/lance-encoding/src/array_encoding/physical/binary.rs index 0bd4d96b45e..294b295def3 100644 --- a/rust/lance-encoding/src/previous/encodings/physical/binary.rs +++ b/rust/lance-encoding/src/array_encoding/physical/binary.rs @@ -12,17 +12,17 @@ use futures::TryFutureExt; use futures::{FutureExt, future::BoxFuture}; +use crate::array_encoding::logical::primitive::PrimitiveFieldDecoder; use crate::buffer::LanceBuffer; use crate::data::{ BlockInfo, DataBlock, FixedWidthDataBlock, NullableDataBlock, VariableWidthBlock, }; +use crate::decoder::LogicalPageDecoder; +use crate::encoder::{ArrayEncoder, EncodedArray}; use crate::encodings::physical::block::{ BufferCompressor, CompressionConfig, GeneralBufferCompressor, }; use crate::format::ProtobufUtils; -use crate::previous::decoder::LogicalPageDecoder; -use crate::previous::encoder::{ArrayEncoder, EncodedArray}; -use crate::previous::encodings::logical::primitive::PrimitiveFieldDecoder; use crate::{ EncodingsIo, decoder::{PageScheduler, PrimitivePageDecoder}, diff --git a/rust/lance-encoding/src/previous/encodings/physical/bitmap.rs b/rust/lance-encoding/src/array_encoding/physical/bitmap.rs similarity index 99% rename from rust/lance-encoding/src/previous/encodings/physical/bitmap.rs rename to rust/lance-encoding/src/array_encoding/physical/bitmap.rs index fc6d295b0f9..2168aff24e2 100644 --- a/rust/lance-encoding/src/previous/encodings/physical/bitmap.rs +++ b/rust/lance-encoding/src/array_encoding/physical/bitmap.rs @@ -131,9 +131,9 @@ mod tests { use bytes::Bytes; use std::{collections::HashMap, sync::Arc}; + use crate::array_encoding::physical::bitmap::BitmapData; use crate::data::{DataBlock, FixedWidthDataBlock}; use crate::decoder::PrimitivePageDecoder; - use crate::previous::encodings::physical::bitmap::BitmapData; use crate::testing::{TestCases, check_basic_random, check_round_trip_encoding_of_data}; use super::BitmapDecoder; diff --git a/rust/lance-encoding/src/previous/encodings/physical/bitpack.rs b/rust/lance-encoding/src/array_encoding/physical/bitpack.rs similarity index 99% rename from rust/lance-encoding/src/previous/encodings/physical/bitpack.rs rename to rust/lance-encoding/src/array_encoding/physical/bitpack.rs index d80dec351d1..856af8eb8b4 100644 --- a/rust/lance-encoding/src/previous/encodings/physical/bitpack.rs +++ b/rust/lance-encoding/src/array_encoding/physical/bitpack.rs @@ -23,8 +23,8 @@ use crate::buffer::LanceBuffer; use crate::data::BlockInfo; use crate::data::{DataBlock, FixedWidthDataBlock, NullableDataBlock}; use crate::decoder::{PageScheduler, PrimitivePageDecoder}; +use crate::encoder::{ArrayEncoder, EncodedArray}; use crate::format::ProtobufUtils; -use crate::previous::encoder::{ArrayEncoder, EncodedArray}; use bytemuck::cast_slice; const LOG_ELEMS_PER_CHUNK: u8 = 10; @@ -1146,7 +1146,6 @@ pub mod test { use crate::{ format::pb, testing::{ArrayGeneratorProvider, TestCases, check_round_trip_encoding_generated}, - version::LanceFileVersion, }; use super::*; @@ -1677,7 +1676,7 @@ pub mod test { for (data_type, array_gen_provider) in bitpacked_test_cases { let field = Field::new("", data_type.clone(), false); - let test_cases = TestCases::basic().with_min_file_version(LanceFileVersion::V2_1); + let test_cases = TestCases::basic().with_structural_encodings(); check_round_trip_encoding_generated(field, array_gen_provider.copy(), test_cases).await; } } diff --git a/rust/lance-encoding/src/previous/encodings/physical/block.rs b/rust/lance-encoding/src/array_encoding/physical/block.rs similarity index 97% rename from rust/lance-encoding/src/previous/encodings/physical/block.rs rename to rust/lance-encoding/src/array_encoding/physical/block.rs index 517bce4e3a7..3bbd966e24a 100644 --- a/rust/lance-encoding/src/previous/encodings/physical/block.rs +++ b/rust/lance-encoding/src/array_encoding/physical/block.rs @@ -5,9 +5,9 @@ use arrow_schema::DataType; use crate::{ data::{BlockInfo, DataBlock, OpaqueBlock}, + encoder::{ArrayEncoder, EncodedArray}, encodings::physical::block::{CompressedBufferEncoder, CompressionConfig, CompressionScheme}, format::ProtobufUtils, - previous::encoder::{ArrayEncoder, EncodedArray}, }; use lance_core::Result; diff --git a/rust/lance-encoding/src/previous/encodings/physical/dictionary.rs b/rust/lance-encoding/src/array_encoding/physical/dictionary.rs similarity index 99% rename from rust/lance-encoding/src/previous/encodings/physical/dictionary.rs rename to rust/lance-encoding/src/array_encoding/physical/dictionary.rs index a4c01ff9ef5..08c5e5d8051 100644 --- a/rust/lance-encoding/src/previous/encodings/physical/dictionary.rs +++ b/rust/lance-encoding/src/array_encoding/physical/dictionary.rs @@ -16,18 +16,18 @@ use lance_arrow::DataTypeExt; use lance_core::{Error, Result}; use std::collections::HashMap; +use crate::array_encoding::logical::primitive::PrimitiveFieldDecoder; use crate::buffer::LanceBuffer; use crate::data::{ BlockInfo, DataBlock, DictionaryDataBlock, FixedWidthDataBlock, NullableDataBlock, VariableWidthBlock, }; +use crate::decoder::LogicalPageDecoder; use crate::format::ProtobufUtils; -use crate::previous::decoder::LogicalPageDecoder; -use crate::previous::encodings::logical::primitive::PrimitiveFieldDecoder; use crate::{ EncodingsIo, decoder::{PageScheduler, PrimitivePageDecoder}, - previous::encoder::{ArrayEncoder, EncodedArray}, + encoder::{ArrayEncoder, EncodedArray}, }; #[derive(Debug)] diff --git a/rust/lance-encoding/src/previous/encodings/physical/fixed_size_binary.rs b/rust/lance-encoding/src/array_encoding/physical/fixed_size_binary.rs similarity index 98% rename from rust/lance-encoding/src/previous/encodings/physical/fixed_size_binary.rs rename to rust/lance-encoding/src/array_encoding/physical/fixed_size_binary.rs index 696edde8c9b..b655271e3aa 100644 --- a/rust/lance-encoding/src/previous/encodings/physical/fixed_size_binary.rs +++ b/rust/lance-encoding/src/array_encoding/physical/fixed_size_binary.rs @@ -13,8 +13,8 @@ use crate::{ buffer::LanceBuffer, data::{BlockInfo, DataBlock, FixedWidthDataBlock, VariableWidthBlock}, decoder::{PageScheduler, PrimitivePageDecoder}, + encoder::{ArrayEncoder, EncodedArray}, format::ProtobufUtils, - previous::encoder::{ArrayEncoder, EncodedArray}, }; /// A scheduler for fixed size binary data @@ -173,9 +173,9 @@ mod tests { use arrow_data::ArrayData; use arrow_schema::{DataType, Field}; + use crate::array_encoding::physical::fixed_size_binary::FixedSizeBinaryDecoder; use crate::data::{DataBlock, FixedWidthDataBlock}; use crate::decoder::PrimitivePageDecoder; - use crate::previous::encodings::physical::fixed_size_binary::FixedSizeBinaryDecoder; use crate::testing::{TestCases, check_basic_random, check_round_trip_encoding_of_data}; #[test_log::test(tokio::test)] diff --git a/rust/lance-encoding/src/previous/encodings/physical/fixed_size_list.rs b/rust/lance-encoding/src/array_encoding/physical/fixed_size_list.rs similarity index 95% rename from rust/lance-encoding/src/previous/encodings/physical/fixed_size_list.rs rename to rust/lance-encoding/src/array_encoding/physical/fixed_size_list.rs index e980301d117..a0f596fd8cf 100644 --- a/rust/lance-encoding/src/previous/encodings/physical/fixed_size_list.rs +++ b/rust/lance-encoding/src/array_encoding/physical/fixed_size_list.rs @@ -12,8 +12,8 @@ use crate::{ EncodingsIo, data::{DataBlock, FixedSizeListBlock}, decoder::{PageScheduler, PrimitivePageDecoder}, + encoder::{ArrayEncoder, EncodedArray}, format::ProtobufUtils, - previous::encoder::{ArrayEncoder, EncodedArray}, }; /// A scheduler for fixed size lists of primitive values @@ -138,10 +138,7 @@ mod tests { use arrow_schema::{DataType, Field}; use lance_datagen::{ArrayGeneratorExt, RowCount, array, gen_array}; - use crate::{ - testing::{TestCases, check_basic_random, check_round_trip_encoding_of_data}, - version::LanceFileVersion, - }; + use crate::testing::{TestCases, check_basic_random, check_round_trip_encoding_of_data}; const PRIMITIVE_TYPES: &[DataType] = &[DataType::Int8, DataType::Float32, DataType::Float64]; @@ -182,7 +179,7 @@ mod tests { .with_indices(vec![0, 1, 2]) .with_indices(vec![1]) .with_indices(vec![2]) - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); check_round_trip_encoding_of_data(vec![list], &test_cases, HashMap::default()).await; } @@ -209,7 +206,7 @@ mod tests { .with_indices(vec![0, 1, 2]) .with_indices(vec![1]) .with_indices(vec![2]) - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); check_round_trip_encoding_of_data(vec![list], &test_cases, HashMap::default()).await; } @@ -260,7 +257,7 @@ mod tests { .with_range(1..3) .with_indices(vec![0, 1, 2]) .with_indices(vec![2]) - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); check_round_trip_encoding_of_data(vec![outer_list], &test_cases, HashMap::default()).await; } diff --git a/rust/lance-encoding/src/previous/encodings/physical/fsst.rs b/rust/lance-encoding/src/array_encoding/physical/fsst.rs similarity index 99% rename from rust/lance-encoding/src/previous/encodings/physical/fsst.rs rename to rust/lance-encoding/src/array_encoding/physical/fsst.rs index e9bb585ed73..fd1f11b860e 100644 --- a/rust/lance-encoding/src/previous/encodings/physical/fsst.rs +++ b/rust/lance-encoding/src/array_encoding/physical/fsst.rs @@ -14,8 +14,8 @@ use crate::{ buffer::LanceBuffer, data::{BlockInfo, DataBlock, NullableDataBlock, VariableWidthBlock}, decoder::{PageScheduler, PrimitivePageDecoder}, + encoder::{ArrayEncoder, EncodedArray}, format::ProtobufUtils, - previous::encoder::{ArrayEncoder, EncodedArray}, }; #[derive(Debug)] diff --git a/rust/lance-encoding/src/previous/encodings/physical/packed_struct.rs b/rust/lance-encoding/src/array_encoding/physical/packed_struct.rs similarity index 99% rename from rust/lance-encoding/src/previous/encodings/physical/packed_struct.rs rename to rust/lance-encoding/src/array_encoding/physical/packed_struct.rs index 0f3a5fc3832..b608071e9b0 100644 --- a/rust/lance-encoding/src/previous/encodings/physical/packed_struct.rs +++ b/rust/lance-encoding/src/array_encoding/physical/packed_struct.rs @@ -18,7 +18,7 @@ use crate::{ buffer::LanceBuffer, data::{DataBlock, FixedWidthDataBlock, StructDataBlock}, decoder::{PageScheduler, PrimitivePageDecoder}, - previous::encoder::{ArrayEncoder, EncodedArray}, + encoder::{ArrayEncoder, EncodedArray}, }; #[derive(Debug)] diff --git a/rust/lance-encoding/src/previous/encodings/physical/value.rs b/rust/lance-encoding/src/array_encoding/physical/value.rs similarity index 99% rename from rust/lance-encoding/src/previous/encodings/physical/value.rs rename to rust/lance-encoding/src/array_encoding/physical/value.rs index 92ec3240a14..9f594a00252 100644 --- a/rust/lance-encoding/src/previous/encodings/physical/value.rs +++ b/rust/lance-encoding/src/array_encoding/physical/value.rs @@ -18,7 +18,7 @@ use crate::format::ProtobufUtils; use crate::{ EncodingsIo, decoder::{PageScheduler, PrimitivePageDecoder}, - previous::encoder::{ArrayEncoder, EncodedArray}, + encoder::{ArrayEncoder, EncodedArray}, }; use lance_core::{Error, Result}; diff --git a/rust/lance-encoding/src/previous/encoder.rs b/rust/lance-encoding/src/array_encoding/strategy.rs similarity index 64% rename from rust/lance-encoding/src/previous/encoder.rs rename to rust/lance-encoding/src/array_encoding/strategy.rs index 9c314ae97e9..2eee3eb4e4a 100644 --- a/rust/lance-encoding/src/previous/encoder.rs +++ b/rust/lance-encoding/src/array_encoding/strategy.rs @@ -3,23 +3,14 @@ use std::{collections::HashMap, env, hash::RandomState, sync::Arc}; -use arrow_array::{ArrayRef, UInt8Array, cast::AsArray}; +#[cfg(test)] +use arrow_array::cast::AsArray; +use arrow_array::{ArrayRef, UInt8Array}; use arrow_schema::DataType; use hyperloglogplus::{HyperLogLog, HyperLogLogPlus}; use crate::{ - buffer::LanceBuffer, - data::DataBlock, - encoder::{ColumnIndexSequence, EncodingOptions, FieldEncoder, FieldEncodingStrategy}, - encodings::{ - logical::r#struct::StructFieldEncoder, - physical::{ - block::{CompressionConfig, CompressionScheme}, - value::ValueEncoder, - }, - }, - format::pb, - previous::encodings::{ + array_encoding::{ logical::{ blob::BlobFieldEncoder, list::ListFieldEncoder, primitive::PrimitiveFieldEncoder, }, @@ -27,98 +18,49 @@ use crate::{ basic::BasicEncoder, binary::BinaryEncoder, dictionary::{AlreadyDictionaryEncoder, DictionaryEncoder}, - fixed_size_binary::FixedSizeBinaryEncoder, fixed_size_list::FslEncoder, fsst::FsstArrayEncoder, packed_struct::PackedStructEncoder, }, }, - version::LanceFileVersion, -}; - -#[cfg(feature = "bitpacking")] -use crate::previous::encodings::physical::bitpack::{ - BitpackedForNonNegArrayEncoder, compute_compressed_bit_width_for_non_neg, -}; - -use crate::constants::{ - COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, PACKED_STRUCT_LEGACY_META_KEY, - PACKED_STRUCT_META_KEY, + constants::{ + COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, PACKED_STRUCT_LEGACY_META_KEY, + PACKED_STRUCT_META_KEY, + }, + encoder::{ + ArrayEncoder, ArrayEncodingStrategy, ColumnIndexSequence, FieldEncoder, + FieldEncodingContext, FieldEncodingStrategy, + }, + encodings::{ + logical::r#struct::StructFieldEncoder, + physical::{ + block::{CompressionConfig, CompressionScheme}, + value::ValueEncoder, + }, + }, }; use lance_arrow::BLOB_META_KEY; use lance_core::datatypes::{BLOB_DESC_FIELD, Field}; use lance_core::{Error, Result}; -/// An encoded array -/// -/// Maps to a single Arrow array -/// -/// This contains the encoded data as well as a description of the encoding that was applied which -/// can be used to decode the data later. +/// Field-to-column composition for the `pb::ArrayEncoding` grammar. #[derive(Debug)] -pub struct EncodedArray { - /// The encoded buffers - pub data: DataBlock, - /// A description of the encoding used to encode the array - pub encoding: pb::ArrayEncoding, +pub struct ArrayFieldEncodingStrategy { + array_encoding_strategy: Arc, } -impl EncodedArray { - pub fn new(data: DataBlock, encoding: pb::ArrayEncoding) -> Self { - Self { data, encoding } - } - - pub fn into_buffers(self) -> (Vec, pb::ArrayEncoding) { - let buffers = self.data.into_buffers(); - (buffers, self.encoding) - } -} - -/// Encodes data from one format to another (hopefully more compact or useful) format -/// -/// The array encoder must be Send + Sync. Encoding is always done on its own -/// thread task in the background and there could potentially be multiple encode -/// tasks running for a column at once. -pub trait ArrayEncoder: std::fmt::Debug + Send + Sync { - /// Encode data +impl ArrayFieldEncodingStrategy { + /// Create the field strategy for the `pb::ArrayEncoding` grammar. /// - /// The result should contain a description of the encoding that was chosen. - /// This can be used to decode the data later. - fn encode( - &self, - data: DataBlock, - data_type: &DataType, - buffer_index: &mut u32, - ) -> Result; -} - -/// A trait to pick which encoding strategy to use for a single page -/// of data -/// -/// Presumably, implementations will make encoding decisions based on -/// array statistics. -pub trait ArrayEncodingStrategy: Send + Sync + std::fmt::Debug { - fn create_array_encoder( - &self, - arrays: &[ArrayRef], - field: &Field, - ) -> Result>; -} - -/// The core field encoding strategy is a set of basic encodings that -/// are generally applicable in most scenarios. -#[derive(Debug)] -pub struct CoreFieldEncodingStrategy { - pub array_encoding_strategy: Arc, - pub version: LanceFileVersion, -} - -impl CoreFieldEncodingStrategy { - pub fn new(version: LanceFileVersion) -> Self { + /// ``` + /// use lance_encoding::array_encoding::ArrayFieldEncodingStrategy; + /// + /// let strategy = ArrayFieldEncodingStrategy::new(); + /// ``` + pub fn new() -> Self { Self { - array_encoding_strategy: Arc::new(CoreArrayEncodingStrategy::new(version)), - version, + array_encoding_strategy: Arc::new(ArrayStrategy), } } @@ -157,14 +99,20 @@ impl CoreFieldEncodingStrategy { } } -impl FieldEncodingStrategy for CoreFieldEncodingStrategy { +impl Default for ArrayFieldEncodingStrategy { + fn default() -> Self { + Self::new() + } +} + +impl FieldEncodingStrategy for ArrayFieldEncodingStrategy { fn create_field_encoder( &self, - encoding_strategy_root: &dyn FieldEncodingStrategy, field: &Field, column_index: &mut ColumnIndexSequence, - options: &EncodingOptions, + context: &FieldEncodingContext<'_>, ) -> Result> { + let options = context.options; let data_type = field.data_type(); if Self::is_primitive_type(&data_type) { let column_index = column_index.next_column_index(field.id as u32); @@ -192,11 +140,10 @@ impl FieldEncodingStrategy for CoreFieldEncodingStrategy { match data_type { DataType::List(_child) | DataType::LargeList(_child) => { let list_idx = column_index.next_column_index(field.id as u32); - let inner_encoding = encoding_strategy_root.create_field_encoder( - encoding_strategy_root, + let inner_encoding = context.strategy.create_field_encoder( &field.children[0], column_index, - options, + context, )?; let offsets_encoder = Arc::new(BasicEncoder::new(Box::new(ValueEncoder::default()))); @@ -227,12 +174,9 @@ impl FieldEncodingStrategy for CoreFieldEncodingStrategy { .children .iter() .map(|field| { - self.create_field_encoder( - encoding_strategy_root, - field, - column_index, - options, - ) + context + .strategy + .create_field_encoder(field, column_index, context) }) .collect::>>()?; Ok(Box::new(StructFieldEncoder::new( @@ -265,33 +209,11 @@ impl FieldEncodingStrategy for CoreFieldEncodingStrategy { } } -/// The core array encoding strategy is a set of basic encodings that -/// are generally applicable in most scenarios. +/// Page-encoding selection for the `pb::ArrayEncoding` grammar. #[derive(Debug)] -pub struct CoreArrayEncodingStrategy { - pub version: LanceFileVersion, -} - -const BINARY_DATATYPES: [DataType; 4] = [ - DataType::Binary, - DataType::LargeBinary, - DataType::Utf8, - DataType::LargeUtf8, -]; - -impl CoreArrayEncodingStrategy { - fn new(version: LanceFileVersion) -> Self { - Self { version } - } -} - -impl CoreArrayEncodingStrategy { - fn can_use_fsst(data_type: &DataType, data_size: u64, version: LanceFileVersion) -> bool { - version >= LanceFileVersion::V2_1 - && matches!(data_type, DataType::Utf8 | DataType::Binary) - && data_size > 4 * 1024 * 1024 - } +struct ArrayStrategy; +impl ArrayStrategy { fn get_field_compression(field_meta: &HashMap) -> Option { let compression = field_meta.get(COMPRESSION_META_KEY)?; let compression_scheme = compression.parse::(); @@ -308,16 +230,14 @@ impl CoreArrayEncodingStrategy { fn default_binary_encoder( arrays: &[ArrayRef], - data_type: &DataType, field_meta: Option<&HashMap>, data_size: u64, - version: LanceFileVersion, ) -> Result> { let bin_indices_encoder = - Self::choose_array_encoder(arrays, &DataType::UInt64, data_size, false, version, None)?; + Self::choose_array_encoder(arrays, &DataType::UInt64, data_size, false, None)?; if let Some(compression) = field_meta.and_then(Self::get_field_compression) { - if compression.scheme == CompressionScheme::Fsst { + if compression.scheme() == CompressionScheme::Fsst { // User requested FSST let raw_encoder = Box::new(BinaryEncoder::try_new(bin_indices_encoder, None)?); Ok(Box::new(FsstArrayEncoder::new(raw_encoder))) @@ -329,13 +249,7 @@ impl CoreArrayEncodingStrategy { )?)) } } else { - // No user-specified compression, use FSST if we can - let bin_encoder = Box::new(BinaryEncoder::try_new(bin_indices_encoder, None)?); - if Self::can_use_fsst(data_type, data_size, version) { - Ok(Box::new(FsstArrayEncoder::new(bin_encoder))) - } else { - Ok(bin_encoder) - } + Ok(Box::new(BinaryEncoder::try_new(bin_indices_encoder, None)?)) } } @@ -344,7 +258,6 @@ impl CoreArrayEncodingStrategy { data_type: &DataType, data_size: u64, use_dict_encoding: bool, - version: LanceFileVersion, field_meta: Option<&HashMap>, ) -> Result> { match data_type { @@ -355,7 +268,6 @@ impl CoreArrayEncodingStrategy { inner.data_type(), data_size, use_dict_encoding, - version, None, )?, *dimension as u32, @@ -363,10 +275,9 @@ impl CoreArrayEncodingStrategy { } DataType::Dictionary(key_type, value_type) => { let key_encoder = - Self::choose_array_encoder(arrays, key_type, data_size, false, version, None)?; - let value_encoder = Self::choose_array_encoder( - arrays, value_type, data_size, false, version, None, - )?; + Self::choose_array_encoder(arrays, key_type, data_size, false, None)?; + let value_encoder = + Self::choose_array_encoder(arrays, value_type, data_size, false, None)?; Ok(Box::new(AlreadyDictionaryEncoder::new( key_encoder, @@ -384,7 +295,6 @@ impl CoreArrayEncodingStrategy { &DataType::UInt8, data_size, false, - version, None, )?; let dict_items_encoder = Self::choose_array_encoder( @@ -392,7 +302,6 @@ impl CoreArrayEncodingStrategy { &DataType::Utf8, data_size, false, - version, None, )?; @@ -400,31 +309,8 @@ impl CoreArrayEncodingStrategy { dict_indices_encoder, dict_items_encoder, ))) - } - // The parent datatype should be binary or utf8 to use the fixed size encoding - // The variable 'data_type' is passed through recursion so comparing with it would be incorrect - else if BINARY_DATATYPES.contains(arrays[0].data_type()) { - if let Some(byte_width) = check_fixed_size_encoding(arrays, version) { - // use FixedSizeBinaryEncoder - let bytes_encoder = Self::choose_array_encoder( - arrays, - &DataType::UInt8, - data_size, - false, - version, - None, - )?; - - Ok(Box::new(BasicEncoder::new(Box::new( - FixedSizeBinaryEncoder::new(bytes_encoder, byte_width as usize), - )))) - } else { - Self::default_binary_encoder( - arrays, data_type, field_meta, data_size, version, - ) - } } else { - Self::default_binary_encoder(arrays, data_type, field_meta, data_size, version) + Self::default_binary_encoder(arrays, field_meta, data_size) } } DataType::Struct(fields) => { @@ -438,7 +324,6 @@ impl CoreArrayEncodingStrategy { inner_datatype, data_size, use_dict_encoding, - version, None, )?; inner_encoders.push(inner_encoder); @@ -446,54 +331,16 @@ impl CoreArrayEncodingStrategy { Ok(Box::new(PackedStructEncoder::new(inner_encoders))) } - DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => { - if version >= LanceFileVersion::V2_1 && arrays[0].data_type() == data_type { - #[cfg(feature = "bitpacking")] - { - let compressed_bit_width = compute_compressed_bit_width_for_non_neg(arrays); - Ok(Box::new(BitpackedForNonNegArrayEncoder::new( - compressed_bit_width as usize, - data_type.clone(), - ))) - } - #[cfg(not(feature = "bitpacking"))] - { - Ok(Box::new(BasicEncoder::new(Box::new( - ValueEncoder::default(), - )))) - } - } else { - Ok(Box::new(BasicEncoder::new(Box::new( - ValueEncoder::default(), - )))) - } - } + DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => Ok( + Box::new(BasicEncoder::new(Box::new(ValueEncoder::default()))), + ), // TODO: for signed integers, I intend to make it a cascaded encoding, a sparse array for the negative values and very wide(bit-width) values, // then a bitpacked array for the narrow(bit-width) values, I need `BitpackedForNeg` to be merged first, I am // thinking about putting this sparse array in the metadata so bitpacking remain using one page buffer only. - DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => { - if version >= LanceFileVersion::V2_1 && arrays[0].data_type() == data_type { - #[cfg(feature = "bitpacking")] - { - let compressed_bit_width = compute_compressed_bit_width_for_non_neg(arrays); - Ok(Box::new(BitpackedForNonNegArrayEncoder::new( - compressed_bit_width as usize, - data_type.clone(), - ))) - } - #[cfg(not(feature = "bitpacking"))] - { - Ok(Box::new(BasicEncoder::new(Box::new( - ValueEncoder::default(), - )))) - } - } else { - Ok(Box::new(BasicEncoder::new(Box::new( - ValueEncoder::default(), - )))) - } - } + DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Box::new( + BasicEncoder::new(Box::new(ValueEncoder::default())), + )), _ => Ok(Box::new(BasicEncoder::new(Box::new( ValueEncoder::default(), )))), @@ -539,8 +386,9 @@ fn check_dict_encoding(arrays: &[ArrayRef], threshold: u64) -> bool { true } -fn check_fixed_size_encoding(arrays: &[ArrayRef], version: LanceFileVersion) -> Option { - if version < LanceFileVersion::V2_1 || arrays.is_empty() { +#[cfg(test)] +fn check_fixed_size_encoding(arrays: &[ArrayRef]) -> Option { + if arrays.is_empty() { return None; } @@ -612,7 +460,7 @@ fn check_fixed_size_encoding(arrays: &[ArrayRef], version: LanceFileVersion) -> } } -impl ArrayEncodingStrategy for CoreArrayEncodingStrategy { +impl ArrayEncodingStrategy for ArrayStrategy { fn create_array_encoder( &self, arrays: &[ArrayRef], @@ -632,7 +480,6 @@ impl ArrayEncodingStrategy for CoreArrayEncodingStrategy { data_type, data_size, use_dict_encoding, - self.version, Some(&field.metadata), ) } @@ -640,12 +487,10 @@ impl ArrayEncodingStrategy for CoreArrayEncodingStrategy { #[cfg(test)] mod tests { - use crate::constants::{COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY}; - use crate::previous::encoder::{ - ArrayEncodingStrategy, CoreArrayEncodingStrategy, check_dict_encoding, - check_fixed_size_encoding, + use super::{ + ArrayEncodingStrategy, ArrayStrategy, check_dict_encoding, check_fixed_size_encoding, }; - use crate::version::LanceFileVersion; + use crate::constants::{COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY}; use arrow_array::{ArrayRef, StringArray}; use arrow_schema::Field; use std::collections::HashMap; @@ -691,10 +536,7 @@ mod tests { assert!(!is_dict_encoding_applicable(vec![Some("a"), Some("a")], 3)); } - fn is_fixed_size_encoding_applicable( - arrays: Vec>>, - version: LanceFileVersion, - ) -> bool { + fn is_fixed_size_encoding_applicable(arrays: Vec>>) -> bool { let mut final_arrays = Vec::new(); for arr in arrays { let arr = StringArray::from(arr); @@ -702,82 +544,78 @@ mod tests { final_arrays.push(arr); } - check_fixed_size_encoding(&final_arrays.clone(), version).is_some() + check_fixed_size_encoding(&final_arrays).is_some() } #[test] fn test_fixed_size_binary_encoding_applicable() { - assert!(!is_fixed_size_encoding_applicable( - vec![vec![]], - LanceFileVersion::V2_1 - )); - - assert!(is_fixed_size_encoding_applicable( - vec![vec![Some("a"), Some("b")]], - LanceFileVersion::V2_1 - )); - - assert!(!is_fixed_size_encoding_applicable( - vec![vec![Some("abc"), Some("de")]], - LanceFileVersion::V2_1 - )); - - assert!(is_fixed_size_encoding_applicable( - vec![vec![Some("pqr"), None]], - LanceFileVersion::V2_1 - )); - - assert!(!is_fixed_size_encoding_applicable( - vec![vec![Some("pqr"), Some("")]], - LanceFileVersion::V2_1 - )); - - assert!(!is_fixed_size_encoding_applicable( - vec![vec![Some(""), Some("")]], - LanceFileVersion::V2_1 - )); + assert!(!is_fixed_size_encoding_applicable(vec![vec![]])); + + assert!(is_fixed_size_encoding_applicable(vec![vec![ + Some("a"), + Some("b") + ]])); + + assert!(!is_fixed_size_encoding_applicable(vec![vec![ + Some("abc"), + Some("de") + ]])); + + assert!(is_fixed_size_encoding_applicable(vec![vec![ + Some("pqr"), + None + ]])); + + assert!(!is_fixed_size_encoding_applicable(vec![vec![ + Some("pqr"), + Some("") + ]])); + + assert!(!is_fixed_size_encoding_applicable(vec![vec![ + Some(""), + Some("") + ]])); } #[test] fn test_fixed_size_binary_encoding_applicable_multiple_arrays() { - assert!(is_fixed_size_encoding_applicable( - vec![vec![Some("a"), Some("b")], vec![Some("c"), Some("d")]], - LanceFileVersion::V2_1 - )); - - assert!(!is_fixed_size_encoding_applicable( - vec![vec![Some("ab"), Some("bc")], vec![Some("c"), Some("d")]], - LanceFileVersion::V2_1 - )); - - assert!(!is_fixed_size_encoding_applicable( - vec![vec![Some("ab"), None], vec![None, Some("d")]], - LanceFileVersion::V2_1 - )); - - assert!(is_fixed_size_encoding_applicable( - vec![vec![Some("a"), None], vec![None, Some("d")]], - LanceFileVersion::V2_1 - )); - - assert!(!is_fixed_size_encoding_applicable( - vec![vec![Some(""), None], vec![None, Some("")]], - LanceFileVersion::V2_1 - )); - - assert!(!is_fixed_size_encoding_applicable( - vec![vec![None, None], vec![None, None]], - LanceFileVersion::V2_1 - )); + assert!(is_fixed_size_encoding_applicable(vec![ + vec![Some("a"), Some("b")], + vec![Some("c"), Some("d")] + ])); + + assert!(!is_fixed_size_encoding_applicable(vec![ + vec![Some("ab"), Some("bc")], + vec![Some("c"), Some("d")] + ])); + + assert!(!is_fixed_size_encoding_applicable(vec![ + vec![Some("ab"), None], + vec![None, Some("d")] + ])); + + assert!(is_fixed_size_encoding_applicable(vec![ + vec![Some("a"), None], + vec![None, Some("d")] + ])); + + assert!(!is_fixed_size_encoding_applicable(vec![ + vec![Some(""), None], + vec![None, Some("")] + ])); + + assert!(!is_fixed_size_encoding_applicable(vec![ + vec![None, None], + vec![None, None] + ])); } fn verify_array_encoder( array: ArrayRef, field_meta: Option>, - version: LanceFileVersion, expected_encoder: &str, ) { - let encoding_strategy = CoreArrayEncodingStrategy { version }; + let encoding_strategy = ArrayStrategy; let mut field = Field::new("test_field", array.data_type().clone(), true); if let Some(field_meta) = field_meta { field.set_metadata(field_meta); @@ -797,7 +635,6 @@ mod tests { COMPRESSION_META_KEY.to_string(), "zstd".to_string(), )])), - LanceFileVersion::V2_1, "BinaryEncoder { indices_encoder: BasicEncoder { values_encoder: ValueEncoder }, compression_config: Some(CompressionConfig { scheme: Zstd, level: None }), buffer_compressor: Some(ZstdBufferCompressor { compression_level: 0 }) }", ); } @@ -810,7 +647,6 @@ mod tests { (COMPRESSION_META_KEY.to_string(), "zstd".to_string()), (COMPRESSION_LEVEL_META_KEY.to_string(), "22".to_string()), ])), - LanceFileVersion::V2_1, "BinaryEncoder { indices_encoder: BasicEncoder { values_encoder: ValueEncoder }, compression_config: Some(CompressionConfig { scheme: Zstd, level: Some(22) }), buffer_compressor: Some(ZstdBufferCompressor { compression_level: 22 }) }", ); } diff --git a/rust/lance-encoding/src/compression.rs b/rust/lance-encoding/src/compression.rs index 5c15d48c25a..fd9fc300e24 100644 --- a/rust/lance-encoding/src/compression.rs +++ b/rust/lance-encoding/src/compression.rs @@ -20,7 +20,7 @@ use crate::encodings::physical::bitpacking::{InlineBitpacking, OutOfLineBitpacking}; use crate::{ buffer::LanceBuffer, - compression_config::{BssMode, CompressionFieldParams, CompressionParams}, + compression_config::{BssMode, CompressionFieldParams}, constants::{ BSS_META_KEY, COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, RLE_THRESHOLD_META_KEY, }, @@ -66,7 +66,6 @@ use crate::{ pb21::{CompressiveEncoding, compressive_encoding::Compression}, }, statistics::{GetStat, Stat}, - version::LanceFileVersion, }; use arrow_array::{cast::AsArray, types::UInt64Type}; @@ -141,14 +140,6 @@ pub trait CompressionStrategy: Send + Sync + std::fmt::Debug { ) -> Result>; } -#[derive(Debug, Default, Clone)] -pub struct DefaultCompressionStrategy { - /// User-configured compression parameters - params: CompressionParams, - /// The lance file version for compatibilities. - version: LanceFileVersion, -} - fn try_bss_for_mini_block( data: &FixedWidthDataBlock, params: &CompressionFieldParams, @@ -169,12 +160,7 @@ fn try_bss_for_mini_block( None } -fn try_rle_for_mini_block( - data: &FixedWidthDataBlock, - version: LanceFileVersion, - params: &CompressionFieldParams, - use_rle_v2: bool, -) -> Option> { +fn rle_is_applicable(data: &FixedWidthDataBlock, params: &CompressionFieldParams) -> Option { let bits = data.bits_per_value; if !matches!(bits, 8 | 16 | 32 | 64) { return None; @@ -197,40 +183,59 @@ fn try_rle_for_mini_block( return None; } - let num_values = data.num_values; - let raw_bytes = (num_values as u128) * (type_size as u128); - let (run_length_width, rle_bytes) = if use_rle_v2 { - estimate_rle_width_and_size_from_data(data, Some(*MAX_MINIBLOCK_VALUES)).ok()? - } else { - ( - RunLengthWidth::U8, - estimate_rle_size_for_width_from_data( - data, - Some(*MAX_MINIBLOCK_VALUES), - RunLengthWidth::U8, - ) - .ok()?, - ) - }; + Some((data.num_values as u128) * (type_size as u128)) +} - let use_child_encodings = version.resolve() >= LanceFileVersion::V2_3; - let child_compression = if use_child_encodings { - rle_child_compression_config(params) - } else { - None - }; - let use_child_bitpacking = use_child_encodings; - let rle_encoder = || { - if use_child_encodings { - RleEncoder::with_child_encoding( - run_length_width, - child_compression, - child_compression, - use_child_bitpacking, - ) - } else { - RleEncoder::with_run_length_width(run_length_width) +fn rle_beats_raw_and_bitpacking( + data: &FixedWidthDataBlock, + encoded_bytes: u128, + raw_bytes: u128, +) -> bool { + if encoded_bytes >= raw_bytes { + return false; + } + + #[cfg(feature = "bitpacking")] + { + if let Some(bitpack_bytes) = estimate_inline_bitpacking_bytes(data).map(u128::from) + && bitpack_bytes < encoded_bytes + { + return false; } + } + true +} + +fn try_fixed_u8_rle_for_mini_block( + data: &FixedWidthDataBlock, + params: &CompressionFieldParams, +) -> Option> { + let raw_bytes = rle_is_applicable(data, params)?; + let rle_bytes = estimate_rle_size_for_width_from_data( + data, + Some(*MAX_MINIBLOCK_VALUES), + RunLengthWidth::U8, + ) + .ok()?; + rle_beats_raw_and_bitpacking(data, rle_bytes, raw_bytes) + .then(|| Box::new(RleEncoder::with_run_length_width(RunLengthWidth::U8)) as _) +} + +fn try_child_rle_for_mini_block( + data: &FixedWidthDataBlock, + params: &CompressionFieldParams, +) -> Option> { + let raw_bytes = rle_is_applicable(data, params)?; + let (run_length_width, estimated_bytes) = + estimate_rle_width_and_size_from_data(data, Some(*MAX_MINIBLOCK_VALUES)).ok()?; + let child_compression = rle_child_compression_config(params); + let encoder = || { + RleEncoder::with_child_encoding( + run_length_width, + child_compression, + child_compression, + true, + ) }; #[cfg(feature = "bitpacking")] @@ -238,23 +243,16 @@ fn try_rle_for_mini_block( #[cfg(not(feature = "bitpacking"))] let bitpack_bytes = None::; - let mut selected_rle_bytes = rle_bytes; - let should_estimate_child_size = use_child_encodings - && (child_compression.is_some() || cfg!(feature = "bitpacking")) - && (rle_bytes >= raw_bytes || bitpack_bytes.is_some_and(|bytes| bytes < rle_bytes)); - if should_estimate_child_size { - selected_rle_bytes = rle_encoder().selected_payload_size(data).ok()?; - } + let should_measure_children = (child_compression.is_some() || cfg!(feature = "bitpacking")) + && (estimated_bytes >= raw_bytes + || bitpack_bytes.is_some_and(|bytes| bytes < estimated_bytes)); + let selected_bytes = if should_measure_children { + encoder().selected_payload_size(data).ok()? + } else { + estimated_bytes + }; - if selected_rle_bytes < raw_bytes { - if let Some(bitpack_bytes) = bitpack_bytes - && bitpack_bytes < selected_rle_bytes - { - return None; - } - return Some(Box::new(rle_encoder())); - } - None + rle_beats_raw_and_bitpacking(data, selected_bytes, raw_bytes).then(|| Box::new(encoder()) as _) } fn rle_child_compression_config(params: &CompressionFieldParams) -> Option { @@ -266,16 +264,12 @@ fn rle_child_compression_config(params: &CompressionFieldParams) -> Option Result, CompressiveEncoding)>> { - if version < LanceFileVersion::V2_2 { - return Ok(None); - } - let bits = data.bits_per_value; if !matches!(bits, 8 | 16 | 32 | 64) { return Ok(None); @@ -296,14 +290,6 @@ fn try_rle_for_block( } let raw_bytes = (data.num_values as u128) * ((bits / 8) as u128); - let (run_length_width, rle_payload_bytes) = if use_rle_v2 { - estimate_rle_width_and_size_from_data(data, None)? - } else { - ( - RunLengthWidth::U8, - estimate_rle_size_for_width_from_data(data, None, RunLengthWidth::U8)?, - ) - }; let rle_bytes = rle_payload_bytes.saturating_add(RLE_BLOCK_HEADER_BYTES); if rle_bytes >= raw_bytes { @@ -327,6 +313,28 @@ fn try_rle_for_block( Ok(Some((compressor, encoding))) } +fn try_fixed_u8_rle_for_block( + data: &FixedWidthDataBlock, + params: &CompressionFieldParams, +) -> Result, CompressiveEncoding)>> { + if !matches!(data.bits_per_value, 8 | 16 | 32 | 64) { + return Ok(None); + } + let encoded_bytes = estimate_rle_size_for_width_from_data(data, None, RunLengthWidth::U8)?; + try_rle_for_block_with_width(data, params, RunLengthWidth::U8, encoded_bytes) +} + +fn try_variable_rle_for_block( + data: &FixedWidthDataBlock, + params: &CompressionFieldParams, +) -> Result, CompressiveEncoding)>> { + if !matches!(data.bits_per_value, 8 | 16 | 32 | 64) { + return Ok(None); + } + let (width, encoded_bytes) = estimate_rle_width_and_size_from_data(data, None)?; + try_rle_for_block_with_width(data, params, width, encoded_bytes) +} + fn estimate_rle_width_and_size_from_data( data: &FixedWidthDataBlock, max_segment_values: Option, @@ -512,7 +520,6 @@ fn maybe_wrap_general_for_mini_block( } fn try_general_compression( - version: LanceFileVersion, field_params: &CompressionFieldParams, data: &DataBlock, ) -> Result, CompressionConfig)>> { @@ -523,9 +530,7 @@ fn try_general_compression( // User-requested compression (unused today but perhaps still used // in the future someday) - if let Some(compression_scheme) = &field_params.compression - && version >= LanceFileVersion::V2_2 - { + if let Some(compression_scheme) = &field_params.compression { let scheme: CompressionScheme = compression_scheme.parse()?; let config = CompressionConfig::new(scheme, field_params.compression_level); let compressor = Box::new(CompressedBufferEncoder::try_new(config)?); @@ -533,9 +538,7 @@ fn try_general_compression( } // Automatic compression for large blocks - if data.data_size() > MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION - && version >= LanceFileVersion::V2_2 - { + if data.data_size() > MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION { let compressor = Box::new(CompressedBufferEncoder::default()); let config = compressor.compressor.config(); return Ok(Some((compressor, config))); @@ -544,355 +547,339 @@ fn try_general_compression( Ok(None) } -impl DefaultCompressionStrategy { - /// Create a new compression strategy with default behavior - pub fn new() -> Self { - Self::default() - } +/// Parse field-level compression metadata without applying format-specific constraints. +pub fn field_metadata_params(field: &Field) -> CompressionFieldParams { + let mut params = CompressionFieldParams::default(); - /// Create a new compression strategy with user-configured parameters - pub fn with_params(params: CompressionParams) -> Self { - Self { - params, - version: LanceFileVersion::default(), - } + if let Some(compression) = field.metadata.get(COMPRESSION_META_KEY) { + params.compression = Some(compression.clone()); } - - /// Override the file version used to make compression decisions - pub fn with_version(mut self, version: LanceFileVersion) -> Self { - self.version = version; - self + if let Some(level) = field.metadata.get(COMPRESSION_LEVEL_META_KEY) { + params.compression_level = level.parse().ok(); } - - fn use_rle_v2(&self) -> bool { - self.version.resolve() >= LanceFileVersion::V2_3 + if let Some(threshold) = field.metadata.get(RLE_THRESHOLD_META_KEY) { + params.rle_threshold = threshold.parse().ok(); } - - /// Parse compression parameters from field metadata - fn parse_field_metadata(field: &Field, version: &LanceFileVersion) -> CompressionFieldParams { - let mut params = CompressionFieldParams::default(); - - // Parse compression method - if let Some(compression) = field.metadata.get(COMPRESSION_META_KEY) { - params.compression = Some(compression.clone()); - } - - // Parse compression level - if let Some(level) = field.metadata.get(COMPRESSION_LEVEL_META_KEY) { - params.compression_level = level.parse().ok(); + if let Some(bss_str) = field.metadata.get(BSS_META_KEY) { + match BssMode::parse(bss_str) { + Some(mode) => params.bss = Some(mode), + None => log::warn!("Invalid BSS mode '{}', using default", bss_str), } - - // Parse RLE threshold - if let Some(threshold) = field.metadata.get(RLE_THRESHOLD_META_KEY) { - params.rle_threshold = threshold.parse().ok(); - } - - // Parse BSS mode - if let Some(bss_str) = field.metadata.get(BSS_META_KEY) { - match BssMode::parse(bss_str) { - Some(mode) => params.bss = Some(mode), - None => { - log::warn!("Invalid BSS mode '{}', using default", bss_str); - } - } - } - - // Parse minichunk size - if let Some(minichunk_size_str) = field - .metadata - .get(super::constants::MINICHUNK_SIZE_META_KEY) - { - if let Ok(minichunk_size) = minichunk_size_str.parse::() { - // for lance v2.1, only 32kb or smaller is supported - if minichunk_size >= 32 * 1024 && *version <= LanceFileVersion::V2_1 { - log::warn!( - "minichunk_size '{}' too large for version '{}', using default", - minichunk_size, - version - ); - } else { - params.minichunk_size = Some(minichunk_size); - } - } else { - log::warn!("Invalid minichunk_size '{}', skipping", minichunk_size_str); - } - } - - params } - - fn build_fixed_width_compressor( - &self, - params: &CompressionFieldParams, - data: &FixedWidthDataBlock, - ) -> Result> { - if params.compression.as_deref() == Some("none") { - return Ok(Box::new(ValueEncoder::default())); + if let Some(minichunk_size_str) = field + .metadata + .get(super::constants::MINICHUNK_SIZE_META_KEY) + { + if let Ok(minichunk_size) = minichunk_size_str.parse::() { + params.minichunk_size = Some(minichunk_size); + } else { + log::warn!("Invalid minichunk_size '{}', skipping", minichunk_size_str); } + } - let base = try_bss_for_mini_block(data, params) - .or_else(|| try_rle_for_mini_block(data, self.version, params, self.use_rle_v2())) - .or_else(|| try_bitpack_for_mini_block(data)) - .unwrap_or_else(|| Box::new(ValueEncoder::default())); + params +} - maybe_wrap_general_for_mini_block(base, params) +/// Apply general-purpose compression requested for a fixed-width miniblock. +pub fn finalize_miniblock_compressor( + data: &DataBlock, + compressor: Box, + params: &CompressionFieldParams, +) -> Result> { + if matches!(data, DataBlock::FixedWidth(_)) { + maybe_wrap_general_for_mini_block(compressor, params) + } else { + Ok(compressor) } +} - /// Build compressor based on parameters for variable-width data - fn build_variable_width_compressor( - &self, - field: &Field, - data: &VariableWidthBlock, - ) -> Result> { - let params = self.get_merged_field_params(field); - let compression = params.compression.as_deref(); - if data.bits_per_offset != 32 && data.bits_per_offset != 64 { - return Err(Error::invalid_input(format!( - "Variable width compression not supported for {} bit offsets", - data.bits_per_offset - ))); - } +/// Honor an explicit `compression = none` request for fixed-width miniblocks. +pub fn try_uncompressed_fixed_width_miniblock( + data: &DataBlock, + params: &CompressionFieldParams, +) -> Option> { + (matches!(data, DataBlock::FixedWidth(_)) && params.compression.as_deref() == Some("none")) + .then(|| Box::new(ValueEncoder::default()) as _) +} - // Get statistics - let data_size = data.expect_single_stat::(Stat::DataSize); - let max_len = data.expect_single_stat::(Stat::MaxLength); +/// Select byte-stream-split compression for an applicable fixed-width miniblock. +pub fn try_byte_stream_split_miniblock( + data: &DataBlock, + params: &CompressionFieldParams, +) -> Option> { + let DataBlock::FixedWidth(data) = data else { + return None; + }; + try_bss_for_mini_block(data, params) +} - // Explicitly disable all compression. - if compression == Some("none") { - return Ok(Box::new(BinaryMiniBlockEncoder::new(params.minichunk_size))); - } +/// Select the original fixed-u8 RLE miniblock grammar. +pub fn try_fixed_u8_rle_miniblock( + data: &DataBlock, + params: &CompressionFieldParams, +) -> Option> { + let DataBlock::FixedWidth(data) = data else { + return None; + }; + try_fixed_u8_rle_for_mini_block(data, params) +} - let use_fsst = compression == Some("fsst") - || (compression.is_none() - && !matches!(field.data_type(), DataType::Binary | DataType::LargeBinary) - && max_len >= FSST_LEAST_INPUT_MAX_LENGTH - && data_size >= FSST_LEAST_INPUT_SIZE as u64); +/// Select variable-width RLE with independently encoded children. +pub fn try_child_rle_miniblock( + data: &DataBlock, + params: &CompressionFieldParams, +) -> Option> { + let DataBlock::FixedWidth(data) = data else { + return None; + }; + try_child_rle_for_mini_block(data, params) +} - // Choose base encoder (FSST or Binary) once. - let mut base_encoder: Box = if use_fsst { - Box::new(FsstMiniBlockEncoder::new(params.minichunk_size)) - } else { - Box::new(BinaryMiniBlockEncoder::new(params.minichunk_size)) - }; +/// Select inline bitpacking for applicable fixed-width miniblocks. +pub fn try_bitpacking_miniblock(data: &DataBlock) -> Option> { + let DataBlock::FixedWidth(data) = data else { + return None; + }; + try_bitpack_for_mini_block(data) +} - // Wrap with general compression when configured (except FSST / none). - if let Some(compression_scheme) = compression.filter(|scheme| *scheme != "fsst") { - let scheme: CompressionScheme = compression_scheme.parse()?; - let config = CompressionConfig::new(scheme, params.compression_level); - base_encoder = Box::new(GeneralMiniBlockCompressor::new(base_encoder, config)); - } +/// Store fixed-width miniblock values without a value codec. +pub fn try_raw_fixed_width_miniblock(data: &DataBlock) -> Option> { + matches!(data, DataBlock::FixedWidth(_)).then(|| Box::new(ValueEncoder::default()) as _) +} - Ok(base_encoder) +/// Encode variable-width miniblocks with binary or FSST encoding. +pub fn try_variable_width_miniblock( + field: &Field, + data: &DataBlock, + params: &CompressionFieldParams, +) -> Result>> { + let DataBlock::VariableWidth(data) = data else { + return Ok(None); + }; + if data.bits_per_offset != 32 && data.bits_per_offset != 64 { + return Err(Error::invalid_input(format!( + "Variable width compression not supported for {} bit offsets", + data.bits_per_offset + ))); } - /// Merge user-configured parameters with field metadata - /// Field metadata has highest priority - fn get_merged_field_params(&self, field: &Field) -> CompressionFieldParams { - let mut field_params = self - .params - .get_field_params(&field.name, &field.data_type()); - - // Override with field metadata if present (highest priority) - let metadata_params = Self::parse_field_metadata(field, &self.version); - field_params.merge(&metadata_params); + let compression = params.compression.as_deref(); + let data_size = data.expect_single_stat::(Stat::DataSize); + let max_len = data.expect_single_stat::(Stat::MaxLength); + if compression == Some("none") { + return Ok(Some(Box::new(BinaryMiniBlockEncoder::new( + params.minichunk_size, + )))); + } - field_params + let use_fsst = compression == Some("fsst") + || (compression.is_none() + && !matches!(field.data_type(), DataType::Binary | DataType::LargeBinary) + && max_len >= FSST_LEAST_INPUT_MAX_LENGTH + && data_size >= FSST_LEAST_INPUT_SIZE as u64); + let mut encoder: Box = if use_fsst { + Box::new(FsstMiniBlockEncoder::new(params.minichunk_size)) + } else { + Box::new(BinaryMiniBlockEncoder::new(params.minichunk_size)) + }; + if let Some(compression_scheme) = compression.filter(|scheme| *scheme != "fsst") { + let scheme: CompressionScheme = compression_scheme.parse()?; + let config = CompressionConfig::new(scheme, params.compression_level); + encoder = Box::new(GeneralMiniBlockCompressor::new(encoder, config)); } + Ok(Some(encoder)) } -impl CompressionStrategy for DefaultCompressionStrategy { - fn create_miniblock_compressor( - &self, - field: &Field, - data: &DataBlock, - ) -> Result> { - match data { - DataBlock::FixedWidth(fixed_width_data) => { - let field_params = self.get_merged_field_params(field); - self.build_fixed_width_compressor(&field_params, fixed_width_data) - } - DataBlock::VariableWidth(variable_width_data) => { - self.build_variable_width_compressor(field, variable_width_data) - } - DataBlock::Struct(struct_data_block) => { - // this condition is actually checked at `PrimitiveStructuralEncoder::do_flush`, - // just being cautious here. - if struct_data_block.has_variable_width_child() { - return Err(Error::invalid_input( - "Packed struct mini-block encoding supports only fixed-width children", - )); - } - Ok(Box::new(PackedStructFixedWidthMiniBlockEncoder::default())) - } - DataBlock::FixedSizeList(_) => { - // Ideally we would compress the list items but this creates something of a challenge. - // We don't want to break lists across chunks and we need to worry about inner validity - // layers. If we try and use a compression scheme then it is unlikely to respect these - // constraints. - // - // For now, we just don't compress. In the future, we might want to consider a more - // sophisticated approach. - Ok(Box::new(ValueEncoder::default())) - } - _ => Err(Error::not_supported_source( - format!( - "Mini-block compression not yet supported for block type {}", - data.name() - ) - .into(), - )), - } +/// Encode fixed-width packed structs as miniblocks. +pub fn try_fixed_packed_struct_miniblock( + data: &DataBlock, +) -> Result>> { + let DataBlock::Struct(data) = data else { + return Ok(None); + }; + if data.has_variable_width_child() { + return Err(Error::invalid_input( + "Packed struct mini-block encoding supports only fixed-width children", + )); } + Ok(Some(Box::new( + PackedStructFixedWidthMiniBlockEncoder::default(), + ))) +} - fn create_per_value( - &self, - field: &Field, - data: &DataBlock, - ) -> Result> { - let field_params = self.get_merged_field_params(field); - - match data { - DataBlock::FixedWidth(_) => Ok(Box::new(ValueEncoder::default())), - DataBlock::FixedSizeList(_) => Ok(Box::new(ValueEncoder::default())), - DataBlock::Struct(struct_block) => { - if field.children.len() != struct_block.children.len() { - return Err(Error::invalid_input( - "Struct field metadata does not match data block children", - )); - } - let has_variable_child = struct_block.has_variable_width_child(); - if has_variable_child { - if self.version < LanceFileVersion::V2_2 { - return Err(Error::not_supported_source("Variable packed struct encoding requires Lance file version 2.2 or later".into())); - } - Ok(Box::new(PackedStructVariablePerValueEncoder::new( - self.clone(), - field.children.clone(), - ))) - } else { - Err(Error::invalid_input( - "Packed struct per-value compression should not be used for fixed-width-only structs", - )) - } - } - DataBlock::VariableWidth(variable_width) => { - let compression = field_params.compression.as_deref(); - // Check for explicit "none" compression - if compression == Some("none") { - return Ok(Box::new(VariableEncoder::default())); - } +/// Store fixed-size-list miniblocks without a value codec. +pub fn try_raw_fixed_size_list_miniblock(data: &DataBlock) -> Option> { + matches!(data, DataBlock::FixedSizeList(_)).then(|| Box::new(ValueEncoder::default()) as _) +} - let max_len = variable_width.expect_single_stat::(Stat::MaxLength); - let data_size = variable_width.expect_single_stat::(Stat::DataSize); +/// Store fixed-width and fixed-size-list values directly in full-zip pages. +pub fn try_raw_per_value(data: &DataBlock) -> Option> { + matches!(data, DataBlock::FixedWidth(_) | DataBlock::FixedSizeList(_)) + .then(|| Box::new(ValueEncoder::default()) as _) +} - // If values are very large then use block compression on a per-value basis - // - // TODO: Could maybe use median here +fn validate_packed_struct(field: &Field, data: &DataBlock) -> Result> { + let DataBlock::Struct(data) = data else { + return Ok(None); + }; + if field.children.len() != data.children.len() { + return Err(Error::invalid_input( + "Struct field metadata does not match data block children", + )); + } + Ok(Some(data.has_variable_width_child())) +} - let per_value_requested = - compression.is_some_and(|compression| compression != "fsst"); +/// Reject variable-width packed structs while preserving the fixed-width error. +pub fn reject_packed_struct_per_value( + field: &Field, + data: &DataBlock, +) -> Result>> { + let Some(has_variable_child) = validate_packed_struct(field, data)? else { + return Ok(None); + }; + if has_variable_child { + return Err(Error::not_supported_source( + "Variable packed struct encoding is not enabled by the selected file format".into(), + )); + } + Err(Error::invalid_input( + "Packed struct per-value compression should not be used for fixed-width-only structs", + )) +} - if (max_len > 32 * 1024 || per_value_requested) - && data_size >= FSST_LEAST_INPUT_SIZE as u64 - { - if compression == Some("zstd") { - let config = CompressionConfig::new( - CompressionScheme::Zstd, - field_params.compression_level, - ); - return Ok(Box::new(CompressedBufferEncoder::try_new(config)?)); - } - return Ok(Box::new(CompressedBufferEncoder::default())); - } +/// Encode variable-width packed structs with the exact strategy recursively. +pub fn try_variable_packed_struct_per_value( + strategy: Arc, + field: &Field, + data: &DataBlock, +) -> Result>> { + let Some(has_variable_child) = validate_packed_struct(field, data)? else { + return Ok(None); + }; + if !has_variable_child { + return Err(Error::invalid_input( + "Packed struct per-value compression should not be used for fixed-width-only structs", + )); + } + Ok(Some(Box::new(PackedStructVariablePerValueEncoder::new( + strategy, + field.children.clone(), + )))) +} - if variable_width.bits_per_offset == 32 || variable_width.bits_per_offset == 64 { - let variable_compression = Box::new(VariableEncoder::default()); - let use_fsst = compression == Some("fsst") - || (compression.is_none() - && !matches!( - field.data_type(), - DataType::Binary | DataType::LargeBinary - ) - && max_len >= FSST_LEAST_INPUT_MAX_LENGTH - && data_size >= FSST_LEAST_INPUT_SIZE as u64); - - // Use FSST if explicitly requested or if data characteristics warrant it. - if use_fsst { - Ok(Box::new(FsstPerValueEncoder::new(variable_compression))) - } else { - Ok(variable_compression) - } - } else { - panic!( - "Does not support MiniBlockCompression for VariableWidth DataBlock with {} bits offsets.", - variable_width.bits_per_offset - ); - } - } - _ => unreachable!( - "Per-value compression not yet supported for block type: {}", - data.name() - ), +/// Encode variable-width values directly, with FSST or per-value compression +/// when applicable. +pub fn try_variable_width_per_value( + field: &Field, + data: &DataBlock, + params: &CompressionFieldParams, +) -> Result>> { + let DataBlock::VariableWidth(data) = data else { + return Ok(None); + }; + let compression = params.compression.as_deref(); + if compression == Some("none") { + return Ok(Some(Box::new(VariableEncoder::default()))); + } + + let max_len = data.expect_single_stat::(Stat::MaxLength); + let data_size = data.expect_single_stat::(Stat::DataSize); + let per_value_requested = compression.is_some_and(|compression| compression != "fsst"); + if (max_len > 32 * 1024 || per_value_requested) && data_size >= FSST_LEAST_INPUT_SIZE as u64 { + if compression == Some("zstd") { + let config = CompressionConfig::new(CompressionScheme::Zstd, params.compression_level); + return Ok(Some(Box::new(CompressedBufferEncoder::try_new(config)?))); } + return Ok(Some(Box::new(CompressedBufferEncoder::default()))); } - fn create_block_compressor( - &self, - field: &Field, - data: &DataBlock, - ) -> Result<(Box, CompressiveEncoding)> { - let field_params = self.get_merged_field_params(field); + if data.bits_per_offset != 32 && data.bits_per_offset != 64 { + return Err(Error::invalid_input(format!( + "Per-value compression does not support variable-width data with {}-bit offsets", + data.bits_per_offset + ))); + } + let encoder = Box::new(VariableEncoder::default()); + let use_fsst = compression == Some("fsst") + || (compression.is_none() + && !matches!(field.data_type(), DataType::Binary | DataType::LargeBinary) + && max_len >= FSST_LEAST_INPUT_MAX_LENGTH + && data_size >= FSST_LEAST_INPUT_SIZE as u64); + Ok(Some(if use_fsst { + Box::new(FsstPerValueEncoder::new(encoder)) + } else { + encoder + })) +} - match data { - DataBlock::FixedWidth(fixed_width) => { - if let Some((compressor, encoding)) = - try_rle_for_block(fixed_width, self.version, &field_params, self.use_rle_v2())? - { - return Ok((compressor, encoding)); - } - if let Some((compressor, encoding)) = try_bitpack_for_block(fixed_width) { - return Ok((compressor, encoding)); - } +/// Select fixed-u8 RLE for block compression. +pub fn try_fixed_u8_rle_block( + data: &DataBlock, + params: &CompressionFieldParams, +) -> Result, CompressiveEncoding)>> { + let DataBlock::FixedWidth(data) = data else { + return Ok(None); + }; + try_fixed_u8_rle_for_block(data, params) +} - // Try general compression (user-requested or automatic over MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION) - if let Some((compressor, config)) = - try_general_compression(self.version, &field_params, data)? - { - let encoding = ProtobufUtils21::wrapped( - config, - ProtobufUtils21::flat(fixed_width.bits_per_value, None), - )?; - return Ok((compressor, encoding)); - } +/// Select variable-width RLE for block compression. +pub fn try_variable_rle_block( + data: &DataBlock, + params: &CompressionFieldParams, +) -> Result, CompressiveEncoding)>> { + let DataBlock::FixedWidth(data) = data else { + return Ok(None); + }; + try_variable_rle_for_block(data, params) +} - let encoder = Box::new(ValueEncoder::default()); - let encoding = ProtobufUtils21::flat(fixed_width.bits_per_value, None); - Ok((encoder, encoding)) - } - DataBlock::VariableWidth(variable_width) => { - // Try general compression - if let Some((compressor, config)) = - try_general_compression(self.version, &field_params, data)? - { - let encoding = ProtobufUtils21::wrapped( - config, - ProtobufUtils21::variable( - ProtobufUtils21::flat(variable_width.bits_per_offset as u64, None), - None, - ), - )?; - return Ok((compressor, encoding)); - } +/// Select block bitpacking for applicable fixed-width values. +pub fn try_bitpacking_block( + data: &DataBlock, +) -> Option<(Box, CompressiveEncoding)> { + let DataBlock::FixedWidth(data) = data else { + return None; + }; + try_bitpack_for_block(data) +} - let encoder = Box::new(VariableEncoder::default()); - let encoding = ProtobufUtils21::variable( - ProtobufUtils21::flat(variable_width.bits_per_offset as u64, None), - None, - ); - Ok((encoder, encoding)) - } - _ => unreachable!(), - } +/// Select explicitly requested or automatic general-purpose block compression. +pub fn try_general_block( + data: &DataBlock, + params: &CompressionFieldParams, +) -> Result, CompressiveEncoding)>> { + let Some((compressor, config)) = try_general_compression(params, data)? else { + return Ok(None); + }; + let inner = match data { + DataBlock::FixedWidth(data) => ProtobufUtils21::flat(data.bits_per_value, None), + DataBlock::VariableWidth(data) => ProtobufUtils21::variable( + ProtobufUtils21::flat(data.bits_per_offset as u64, None), + None, + ), + _ => return Ok(None), + }; + Ok(Some((compressor, ProtobufUtils21::wrapped(config, inner)?))) +} + +/// Store fixed- and variable-width block values without block compression. +pub fn try_raw_block(data: &DataBlock) -> Option<(Box, CompressiveEncoding)> { + match data { + DataBlock::FixedWidth(data) => Some(( + Box::new(ValueEncoder::default()) as Box, + ProtobufUtils21::flat(data.bits_per_value, None), + )), + DataBlock::VariableWidth(data) => Some(( + Box::new(VariableEncoder::default()) as Box, + ProtobufUtils21::variable( + ProtobufUtils21::flat(data.bits_per_offset as u64, None), + None, + ), + )), + _ => None, } } @@ -1366,12 +1353,21 @@ fn compression_name(compression: &Compression) -> &'static str { mod tests { use super::*; use crate::buffer::LanceBuffer; + use crate::compression_config::CompressionParams; use crate::data::{BlockInfo, DataBlock, FixedWidthDataBlock}; use crate::statistics::ComputeStat; - use crate::testing::extract_array_encoding_chain; + use crate::testing::{TestEncoding, extract_array_encoding_chain, test_compression_strategy}; use arrow_schema::{DataType, Field as ArrowField}; use std::collections::HashMap; + fn strategy(encoding: TestEncoding, params: CompressionParams) -> Arc { + test_compression_strategy(encoding, params) + } + + fn baseline_strategy(params: CompressionParams) -> Arc { + strategy(TestEncoding::StructuralU16, params) + } + fn create_test_field(name: &str, data_type: DataType) -> Field { let arrow_field = ArrowField::new(name, data_type, true); let mut field = Field::try_from(&arrow_field).unwrap(); @@ -1560,7 +1556,7 @@ mod tests { }, ); - let strategy = DefaultCompressionStrategy::with_params(params); + let strategy = baseline_strategy(params); let field = create_test_field("user_id", DataType::Int32); // Create data with low run count for RLE @@ -1592,7 +1588,7 @@ mod tests { }, ); - let strategy = DefaultCompressionStrategy::with_params(params); + let strategy = baseline_strategy(params); let field = create_test_field("some_column", DataType::Int32); // Create data with very low run count (50 runs for 1000 values = 0.05 ratio) let data = create_fixed_width_block_with_stats(32, 1000, 50); @@ -1607,7 +1603,7 @@ mod tests { #[test] #[cfg(feature = "bitpacking")] fn test_block_bitpacks_with_zero_segment() { - let strategy = DefaultCompressionStrategy::new(); + let strategy = baseline_strategy(CompressionParams::default()); let field = create_test_field("levels", DataType::UInt16); // First 1024 zeros, then 1024 ones; max bit width is 1. @@ -1632,7 +1628,7 @@ mod tests { #[test] fn test_rle_block_accounts_for_header_before_selecting() { - let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); let field = create_test_field("small_constant", DataType::Int32); let values = vec![42i32; 2]; let mut block = FixedWidthDataBlock { @@ -1656,7 +1652,7 @@ mod tests { #[test] #[cfg(feature = "bitpacking")] fn test_rle_block_prefers_bitpacking_when_smaller() { - let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); let field = create_test_field("levels", DataType::UInt16); let mut values = Vec::with_capacity(2048); @@ -1687,7 +1683,7 @@ mod tests { #[test] #[cfg(feature = "bitpacking")] fn test_low_cardinality_prefers_bitpacking_over_rle() { - let strategy = DefaultCompressionStrategy::new(); + let strategy = baseline_strategy(CompressionParams::default()); let field = create_test_field("int_score", DataType::Int64); // Low cardinality values (3/4/5) but with moderate run count: @@ -1750,7 +1746,7 @@ mod tests { }, ); - let strategy = DefaultCompressionStrategy::with_params(params); + let strategy = baseline_strategy(params); let field = create_test_field("embeddings", DataType::Float32); let fixed_data = create_fixed_width_block(32, 1000); let variable_data = create_variable_width_block(32, 10, 32 * 1024); @@ -1785,7 +1781,7 @@ mod tests { arrow_field = arrow_field.with_metadata(metadata); let field = Field::try_from(&arrow_field).unwrap(); - let strategy = DefaultCompressionStrategy::with_params(CompressionParams::new()); + let strategy = baseline_strategy(CompressionParams::new()); // Test miniblock let fixed_data = create_fixed_width_block(32, 1000); @@ -1815,7 +1811,7 @@ mod tests { #[test] fn test_auto_fsst_disabled_for_binary_fields() { - let strategy = DefaultCompressionStrategy::new(); + let strategy = baseline_strategy(CompressionParams::default()); let field = create_test_field("bytes", DataType::Binary); let variable_data = create_fsst_candidate_variable_width_block(); @@ -1846,7 +1842,7 @@ mod tests { #[test] fn test_auto_fsst_still_enabled_for_utf8_fields() { - let strategy = DefaultCompressionStrategy::new(); + let strategy = baseline_strategy(CompressionParams::default()); let field = create_test_field("text", DataType::Utf8); let variable_data = create_fsst_candidate_variable_width_block(); @@ -1878,7 +1874,7 @@ mod tests { }, ); - let strategy = DefaultCompressionStrategy::with_params(params); + let strategy = baseline_strategy(params); let field = create_test_field("bytes", DataType::Binary); let variable_data = create_fsst_candidate_variable_width_block(); @@ -1911,7 +1907,7 @@ mod tests { ..Default::default() }, ); - let strategy = DefaultCompressionStrategy::with_params(params); + let strategy = baseline_strategy(params); let field = create_test_field("html", DataType::Utf8); let large = create_variable_width_block(32, 64, 40 * 1024); @@ -1949,12 +1945,8 @@ mod tests { }, ); - let strategy = DefaultCompressionStrategy::with_params(params); - // Get merged params - let merged = strategy - .params - .get_field_params("user_id", &DataType::Int32); + let merged = params.get_field_params("user_id", &DataType::Int32); // Column params should override type params assert_eq!(merged.rle_threshold, Some(0.2)); @@ -1962,9 +1954,7 @@ mod tests { assert_eq!(merged.compression_level, Some(6)); // Test field with only type params - let merged = strategy - .params - .get_field_params("other_field", &DataType::Int32); + let merged = params.get_field_params("other_field", &DataType::Int32); assert_eq!(merged.rle_threshold, Some(0.5)); assert_eq!(merged.compression, Some("lz4".to_string())); assert_eq!(merged.compression_level, None); @@ -1984,26 +1974,20 @@ mod tests { }, ); - let strategy = DefaultCompressionStrategy::with_params(params); - // Should match pattern - let merged = strategy - .params - .get_field_params("log_messages", &DataType::Utf8); + let merged = params.get_field_params("log_messages", &DataType::Utf8); assert_eq!(merged.compression, Some("zstd".to_string())); assert_eq!(merged.compression_level, Some(6)); // Should not match - let merged = strategy - .params - .get_field_params("messages_log", &DataType::Utf8); + let merged = params.get_field_params("messages_log", &DataType::Utf8); assert_eq!(merged.compression, None); } #[test] fn test_legacy_metadata_support() { let params = CompressionParams::new(); - let strategy = DefaultCompressionStrategy::with_params(params); + let strategy = baseline_strategy(params); // Test field with "none" compression metadata let mut metadata = HashMap::new(); @@ -2022,7 +2006,7 @@ mod tests { fn test_default_behavior() { // Empty params should fall back to default behavior let params = CompressionParams::new(); - let strategy = DefaultCompressionStrategy::with_params(params); + let strategy = baseline_strategy(params); let field = create_test_field("random_column", DataType::Int32); // Create data with high run count that won't trigger RLE (600 runs for 1000 values = 0.6 ratio) @@ -2037,7 +2021,7 @@ mod tests { #[test] fn test_field_metadata_compression() { let params = CompressionParams::new(); - let strategy = DefaultCompressionStrategy::with_params(params); + let strategy = baseline_strategy(params); // Test field with compression metadata let mut metadata = HashMap::new(); @@ -2057,7 +2041,7 @@ mod tests { #[test] fn test_field_metadata_rle_threshold() { let params = CompressionParams::new(); - let strategy = DefaultCompressionStrategy::with_params(params); + let strategy = baseline_strategy(params); // Test field with RLE threshold metadata let mut metadata = HashMap::new(); @@ -2095,7 +2079,7 @@ mod tests { data.compute_stat(); let data = DataBlock::FixedWidth(data); - let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); let (_compressed, encoding) = compressor.compress(data).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 16); @@ -2103,7 +2087,7 @@ mod tests { #[test] fn test_rle_v2_miniblock_keeps_u8_run_lengths_before_v2_3() { - for version in [LanceFileVersion::V2_1, LanceFileVersion::V2_2] { + for version in [TestEncoding::StructuralU16, TestEncoding::StructuralU32] { let mut metadata = HashMap::new(); metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "1.0".to_string()); metadata.insert(BSS_META_KEY.to_string(), "off".to_string()); @@ -2120,7 +2104,7 @@ mod tests { data.compute_stat(); let data = DataBlock::FixedWidth(data); - let strategy = DefaultCompressionStrategy::new().with_version(version); + let strategy = strategy(version, CompressionParams::default()); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); let (_compressed, encoding) = compressor.compress(data).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 8, "version={version}"); @@ -2145,7 +2129,7 @@ mod tests { data.compute_stat(); let data = DataBlock::FixedWidth(data); - let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); let debug_str = format!("{compressor:?}"); assert!(debug_str.contains("RleEncoder")); @@ -2171,7 +2155,7 @@ mod tests { data.compute_stat(); let data = DataBlock::FixedWidth(data); - let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); let (_compressed, encoding) = compressor.compress(data).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 16); @@ -2194,7 +2178,7 @@ mod tests { data.compute_stat(); let data = DataBlock::FixedWidth(data); - let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); let (_compressed, encoding) = compressor.compress(data).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 8); @@ -2203,7 +2187,7 @@ mod tests { #[test] #[cfg(any(feature = "lz4", feature = "zstd"))] fn test_rle_miniblock_released_versions_keep_flat_children_when_compression_requested() { - for version in [LanceFileVersion::V2_1, LanceFileVersion::V2_2] { + for version in [TestEncoding::StructuralU16, TestEncoding::StructuralU32] { let mut params = CompressionParams::new(); params.columns.insert( "dict_indices".to_string(), @@ -2216,7 +2200,7 @@ mod tests { ..Default::default() }, ); - let strategy = DefaultCompressionStrategy::with_params(params).with_version(version); + let strategy = strategy(version, params); let field = create_test_field("dict_indices", DataType::UInt32); let mut values = Vec::with_capacity(8192 * 4); @@ -2276,7 +2260,7 @@ mod tests { data.compute_stat(); let data = DataBlock::FixedWidth(data); - let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); let debug_str = format!("{compressor:?}"); assert!(debug_str.contains("RleEncoder")); @@ -2323,7 +2307,7 @@ mod tests { data.compute_stat(); let data = DataBlock::FixedWidth(data); - let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default()); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); let debug_str = format!("{compressor:?}"); assert!( @@ -2363,7 +2347,7 @@ mod tests { }, ); - let strategy = DefaultCompressionStrategy::with_params(params); + let strategy = baseline_strategy(params); // Field metadata should override params let mut metadata = HashMap::new(); @@ -2391,7 +2375,7 @@ mod tests { }, ); - let strategy = DefaultCompressionStrategy::with_params(params); + let strategy = baseline_strategy(params); // Field metadata provides partial override let mut metadata = HashMap::new(); @@ -2410,7 +2394,7 @@ mod tests { #[test] fn test_bss_field_metadata() { let params = CompressionParams::new(); - let strategy = DefaultCompressionStrategy::with_params(params); + let strategy = baseline_strategy(params); // Test BSS "on" mode with compression enabled (BSS requires compression to be effective) let mut metadata = HashMap::new(); @@ -2431,7 +2415,7 @@ mod tests { #[test] fn test_bss_with_compression() { let params = CompressionParams::new(); - let strategy = DefaultCompressionStrategy::with_params(params); + let strategy = baseline_strategy(params); // Test BSS with LZ4 compression let mut metadata = HashMap::new(); @@ -2464,8 +2448,7 @@ mod tests { }, ); - let mut strategy = DefaultCompressionStrategy::with_params(params); - strategy.version = LanceFileVersion::V2_2; + let strategy = strategy(TestEncoding::StructuralU32, params); let field = create_test_field("dict_values", DataType::FixedSizeBinary(3)); let data = create_fixed_width_block(24, 1024); @@ -2518,8 +2501,7 @@ mod tests { }, ); - let strategy = - DefaultCompressionStrategy::with_params(params).with_version(LanceFileVersion::V2_1); + let strategy = strategy(TestEncoding::StructuralU16, params); let field = create_test_field("dict_values", DataType::FixedSizeBinary(3)); let data = create_fixed_width_block(24, 1024); @@ -2544,8 +2526,7 @@ mod tests { }, ); - let strategy = - DefaultCompressionStrategy::with_params(params).with_version(LanceFileVersion::V2_2); + let strategy = strategy(TestEncoding::StructuralU32, params); let field = create_test_field("dict_values", DataType::FixedSizeBinary(3)); let data = create_fixed_width_block(24, 20_000); @@ -2577,8 +2558,7 @@ mod tests { block.compute_stat(); let data = DataBlock::FixedWidth(block); - let strategy = DefaultCompressionStrategy::with_params(CompressionParams::new()) - .with_version(LanceFileVersion::V2_3); + let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::new()); let (compressor, encoding) = strategy.create_block_compressor(&field, &data).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 32); @@ -2612,8 +2592,7 @@ mod tests { block.compute_stat(); let data = DataBlock::FixedWidth(block); - let strategy = DefaultCompressionStrategy::with_params(CompressionParams::new()) - .with_version(LanceFileVersion::V2_2); + let strategy = strategy(TestEncoding::StructuralU32, CompressionParams::new()); let (_compressor, encoding) = strategy.create_block_compressor(&field, &data).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 8); } @@ -2642,8 +2621,7 @@ mod tests { let data_block = DataBlock::FixedWidth(block); - let strategy = DefaultCompressionStrategy::with_params(CompressionParams::new()) - .with_version(LanceFileVersion::V2_2); + let strategy = strategy(TestEncoding::StructuralU32, CompressionParams::new()); let (compressor, _) = strategy .create_block_compressor(&field, &data_block) @@ -2677,8 +2655,7 @@ mod tests { let data_block = DataBlock::FixedWidth(block); - let strategy = DefaultCompressionStrategy::with_params(CompressionParams::new()) - .with_version(LanceFileVersion::V2_1); + let strategy = strategy(TestEncoding::StructuralU16, CompressionParams::new()); let (compressor, _) = strategy .create_block_compressor(&field, &data_block) diff --git a/rust/lance-encoding/src/decoder.rs b/rust/lance-encoding/src/decoder.rs index 108a848db64..dca8fe49dd0 100644 --- a/rust/lance-encoding/src/decoder.rs +++ b/rust/lance-encoding/src/decoder.rs @@ -240,6 +240,12 @@ use lance_core::error::LanceOptionExt; use lance_core::{ArrowResult, Error, Result}; use tracing::instrument; +use crate::array_encoding::logical::list::OffsetPageInfo; +use crate::array_encoding::logical::r#struct::{SimpleStructDecoder, SimpleStructScheduler}; +use crate::array_encoding::logical::{ + binary::BinaryFieldScheduler, blob::BlobFieldScheduler, list::ListFieldScheduler, + primitive::PrimitiveFieldScheduler, +}; use crate::compression::{DecompressionStrategy, DefaultDecompressionStrategy}; use crate::data::DataBlock; use crate::encoder::EncodedBatch; @@ -250,17 +256,72 @@ use crate::encodings::logical::primitive::StructuralPrimitiveFieldScheduler; use crate::encodings::logical::r#struct::{StructuralStructDecoder, StructuralStructScheduler}; use crate::format::pb::{self, column_encoding}; use crate::format::pb21; -use crate::previous::decoder::LogicalPageDecoder; -use crate::previous::encodings::logical::list::OffsetPageInfo; -use crate::previous::encodings::logical::r#struct::{SimpleStructDecoder, SimpleStructScheduler}; -use crate::previous::encodings::logical::{ - binary::BinaryFieldScheduler, blob::BlobFieldScheduler, list::ListFieldScheduler, - primitive::PrimitiveFieldScheduler, -}; use crate::repdef::{CompositeRepDefUnraveler, RepDefUnraveler}; -use crate::version::LanceFileVersion; use crate::{BufferScheduler, EncodingsIo}; +pub trait SchedulingJob: std::fmt::Debug { + fn schedule_next( + &mut self, + context: &mut SchedulerContext, + priority: &dyn PriorityRange, + ) -> Result; + + fn num_rows(&self) -> u64; +} + +/// Schedules the I/O needed to decode one field. +pub trait FieldScheduler: Send + Sync + std::fmt::Debug { + fn initialize<'a>( + &'a self, + filter: &'a FilterExpression, + context: &'a SchedulerContext, + ) -> BoxFuture<'a, Result<()>>; + + fn schedule_ranges<'a>( + &'a self, + ranges: &[Range], + filter: &FilterExpression, + ) -> Result>; + + fn num_rows(&self) -> u64; +} + +#[derive(Debug)] +pub struct DecoderReady { + pub decoder: Box, + pub path: VecDeque, +} + +/// Stateful decoder for one logical page. +pub trait LogicalPageDecoder: std::fmt::Debug + Send { + fn accept_child(&mut self, _child: DecoderReady) -> Result<()> { + Err(Error::internal(format!( + "The decoder {:?} does not expect children but received a child", + self + ))) + } + + fn wait_for_loaded(&'_ mut self, loaded_need: u64) -> BoxFuture<'_, Result<()>>; + + fn rows_loaded(&self) -> u64; + + fn rows_unloaded(&self) -> u64 { + self.num_rows() - self.rows_loaded() + } + + fn num_rows(&self) -> u64; + + fn rows_drained(&self) -> u64; + + fn rows_left(&self) -> u64 { + self.num_rows() - self.rows_drained() + } + + fn drain(&mut self, num_rows: u64) -> Result; + + fn data_type(&self) -> &DataType; +} + // If users are getting batches over 10MiB large then it's time to reduce the batch size const BATCH_SIZE_BYTES_WARNING: u64 = 10 * 1024 * 1024; const ENV_LANCE_STRUCTURAL_BATCH_DECODE_SPAWN_MODE: &str = @@ -288,39 +349,39 @@ fn inline_scheduling_threshold() -> u64 { }) } -/// Top-level encoding message for a page. Wraps both the -/// legacy pb::ArrayEncoding and the newer pb::PageLayout +/// Top-level encoding message for a page. Wraps both the v2.0 +/// [`pb::ArrayEncoding`] grammar and the structural [`pb21::PageLayout`] grammar. /// /// A file should only use one or the other and never both. /// 2.0 decoders can always assume this is pb::ArrayEncoding /// and 2.1+ decoders can always assume this is pb::PageLayout -#[derive(Debug)] +#[derive(Debug, Clone)] pub enum PageEncoding { - Legacy(pb::ArrayEncoding), + Array(pb::ArrayEncoding), Structural(pb21::PageLayout), } impl DeepSizeOf for PageEncoding { fn deep_size_of_children(&self, _context: &mut Context) -> usize { match self { - Self::Legacy(encoding) => encoding.encoded_len() * 4, + Self::Array(encoding) => encoding.encoded_len() * 4, Self::Structural(encoding) => encoding.encoded_len() * 4, } } } impl PageEncoding { - pub fn as_legacy(&self) -> &pb::ArrayEncoding { + pub fn as_array(&self) -> &pb::ArrayEncoding { match self { - Self::Legacy(enc) => enc, - Self::Structural(_) => panic!("Expected a legacy encoding"), + Self::Array(enc) => enc, + Self::Structural(_) => panic!("Expected an array encoding"), } } pub fn as_structural(&self) -> &pb21::PageLayout { match self { Self::Structural(enc) => enc, - Self::Legacy(_) => panic!("Expected a structural encoding"), + Self::Array(_) => panic!("Expected a structural encoding"), } } @@ -402,21 +463,21 @@ impl ColumnInfo { enum RootScheduler { Structural(Box), - Legacy(Arc), + Array(Arc), } impl RootScheduler { - fn as_legacy(&self) -> &Arc { + fn as_array(&self) -> &Arc { match self { - Self::Structural(_) => panic!("Expected a legacy scheduler"), - Self::Legacy(s) => s, + Self::Structural(_) => panic!("Expected an array scheduler"), + Self::Array(s) => s, } } fn as_structural(&self) -> &dyn StructuralFieldScheduler { match self { Self::Structural(s) => s.as_ref(), - Self::Legacy(_) => panic!("Expected a structural scheduler"), + Self::Array(_) => panic!("Expected a structural scheduler"), } } } @@ -606,14 +667,14 @@ impl CoreFieldDecoderStrategy { } } - fn is_primitive_legacy(data_type: &DataType) -> bool { + fn is_array_primitive(data_type: &DataType) -> bool { if data_type.is_primitive() { true } else { match data_type { // DataType::is_primitive doesn't consider these primitive but we do DataType::Boolean | DataType::Null | DataType::FixedSizeBinary(_) => true, - DataType::FixedSizeList(inner, _) => Self::is_primitive_legacy(inner.data_type()), + DataType::FixedSizeList(inner, _) => Self::is_array_primitive(inner.data_type()), _ => false, } } @@ -624,7 +685,7 @@ impl CoreFieldDecoderStrategy { field: &Field, column: &ColumnInfo, buffers: FileBuffers, - ) -> Result> { + ) -> Result> { Self::ensure_values_encoded(column, &field.name)?; // Primitive fields map to a single column let column_buffers = ColumnBuffers { @@ -647,7 +708,7 @@ impl CoreFieldDecoderStrategy { return Err(Error::invalid_input_source(format!("Due to schema we expected a struct column but we received a column with {} pages and right now we only support struct columns with 1 page", column_info.page_infos.len()).into())); } let encoding = &column_info.page_infos[0].encoding; - match encoding.as_legacy().array_encoding.as_ref().unwrap() { + match encoding.as_array().array_encoding.as_ref().unwrap() { pb::array_encoding::ArrayEncoding::Struct(_) => Ok(()), _ => Err(Error::invalid_input_source(format!("Expected a struct encoding because we have a struct field in the schema but got the encoding {:?}", encoding).into())), } @@ -656,7 +717,7 @@ impl CoreFieldDecoderStrategy { fn check_packed_struct(column_info: &ColumnInfo) -> bool { let encoding = &column_info.page_infos[0].encoding; matches!( - encoding.as_legacy().array_encoding.as_ref().unwrap(), + encoding.as_array().array_encoding.as_ref().unwrap(), pb::array_encoding::ArrayEncoding::PackedStruct(_) ) } @@ -667,14 +728,14 @@ impl CoreFieldDecoderStrategy { column_infos: &mut ColumnInfoIter, buffers: FileBuffers, offsets_column: &ColumnInfo, - ) -> Result> { + ) -> Result> { Self::ensure_values_encoded(offsets_column, &list_field.name)?; let offsets_column_buffers = ColumnBuffers { file_buffers: buffers, positions_and_sizes: &offsets_column.buffer_offsets_and_sizes, }; let items_scheduler = - self.create_legacy_field_scheduler(&list_field.children[0], column_infos, buffers)?; + self.create_array_field_scheduler(&list_field.children[0], column_infos, buffers)?; let (inner_infos, null_offset_adjustments): (Vec<_>, Vec<_>) = offsets_column .page_infos @@ -682,11 +743,11 @@ impl CoreFieldDecoderStrategy { .filter(|offsets_page| offsets_page.num_rows > 0) .map(|offsets_page| { if let Some(pb::array_encoding::ArrayEncoding::List(list_encoding)) = - &offsets_page.encoding.as_legacy().array_encoding + &offsets_page.encoding.as_array().array_encoding { let inner = PageInfo { buffer_offsets_and_sizes: offsets_page.buffer_offsets_and_sizes.clone(), - encoding: PageEncoding::Legacy( + encoding: PageEncoding::Array( list_encoding.offsets.as_ref().unwrap().as_ref().clone(), ), num_rows: offsets_page.num_rows, @@ -712,7 +773,7 @@ impl CoreFieldDecoderStrategy { Arc::from(inner_infos.into_boxed_slice()), offsets_column_buffers, self.validate_data, - )) as Arc; + )) as Arc; let items_field = match list_field.data_type() { DataType::List(inner) => inner, DataType::LargeList(inner) => inner, @@ -853,15 +914,15 @@ impl CoreFieldDecoderStrategy { } } - fn create_legacy_field_scheduler( + fn create_array_field_scheduler( &self, field: &Field, column_infos: &mut ColumnInfoIter, buffers: FileBuffers, - ) -> Result> { + ) -> Result> { let data_type = field.data_type(); validate_fixed_size_list_dimensions(&field.name, &data_type)?; - if Self::is_primitive_legacy(&data_type) { + if Self::is_array_primitive(&data_type) { let column_info = column_infos.expect_next()?; let scheduler = self.create_primitive_scheduler(field, column_info, buffers)?; return Ok(scheduler); @@ -876,7 +937,7 @@ impl CoreFieldDecoderStrategy { } if let Some(page_info) = column_info.page_infos.first() { if matches!( - page_info.encoding.as_legacy(), + page_info.encoding.as_array(), pb::ArrayEncoding { array_encoding: Some(pb::array_encoding::ArrayEncoding::List(..)) } @@ -920,7 +981,7 @@ impl CoreFieldDecoderStrategy { DataType::FixedSizeList(inner, _dimension) => { // A fixed size list column could either be a physical or a logical decoder // depending on the child data type. - if Self::is_primitive_legacy(inner.data_type()) { + if Self::is_array_primitive(inner.data_type()) { let primitive_col = column_infos.expect_next()?; let scheduler = self.create_primitive_scheduler(field, primitive_col, buffers)?; @@ -930,7 +991,7 @@ impl CoreFieldDecoderStrategy { } } DataType::Dictionary(_key_type, value_type) => { - if Self::is_primitive_legacy(value_type) || value_type.is_binary_like() { + if Self::is_array_primitive(value_type) || value_type.is_binary_like() { let primitive_col = column_infos.expect_next()?; let scheduler = self.create_primitive_scheduler(field, primitive_col, buffers)?; @@ -974,7 +1035,7 @@ impl CoreFieldDecoderStrategy { for field in &field.children { column_infos.next_top_level(); let field_scheduler = - self.create_legacy_field_scheduler(field, column_infos, buffers)?; + self.create_array_field_scheduler(field, column_infos, buffers)?; child_schedulers.push(Arc::from(field_scheduler)); } @@ -1003,12 +1064,12 @@ fn root_column(num_rows: u64) -> ColumnInfo { } else { u64::MAX }, - encoding: PageEncoding::Legacy(pb::ArrayEncoding { + encoding: PageEncoding::Array(pb::ArrayEncoding { array_encoding: Some(pb::array_encoding::ArrayEncoding::Struct( pb::SimpleStruct {}, )), }), - priority: 0, // not used in legacy scheduler + priority: 0, // not used by the array scheduler buffer_offsets_and_sizes: Arc::new([]), }) .collect::>(); @@ -1024,21 +1085,21 @@ fn root_column(num_rows: u64) -> ColumnInfo { pub enum RootDecoder { Structural(StructuralStructDecoder), - Legacy(SimpleStructDecoder), + Array(SimpleStructDecoder), } impl RootDecoder { pub fn into_structural(self) -> StructuralStructDecoder { match self { Self::Structural(decoder) => decoder, - Self::Legacy(_) => panic!("Expected a structural decoder"), + Self::Array(_) => panic!("Expected a structural decoder"), } } - pub fn into_legacy(self) -> SimpleStructDecoder { + pub fn into_array(self) -> SimpleStructDecoder { match self { - Self::Legacy(decoder) => decoder, - Self::Structural(_) => panic!("Expected a legacy decoder"), + Self::Array(decoder) => decoder, + Self::Structural(_) => panic!("Expected an array decoder"), } } } @@ -1104,27 +1165,27 @@ impl DecodeBatchScheduler { let mut column_iter = ColumnInfoIter::new(columns, &adjusted_column_indices); let strategy = CoreFieldDecoderStrategy::from_decoder_config(decoder_config); let root_scheduler = - strategy.create_legacy_field_scheduler(&root_field, &mut column_iter, buffers)?; + strategy.create_array_field_scheduler(&root_field, &mut column_iter, buffers)?; let context = SchedulerContext::new(io, cache.clone()); root_scheduler.initialize(filter, &context).await?; Ok(Self { - root_scheduler: RootScheduler::Legacy(root_scheduler.into()), + root_scheduler: RootScheduler::Array(root_scheduler.into()), root_fields, cache, }) } } - #[deprecated(since = "0.29.1", note = "This is for legacy 2.0 paths")] + #[deprecated(since = "0.29.1", note = "This is for v2.0 array-encoding paths")] pub fn from_scheduler( - root_scheduler: Arc, + root_scheduler: Arc, root_fields: Fields, cache: Arc, ) -> Self { Self { - root_scheduler: RootScheduler::Legacy(root_scheduler), + root_scheduler: RootScheduler::Array(root_scheduler), root_fields, cache, } @@ -1174,7 +1235,7 @@ impl DecodeBatchScheduler { } } - fn do_schedule_ranges_legacy( + fn do_schedule_ranges_array( &mut self, ranges: &[Range], filter: &FilterExpression, @@ -1185,7 +1246,7 @@ impl DecodeBatchScheduler { // tasks are scheduled at the same top level row. priority: Option>, ) { - let root_scheduler = self.root_scheduler.as_legacy(); + let root_scheduler = self.root_scheduler.as_array(); let rows_requested = ranges.iter().map(|r| r.end - r.start).sum::(); trace!( "Scheduling {} ranges across {}..{} ({} rows){}", @@ -1249,8 +1310,8 @@ impl DecodeBatchScheduler { priority: Option>, ) { match &self.root_scheduler { - RootScheduler::Legacy(_) => { - self.do_schedule_ranges_legacy(ranges, filter, io, schedule_action, priority) + RootScheduler::Array(_) => { + self.do_schedule_ranges_array(ranges, filter, io, schedule_action, priority) } RootScheduler::Structural(_) => { self.do_schedule_ranges_structural(ranges, filter, io, schedule_action) @@ -1422,7 +1483,7 @@ impl BatchDecodeStream { } } - fn accept_decoder(&mut self, decoder: crate::previous::decoder::DecoderReady) -> Result<()> { + fn accept_decoder(&mut self, decoder: DecoderReady) -> Result<()> { if decoder.path.is_empty() { // The root decoder we can ignore Ok(()) @@ -1443,7 +1504,7 @@ impl BatchDecodeStream { let scan_line = scan_line?; self.rows_scheduled = scan_line.scheduled_so_far; for message in scan_line.decoders { - self.accept_decoder(message.into_legacy())?; + self.accept_decoder(message.into_array())?; } } None => { @@ -1552,7 +1613,7 @@ impl BatchDecodeStream { // we can have a single implementation of the batch decode iterator enum RootDecoderMessage { LoadedPage(LoadedPageShard), - LegacyPage(crate::previous::decoder::DecoderReady), + ArrayPage(DecoderReady), } trait RootDecoderType { fn accept_message(&mut self, message: RootDecoderMessage) -> Result<()>; @@ -1576,10 +1637,10 @@ impl RootDecoderType for StructuralStructDecoder { } impl RootDecoderType for SimpleStructDecoder { fn accept_message(&mut self, message: RootDecoderMessage) -> Result<()> { - let RootDecoderMessage::LegacyPage(legacy_page) = message else { + let RootDecoderMessage::ArrayPage(array_page) = message else { unreachable!() }; - self.accept_child(legacy_page) + self.accept_child(array_page) } fn drain_batch(&mut self, num_rows: u64) -> Result { self.drain(num_rows) @@ -1663,7 +1724,7 @@ impl BatchDecodeIterator { // The root decoder we can ignore if !decoder_ready.path.is_empty() { self.root_decoder - .accept_message(RootDecoderMessage::LegacyPage(decoder_ready))?; + .accept_message(RootDecoderMessage::ArrayPage(decoder_ready))?; } } } @@ -2631,17 +2692,14 @@ impl SchedulerContext { VecDeque::from_iter(self.path.iter().copied()) } - #[deprecated(since = "0.29.1", note = "This is for legacy 2.0 paths")] - pub fn locate_decoder( - &mut self, - decoder: Box, - ) -> crate::previous::decoder::DecoderReady { + #[deprecated(since = "0.29.1", note = "This is for v2.0 array-encoding paths")] + pub fn locate_decoder(&mut self, decoder: Box) -> DecoderReady { trace!( "Scheduling decoder of type {:?} for {:?}", decoder.data_type(), self.path, ); - crate::previous::decoder::DecoderReady { + DecoderReady { decoder, path: self.current_path(), } @@ -2762,7 +2820,7 @@ pub enum MessageType { // decoder itself. The messages were not sent in priority order and the decoder // had to wait for I/O, figuring out the correct priority. This was a lot of // complexity. - DecoderReady(crate::previous::decoder::DecoderReady), + DecoderReady(DecoderReady), // Starting in 2.1 we use a simpler scheme where the scheduling happens in priority // order and the message is an unloaded decoder. These can be awaited, in order, and // the decoder does not have to worry about waiting for I/O. @@ -2770,7 +2828,7 @@ pub enum MessageType { } impl MessageType { - pub fn into_legacy(self) -> crate::previous::decoder::DecoderReady { + pub fn into_array(self) -> DecoderReady { match self { Self::DecoderReady(decoder) => decoder, Self::UnloadedPage(_) => { @@ -2870,13 +2928,22 @@ pub trait StructuralFieldDecoder: std::fmt::Debug + Send { #[derive(Debug, Default)] pub struct DecoderPlugins {} +/// The top-level column layout used by an in-memory encoded batch. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum EncodedBatchLayout { + /// Array pages include structural columns. + Array, + /// Structural pages include only leaf columns. + Structural, +} + /// Decodes a batch of data from an in-memory structure created by [`crate::encoder::encode_batch`] pub async fn decode_batch( batch: &EncodedBatch, filter: &FilterExpression, decoder_plugins: Arc, should_validate: bool, - version: LanceFileVersion, + layout: EncodedBatchLayout, cache: Option>, ) -> Result { // The io is synchronous so it shouldn't be possible for any async stuff to still be in progress @@ -2906,7 +2973,7 @@ pub async fn decode_batch( .await?; let (tx, rx) = unbounded_channel(); decode_scheduler.schedule_range(0..batch.num_rows, filter, tx, io_scheduler); - let is_structural = version >= LanceFileVersion::V2_1; + let is_structural = layout == EncodedBatchLayout::Structural; let mode = std::env::var(ENV_LANCE_STRUCTURAL_BATCH_DECODE_SPAWN_MODE); let spawn_structural_batch_decode_tasks = !matches!(mode.ok().as_deref(), Some("never")); let mut decode_stream = create_decode_stream( @@ -2926,7 +2993,6 @@ pub async fn decode_batch( // test coalesce indices to ranges mod tests { use super::*; - use crate::previous::decoder::{DecoderReady, LogicalPageDecoder}; use std::collections::VecDeque; #[derive(Debug)] @@ -3028,11 +3094,11 @@ mod tests { err ); - let mut legacy_columns = ColumnInfoIter::new(vec![], &[]); + let mut array_columns = ColumnInfoIter::new(vec![], &[]); let err = strategy - .create_legacy_field_scheduler( + .create_array_field_scheduler( &field, - &mut legacy_columns, + &mut array_columns, FileBuffers { positions_and_sizes: &[], }, @@ -3047,7 +3113,7 @@ mod tests { } #[tokio::test] - async fn test_legacy_stream_stops_on_load_error() { + async fn test_array_stream_stops_on_load_error() { use arrow_schema::Field as ArrowField; let rows_per_batch = 1; @@ -3081,7 +3147,7 @@ mod tests { let err = batches .next() .await - .expect("stream should emit the legacy page-load error") + .expect("stream should emit the array page-load error") .unwrap_err(); assert!( err.to_string().contains(load_error_message), @@ -3090,7 +3156,7 @@ mod tests { ); assert!( batches.next().await.is_none(), - "stream should stop after the legacy page-load error" + "stream should stop after the array page-load error" ); } @@ -3186,15 +3252,14 @@ mod tests { batch_size: u32, batch_size_bytes: Option, ) -> Vec { - use crate::encoder::{EncodingOptions, default_encoding_strategy, encode_batch}; - use crate::version::LanceFileVersion; - - let version = LanceFileVersion::V2_1; - let options = EncodingOptions { - version, - ..Default::default() + use crate::{ + encoder::{EncodingOptions, encode_batch}, + testing::{TestEncoding, test_encoding_strategy}, }; - let strategy = default_encoding_strategy(version); + + let version = TestEncoding::StructuralU16; + let options = EncodingOptions::default(); + let strategy = test_encoding_strategy(version); let schema = Schema::try_from(batch.schema().as_ref()).unwrap(); let encoded = encode_batch(batch, Arc::new(schema.clone()), strategy.as_ref(), &options) .await diff --git a/rust/lance-encoding/src/encoder.rs b/rust/lance-encoding/src/encoder.rs index cd0731fd06f..92fcb7f539c 100644 --- a/rust/lance-encoding/src/encoder.rs +++ b/rust/lance-encoding/src/encoder.rs @@ -6,44 +6,70 @@ //! Lance files are encoded using a [`FieldEncodingStrategy`] which choose //! what encoder to use for each field. //! -//! The current strategy is the [`StructuralEncodingStrategy`] which uses "structural" -//! encoding. A tree of encoders is built up for each field. The struct & list encoders -//! simply pull off the validity and offsets and collect them. Then, in the primitive leaf -//! encoder the validity, offsets, and values are accumulated in an accumulation buffer. Once -//! enough data has been collected the primitive encoder will either use a miniblock encoding -//! or a full zip encoding to create a page of data from the accumulation buffer. +//! Structural strategies build a tree of encoders for each field from the +//! version-free builders in [`structural`]. Struct and list encoders collect +//! validity and offsets; primitive leaf encoders accumulate values and emit +//! miniblock or full-zip pages. use std::{collections::HashMap, sync::Arc}; use arrow_array::{Array, ArrayRef, RecordBatch}; -use arrow_schema::DataType; use bytes::{Bytes, BytesMut}; use futures::future::BoxFuture; use lance_core::datatypes::{Field, Schema}; -use lance_core::error::LanceOptionExt; use lance_core::utils::bit::{is_pwr_two, pad_bytes_to}; use lance_core::{Error, Result}; use crate::buffer::LanceBuffer; -use crate::compression::{CompressionStrategy, DefaultCompressionStrategy}; -use crate::compression_config::CompressionParams; +use crate::data::DataBlock; use crate::decoder::PageEncoding; -use crate::encodings::logical::blob::{BlobStructuralEncoder, BlobV2StructuralEncoder}; -use crate::encodings::logical::fixed_size_list::FixedSizeListStructuralEncoder; -use crate::encodings::logical::list::ListStructuralEncoder; -use crate::encodings::logical::map::MapStructuralEncoder; -use crate::encodings::logical::primitive::PrimitiveStructuralEncoder; -use crate::encodings::logical::r#struct::StructStructuralEncoder; use crate::repdef::RepDefBuilder; -use crate::version::LanceFileVersion; use crate::{ decoder::{ColumnInfo, PageInfo}, format::pb, }; +pub mod structural; + /// The minimum alignment for a page buffer. Writers must respect this. pub const MIN_PAGE_BUFFER_ALIGNMENT: u64 = 8; +/// An array encoded with the `pb::ArrayEncoding` grammar. +#[derive(Debug)] +pub struct EncodedArray { + pub data: DataBlock, + pub encoding: pb::ArrayEncoding, +} + +impl EncodedArray { + pub fn new(data: DataBlock, encoding: pb::ArrayEncoding) -> Self { + Self { data, encoding } + } + + pub fn into_buffers(self) -> (Vec, pb::ArrayEncoding) { + (self.data.into_buffers(), self.encoding) + } +} + +/// Encodes one data block and describes it with `pb::ArrayEncoding`. +pub trait ArrayEncoder: std::fmt::Debug + Send + Sync { + fn encode( + &self, + data: DataBlock, + data_type: &arrow_schema::DataType, + buffer_index: &mut u32, + ) -> Result; +} + +/// Selects an `ArrayEncoder` for one page. +pub trait ArrayEncodingStrategy: Send + Sync + std::fmt::Debug { + fn create_array_encoder( + &self, + arrays: &[ArrayRef], + field: &Field, + ) -> Result>; +} + /// An encoded page of data /// /// Maps to a top-level array @@ -235,9 +261,6 @@ pub struct EncodingOptions { /// The encoder needs to know this so it figures the position of out-of-line /// buffers correctly pub buffer_alignment: u64, - - /// The Lance file version being written - pub version: LanceFileVersion, } impl Default for EncodingOptions { @@ -247,20 +270,10 @@ impl Default for EncodingOptions { max_page_bytes: 32 * 1024 * 1024, keep_original_array: true, buffer_alignment: 64, - version: LanceFileVersion::default(), } } } -impl EncodingOptions { - /// If true (for Lance file version 2.2+), miniblock chunk sizes are u32, - /// to allow storing larger chunks and their sizes for better compression. - /// For Lance file version 2.1, miniblock chunk sizes are u16. - pub fn support_large_chunk(&self) -> bool { - self.version >= LanceFileVersion::V2_2 - } -} - /// A trait to pick which kind of field encoding to use for a field /// /// Unlike the ArrayEncodingStrategy, the field encoding strategy is @@ -273,324 +286,23 @@ pub trait FieldEncodingStrategy: Send + Sync + std::fmt::Debug { /// The field encoder can be chosen on the data type as well as /// any metadata that is attached to the field. /// - /// The `encoding_strategy_root` is the encoder that should be - /// used to encode any inner data in struct / list / etc. fields. - /// - /// Initially it is the same as `self` and generally should be - /// forwarded to any inner encoding strategy. fn create_field_encoder( &self, - encoding_strategy_root: &dyn FieldEncodingStrategy, field: &Field, column_index: &mut ColumnIndexSequence, - options: &EncodingOptions, + context: &FieldEncodingContext<'_>, ) -> Result>; } -pub fn default_encoding_strategy(version: LanceFileVersion) -> Box { - match version.resolve() { - LanceFileVersion::Legacy => panic!(), - LanceFileVersion::V2_0 => Box::new( - crate::previous::encoder::CoreFieldEncodingStrategy::new(version), - ), - _ => Box::new(StructuralEncodingStrategy::with_version(version)), - } -} - -/// Create an encoding strategy with user-configured compression parameters -pub fn default_encoding_strategy_with_params( - version: LanceFileVersion, - params: CompressionParams, -) -> Result> { - match version.resolve() { - LanceFileVersion::Legacy | LanceFileVersion::V2_0 => Err(Error::invalid_input( - "Compression parameters are only supported in Lance file version 2.1 and later", - )), - _ => { - let compression_strategy = - Arc::new(DefaultCompressionStrategy::with_params(params).with_version(version)); - Ok(Box::new(StructuralEncodingStrategy { - compression_strategy, - version, - })) - } - } -} - -/// An encoding strategy used for 2.1+ files -#[derive(Debug)] -pub struct StructuralEncodingStrategy { - pub compression_strategy: Arc, - pub version: LanceFileVersion, -} - -// For some reason, clippy thinks we can add Default to the above derive but -// rustc doesn't agree (no default for Arc) -#[allow(clippy::derivable_impls)] -impl Default for StructuralEncodingStrategy { - fn default() -> Self { - Self { - compression_strategy: Arc::new(DefaultCompressionStrategy::new()), - version: LanceFileVersion::default(), - } - } -} - -impl StructuralEncodingStrategy { - pub fn with_version(version: LanceFileVersion) -> Self { - Self { - compression_strategy: Arc::new(DefaultCompressionStrategy::new().with_version(version)), - version, - } - } - - fn is_primitive_type(data_type: &DataType) -> bool { - match data_type { - DataType::FixedSizeList(inner, _) => Self::is_primitive_type(inner.data_type()), - _ => matches!( - data_type, - DataType::Boolean - | DataType::Date32 - | DataType::Date64 - | DataType::Decimal128(_, _) - | DataType::Decimal256(_, _) - | DataType::Duration(_) - | DataType::Float16 - | DataType::Float32 - | DataType::Float64 - | DataType::Int16 - | DataType::Int32 - | DataType::Int64 - | DataType::Int8 - | DataType::Interval(_) - | DataType::Null - | DataType::Time32(_) - | DataType::Time64(_) - | DataType::Timestamp(_, _) - | DataType::UInt16 - | DataType::UInt32 - | DataType::UInt64 - | DataType::UInt8 - | DataType::FixedSizeBinary(_) - | DataType::Binary - | DataType::LargeBinary - | DataType::Utf8 - | DataType::LargeUtf8, - ), - } - } - - fn do_create_field_encoder( - &self, - _encoding_strategy_root: &dyn FieldEncodingStrategy, - field: &Field, - column_index: &mut ColumnIndexSequence, - options: &EncodingOptions, - root_field_metadata: &HashMap, - ) -> Result> { - let data_type = field.data_type(); - - // Check if field is marked as blob - if field.is_blob() { - match data_type { - DataType::Binary | DataType::LargeBinary => { - return Ok(Box::new(BlobStructuralEncoder::new( - field, - column_index.next_column_index(field.id as u32), - options, - self.compression_strategy.clone(), - )?)); - } - DataType::Struct(_) if self.version >= LanceFileVersion::V2_2 => { - return Ok(Box::new(BlobV2StructuralEncoder::new( - field, - column_index.next_column_index(field.id as u32), - options, - self.compression_strategy.clone(), - )?)); - } - DataType::Struct(_) => { - return Err(Error::invalid_input_source( - "Blob v2 struct input requires file version >= 2.2".into(), - )); - } - _ => { - return Err(Error::invalid_input_source( - format!( - "Blob encoding only supports Binary/LargeBinary or v2 Struct, got {}", - data_type - ) - .into(), - )); - } - } - } - - if Self::is_primitive_type(&data_type) { - Ok(Box::new(PrimitiveStructuralEncoder::try_new( - options, - self.compression_strategy.clone(), - column_index.next_column_index(field.id as u32), - field.clone(), - Arc::new(root_field_metadata.clone()), - )?)) - } else { - match data_type { - DataType::List(_) | DataType::LargeList(_) => { - let child = field.children.first().expect_ok()?; - let child_encoder = self.do_create_field_encoder( - _encoding_strategy_root, - child, - column_index, - options, - root_field_metadata, - )?; - Ok(Box::new(ListStructuralEncoder::new( - options.keep_original_array, - child_encoder, - ))) - } - DataType::FixedSizeList(inner, _) - if matches!(inner.data_type(), DataType::Struct(_)) => - { - if self.version < LanceFileVersion::V2_2 { - return Err(Error::not_supported_source(format!( - "FixedSizeList is only supported in Lance file format 2.2+, current version: {}", - self.version - ) - .into())); - } - // Complex FixedSizeList needs structural encoding - let child = field.children.first().expect_ok()?; - let child_encoder = self.do_create_field_encoder( - _encoding_strategy_root, - child, - column_index, - options, - root_field_metadata, - )?; - Ok(Box::new(FixedSizeListStructuralEncoder::new( - options.keep_original_array, - child_encoder, - ))) - } - DataType::Map(_, keys_sorted) => { - // TODO: We only support keys_sorted=false for now, - // because converting a rust arrow map field to the python arrow field will - // lose the keys_sorted property. - if keys_sorted { - return Err(Error::not_supported_source(format!("Map data type is not supported with keys_sorted=true now, current value is {}", keys_sorted).into())); - } - if self.version < LanceFileVersion::V2_2 { - return Err(Error::not_supported_source(format!( - "Map data type is only supported in Lance file format 2.2+, current version: {}", - self.version - ) - .into())); - } - let entries_child = field.children.first().ok_or_else(|| { - Error::schema("Map should have an entries child".to_string()) - })?; - let DataType::Struct(struct_fields) = entries_child.data_type() else { - return Err(Error::schema( - "Map entries field must be a Struct".to_string(), - )); - }; - if struct_fields.len() < 2 { - return Err(Error::schema( - "Map entries struct must contain both key and value fields".to_string(), - )); - } - let key_field = &struct_fields[0]; - if key_field.is_nullable() { - return Err(Error::schema(format!( - "Map key field '{}' must be non-nullable according to Arrow Map specification", - key_field.name() - ))); - } - let child_encoder = self.do_create_field_encoder( - _encoding_strategy_root, - entries_child, - column_index, - options, - root_field_metadata, - )?; - Ok(Box::new(MapStructuralEncoder::new( - options.keep_original_array, - child_encoder, - ))) - } - DataType::Struct(fields) => { - if field.is_packed_struct() || fields.is_empty() { - // Both packed structs and empty structs are encoded as primitive - Ok(Box::new(PrimitiveStructuralEncoder::try_new( - options, - self.compression_strategy.clone(), - column_index.next_column_index(field.id as u32), - field.clone(), - Arc::new(root_field_metadata.clone()), - )?)) - } else { - let children_encoders = field - .children - .iter() - .map(|field| { - self.do_create_field_encoder( - _encoding_strategy_root, - field, - column_index, - options, - root_field_metadata, - ) - }) - .collect::>>()?; - Ok(Box::new(StructStructuralEncoder::new( - options.keep_original_array, - children_encoders, - ))) - } - } - DataType::Dictionary(_, value_type) => { - // A dictionary of primitive is, itself, primitive - if Self::is_primitive_type(&value_type) { - Ok(Box::new(PrimitiveStructuralEncoder::try_new( - options, - self.compression_strategy.clone(), - column_index.next_column_index(field.id as u32), - field.clone(), - Arc::new(root_field_metadata.clone()), - )?)) - } else { - // A dictionary of logical is, itself, logical and we don't support that today - // It could be possible (e.g. store indices in one column and values in remaining columns) - // but would be a significant amount of work - // - // An easier fallback implementation would be to decode-on-write and encode-on-read - Err(Error::not_supported_source(format!("cannot encode a dictionary column whose value type is a logical type ({})", value_type).into())) - } - } - _ => todo!("Implement encoding for field {}", field), - } - } - } -} - -impl FieldEncodingStrategy for StructuralEncodingStrategy { - fn create_field_encoder( - &self, - encoding_strategy_root: &dyn FieldEncodingStrategy, - field: &Field, - column_index: &mut ColumnIndexSequence, - options: &EncodingOptions, - ) -> Result> { - self.do_create_field_encoder( - encoding_strategy_root, - field, - column_index, - options, - &field.metadata, - ) - } +/// Context shared while one top-level field and all of its children are mapped +/// to concrete field encoders. +pub struct FieldEncodingContext<'a> { + /// The complete strategy composition used for recursive child fields. + pub strategy: &'a dyn FieldEncodingStrategy, + /// Runtime-only writer options. + pub options: &'a EncodingOptions, + /// Metadata inherited from the top-level field. + pub root_field_metadata: &'a HashMap, } /// A batch encoder that encodes RecordBatch objects by delegating @@ -612,12 +324,13 @@ impl BatchEncoder { .fields .iter() .map(|field| { - let encoder = strategy.create_field_encoder( + let context = FieldEncodingContext { strategy, - field, - &mut col_idx_sequence, options, - )?; + root_field_metadata: &field.metadata, + }; + let encoder = + strategy.create_field_encoder(field, &mut col_idx_sequence, &context)?; col_idx += encoder.as_ref().num_columns(); Ok(encoder) }) @@ -768,50 +481,9 @@ pub async fn encode_batch( #[cfg(test)] mod tests { use super::*; - use crate::compression_config::{CompressionFieldParams, CompressionParams}; + use crate::testing::{TestEncoding, create_test_field_encoder, test_encoding_strategy}; use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Fields as ArrowFields}; - #[test] - fn test_configured_encoding_strategy() { - // Create test parameters - let mut params = CompressionParams::new(); - params.columns.insert( - "*_id".to_string(), - CompressionFieldParams { - rle_threshold: Some(0.5), - compression: Some("lz4".to_string()), - compression_level: None, - bss: None, - minichunk_size: None, - }, - ); - - // Test with V2.1 - should succeed - let strategy = - default_encoding_strategy_with_params(LanceFileVersion::V2_1, params.clone()) - .expect("Should succeed for V2.1"); - - // Verify it's a StructuralEncodingStrategy - assert!(format!("{:?}", strategy).contains("StructuralEncodingStrategy")); - assert!(format!("{:?}", strategy).contains("DefaultCompressionStrategy")); - - // Test with V2.0 - should fail - let err = default_encoding_strategy_with_params(LanceFileVersion::V2_0, params.clone()) - .expect_err("Should fail for V2.0"); - assert!( - err.to_string() - .contains("only supported in Lance file version 2.1") - ); - - // Test with Legacy - should fail - let err = default_encoding_strategy_with_params(LanceFileVersion::Legacy, params) - .expect_err("Should fail for Legacy"); - assert!( - err.to_string() - .contains("only supported in Lance file version 2.1") - ); - } - #[test] fn test_fixed_size_list_struct_requires_v2_2() { let list_item = ArrowField::new( @@ -830,11 +502,12 @@ mod tests { ); let field = Field::try_from(&arrow_field).unwrap(); - let strategy = StructuralEncodingStrategy::with_version(LanceFileVersion::V2_1); + let strategy = test_encoding_strategy(TestEncoding::StructuralU16); let mut column_index = ColumnIndexSequence::default(); let options = EncodingOptions::default(); - let result = strategy.create_field_encoder(&strategy, &field, &mut column_index, &options); + let result = + create_test_field_encoder(strategy.as_ref(), &field, &mut column_index, &options); assert!( result.is_err(), "FixedSizeList should be rejected for file version 2.1" @@ -843,7 +516,7 @@ mod tests { assert!( err.to_string() - .contains("FixedSizeList is only supported in Lance file format 2.2+") + .contains("FixedSizeList is not enabled by the selected file format") ); } } diff --git a/rust/lance-encoding/src/encoder/structural.rs b/rust/lance-encoding/src/encoder/structural.rs new file mode 100644 index 00000000000..3d4e49f42c8 --- /dev/null +++ b/rust/lance-encoding/src/encoder/structural.rs @@ -0,0 +1,285 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Version-free structural field encoder builders. + +use std::sync::Arc; + +use arrow_schema::DataType; +use lance_core::{Error, Result, datatypes::Field, error::LanceOptionExt}; + +pub use crate::encodings::logical::primitive::PrimitivePageEncoding; + +use crate::encodings::logical::{ + blob::{BlobStructuralEncoder, BlobV2StructuralEncoder}, + fixed_size_list::FixedSizeListStructuralEncoder, + list::ListStructuralEncoder, + map::MapStructuralEncoder, + primitive::PrimitiveStructuralEncoder, + r#struct::StructStructuralEncoder, +}; + +use super::{ColumnIndexSequence, FieldEncoder, FieldEncodingContext}; + +/// Encode primitive leaves, primitive fixed-size lists, dictionaries, and +/// packed or empty structs using one concrete primitive page grammar. +#[derive(Debug, Clone)] +pub struct PrimitiveFieldEncoding { + page_encodings: Arc<[PrimitivePageEncoding]>, +} + +impl PrimitiveFieldEncoding { + /// Create a primitive field mechanism from ordered executable page behaviors. + pub fn new(page_encodings: impl IntoIterator) -> Self { + Self { + page_encodings: page_encodings.into_iter().collect(), + } + } + + fn is_primitive_type(data_type: &DataType) -> bool { + match data_type { + DataType::FixedSizeList(inner, _) => Self::is_primitive_type(inner.data_type()), + _ => matches!( + data_type, + DataType::Boolean + | DataType::Date32 + | DataType::Date64 + | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) + | DataType::Duration(_) + | DataType::Float16 + | DataType::Float32 + | DataType::Float64 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::Int8 + | DataType::Interval(_) + | DataType::Null + | DataType::Time32(_) + | DataType::Time64(_) + | DataType::Timestamp(_, _) + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::UInt8 + | DataType::FixedSizeBinary(_) + | DataType::Binary + | DataType::LargeBinary + | DataType::Utf8 + | DataType::LargeUtf8, + ), + } + } + + fn create_at( + &self, + field: Field, + column_index: u32, + context: &FieldEncodingContext<'_>, + ) -> Result> { + Ok(Box::new(PrimitiveStructuralEncoder::try_new( + context.options, + self.page_encodings.clone(), + column_index, + field, + Arc::new(context.root_field_metadata.clone()), + )?)) + } + + /// Create a primitive field encoder when this mechanism recognizes `field`. + pub fn try_create( + &self, + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, + ) -> Result>> { + if field.is_blob() { + return Ok(None); + } + + let data_type = field.data_type(); + let is_primitive = Self::is_primitive_type(&data_type); + let is_packed_or_empty_struct = matches!( + &data_type, + DataType::Struct(fields) if field.is_packed_struct() || fields.is_empty() + ); + let is_primitive_dictionary = matches!( + &data_type, + DataType::Dictionary(_, value_type) if Self::is_primitive_type(value_type) + ); + + if !is_primitive && !is_packed_or_empty_struct && !is_primitive_dictionary { + if let DataType::Dictionary(_, value_type) = data_type { + return Err(Error::not_supported_source( + format!( + "cannot encode a dictionary column whose value type is a logical type ({})", + value_type + ) + .into(), + )); + } + return Ok(None); + } + + Ok(Some(self.create_at( + field.clone(), + column_index.next_column_index(field.id as u32), + context, + )?)) + } +} + +/// Create the original binary blob descriptor when `field` matches. +pub fn try_create_binary_blob( + primitive: &PrimitiveFieldEncoding, + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, +) -> Result>> { + if !field.is_blob() || !matches!(field.data_type(), DataType::Binary | DataType::LargeBinary) { + return Ok(None); + } + let descriptor_column_index = column_index.next_column_index(field.id as u32); + Ok(Some(Box::new(BlobStructuralEncoder::new( + field, + |descriptor_field| primitive.create_at(descriptor_field, descriptor_column_index, context), + )?))) +} + +/// Create the structural blob descriptor when `field` matches. +pub fn try_create_structural_blob( + primitive: &PrimitiveFieldEncoding, + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, +) -> Result>> { + if !field.is_blob() || !matches!(field.data_type(), DataType::Struct(_)) { + return Ok(None); + } + let descriptor_column_index = column_index.next_column_index(field.id as u32); + Ok(Some(Box::new(BlobV2StructuralEncoder::new( + field, + |descriptor_field| primitive.create_at(descriptor_field, descriptor_column_index, context), + )?))) +} + +/// Create a variable-size list encoder when `field` matches. +pub fn try_create_list( + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, +) -> Result>> { + if !matches!( + field.data_type(), + DataType::List(_) | DataType::LargeList(_) + ) { + return Ok(None); + } + let child = field.children.first().expect_ok()?; + let child_encoder = context + .strategy + .create_field_encoder(child, column_index, context)?; + Ok(Some(Box::new(ListStructuralEncoder::new( + context.options.keep_original_array, + child_encoder, + )))) +} + +/// Create a fixed-size-list encoder whose child is a struct when applicable. +pub fn try_create_structural_fixed_size_list( + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, +) -> Result>> { + if !matches!( + field.data_type(), + DataType::FixedSizeList(inner, _) if matches!(inner.data_type(), DataType::Struct(_)) + ) { + return Ok(None); + } + let child = field.children.first().expect_ok()?; + let child_encoder = context + .strategy + .create_field_encoder(child, column_index, context)?; + Ok(Some(Box::new(FixedSizeListStructuralEncoder::new( + context.options.keep_original_array, + child_encoder, + )))) +} + +/// Create an Arrow map encoder when `field` matches. +pub fn try_create_map( + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, +) -> Result>> { + let DataType::Map(_, keys_sorted) = field.data_type() else { + return Ok(None); + }; + if keys_sorted { + return Err(Error::not_supported_source( + format!( + "Map data type is not supported with keys_sorted=true now, current value is {}", + keys_sorted + ) + .into(), + )); + } + let entries_child = field + .children + .first() + .ok_or_else(|| Error::schema("Map should have an entries child".to_string()))?; + let DataType::Struct(struct_fields) = entries_child.data_type() else { + return Err(Error::schema( + "Map entries field must be a Struct".to_string(), + )); + }; + if struct_fields.len() < 2 { + return Err(Error::schema( + "Map entries struct must contain both key and value fields".to_string(), + )); + } + let key_field = &struct_fields[0]; + if key_field.is_nullable() { + return Err(Error::schema(format!( + "Map key field '{}' must be non-nullable according to Arrow Map specification", + key_field.name() + ))); + } + let child_encoder = + context + .strategy + .create_field_encoder(entries_child, column_index, context)?; + Ok(Some(Box::new(MapStructuralEncoder::new( + context.options.keep_original_array, + child_encoder, + )))) +} + +/// Create a non-packed, non-empty struct encoder when `field` matches. +pub fn try_create_struct( + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, +) -> Result>> { + let DataType::Struct(fields) = field.data_type() else { + return Ok(None); + }; + if field.is_blob() || field.is_packed_struct() || fields.is_empty() { + return Ok(None); + } + let children_encoders = field + .children + .iter() + .map(|child| { + context + .strategy + .create_field_encoder(child, column_index, context) + }) + .collect::>>()?; + Ok(Some(Box::new(StructStructuralEncoder::new( + context.options.keep_original_array, + children_encoders, + )))) +} diff --git a/rust/lance-encoding/src/encodings/fuzz_tests.rs b/rust/lance-encoding/src/encodings/fuzz_tests.rs index b92bac09cca..3521d633ac3 100644 --- a/rust/lance-encoding/src/encodings/fuzz_tests.rs +++ b/rust/lance-encoding/src/encodings/fuzz_tests.rs @@ -16,7 +16,6 @@ use arrow_schema::{DataType, Field}; use proptest::prelude::*; use crate::testing::{TestCases, check_round_trip_encoding_of_data}; -use crate::version::LanceFileVersion; use lance_core::Result; use lance_datagen::{ArrayGenerator, ByteCount, Dimension, RowCount, Seed, array, gen_batch}; @@ -280,7 +279,7 @@ proptest! { } let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_1) + .with_structural_encodings() .with_batch_size(100) .with_range(0..num_rows.min(500) as u64) .with_indices(vec![0, num_rows as u64 / 2, (num_rows - 1) as u64]); @@ -301,7 +300,7 @@ async fn test_edge_cases_single_value() { let single_int32 = Arc::new(Int32Array::from(vec![42])) as Arc; let single_string = Arc::new(StringArray::from(vec!["test"])) as Arc; - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); + let test_cases = TestCases::default().with_structural_encodings(); check_round_trip_encoding_of_data(vec![single_int32], &test_cases, HashMap::new()).await; @@ -317,7 +316,7 @@ async fn test_edge_cases_all_nulls() { vec![None, None, None] as Vec> )) as Arc; - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); + let test_cases = TestCases::default().with_structural_encodings(); check_round_trip_encoding_of_data(vec![all_nulls_int32], &test_cases, HashMap::new()).await; @@ -347,7 +346,7 @@ proptest! { let list_array = Arc::new(list_builder.finish()) as Arc; let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_1) + .with_structural_encodings() .with_range(0..list_sizes.len().min(50) as u64); check_round_trip_encoding_of_data( @@ -381,7 +380,7 @@ proptest! { .expect("Failed to generate test data"); let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); check_round_trip_encoding_of_data( vec![test_data], @@ -424,7 +423,7 @@ async fn test_list_dict_empty_batch() { let list_array = Arc::new(list_builder.finish()); let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_1) + .with_structural_encodings() // Read only the empty/null lists (rows 50-99) // This batch will have 0 underlying values .with_range(50..100); @@ -539,7 +538,7 @@ async fn test_all_valid_combinations() { let test_data = generate_test_data_for_config(&config, 100, 42).expect("Failed to generate test data"); - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); + let test_cases = TestCases::default().with_structural_encodings(); check_round_trip_encoding_of_data(vec![test_data], &test_cases, HashMap::new()).await; } diff --git a/rust/lance-encoding/src/encodings/logical/blob.rs b/rust/lance-encoding/src/encodings/logical/blob.rs index d1798b1ef69..72ea44c2e64 100644 --- a/rust/lance-encoding/src/encodings/logical/blob.rs +++ b/rust/lance-encoding/src/encodings/logical/blob.rs @@ -21,7 +21,6 @@ use crate::{ constants::PACKED_STRUCT_META_KEY, decoder::PageEncoding, encoder::{EncodeTask, EncodedColumn, EncodedPage, FieldEncoder, OutOfLineBuffers}, - encodings::logical::primitive::PrimitiveStructuralEncoder, format::ProtobufUtils21, repdef::{DefinitionInterpretation, RepDefBuilder}, }; @@ -42,9 +41,7 @@ pub struct BlobStructuralEncoder { impl BlobStructuralEncoder { pub fn new( field: &Field, - column_index: u32, - options: &crate::encoder::EncodingOptions, - compression_strategy: Arc, + make_descriptor_encoder: impl FnOnce(Field) -> Result>, ) -> Result { // Create descriptor field: struct // Preserve the original field's metadata for packed struct @@ -63,13 +60,7 @@ impl BlobStructuralEncoder { )?; // Use PrimitiveStructuralEncoder to handle the descriptor - let descriptor_encoder = Box::new(PrimitiveStructuralEncoder::try_new( - options, - compression_strategy, - column_index, - descriptor_field, - Arc::new(HashMap::new()), - )?); + let descriptor_encoder = make_descriptor_encoder(descriptor_field)?; Ok(Self { descriptor_encoder, @@ -236,9 +227,7 @@ pub struct BlobV2StructuralEncoder { impl BlobV2StructuralEncoder { pub fn new( field: &Field, - column_index: u32, - options: &crate::encoder::EncodingOptions, - compression_strategy: Arc, + make_descriptor_encoder: impl FnOnce(Field) -> Result>, ) -> Result { let mut descriptor_metadata = HashMap::with_capacity(1); descriptor_metadata.insert(PACKED_STRUCT_META_KEY.to_string(), "true".to_string()); @@ -250,13 +239,7 @@ impl BlobV2StructuralEncoder { .with_metadata(descriptor_metadata), )?; - let descriptor_encoder = Box::new(PrimitiveStructuralEncoder::try_new( - options, - compression_strategy, - column_index, - descriptor_field, - Arc::new(HashMap::new()), - )?); + let descriptor_encoder = make_descriptor_encoder(descriptor_field)?; Ok(Self { descriptor_encoder }) } @@ -430,13 +413,12 @@ impl FieldEncoder for BlobV2StructuralEncoder { mod tests { use super::*; use crate::{ - compression::DefaultCompressionStrategy, encoder::{ColumnIndexSequence, EncodingOptions}, testing::{ - TestCases, check_round_trip_encoding_of_data, - check_round_trip_encoding_of_data_with_expected, + TestCases, TestEncoding, check_round_trip_encoding_of_data, + check_round_trip_encoding_of_data_with_expected, create_test_field_encoder, + test_encoding_strategy, }, - version::LanceFileVersion, }; use arrow_array::{ ArrayRef, LargeBinaryArray, StringArray, StructArray, UInt8Array, UInt32Array, UInt64Array, @@ -445,14 +427,18 @@ mod tests { #[test] fn test_blob_encoder_creation() { - let field = - Field::try_from(ArrowField::new("blob_field", DataType::LargeBinary, true)).unwrap(); + let field = Field::try_from( + ArrowField::new("blob_field", DataType::LargeBinary, true).with_metadata( + HashMap::from([(lance_arrow::BLOB_META_KEY.to_string(), "true".to_string())]), + ), + ) + .unwrap(); let mut column_index = ColumnIndexSequence::default(); - let column_idx = column_index.next_column_index(0); let options = EncodingOptions::default(); - let compression = Arc::new(DefaultCompressionStrategy::new()); + let strategy = test_encoding_strategy(TestEncoding::StructuralU16); - let encoder = BlobStructuralEncoder::new(&field, column_idx, &options, compression); + let encoder = + create_test_field_encoder(strategy.as_ref(), &field, &mut column_index, &options); assert!(encoder.is_ok()); } @@ -466,12 +452,12 @@ mod tests { ) .unwrap(); let mut column_index = ColumnIndexSequence::default(); - let column_idx = column_index.next_column_index(0); let options = EncodingOptions::default(); - let compression = Arc::new(DefaultCompressionStrategy::new()); + let strategy = test_encoding_strategy(TestEncoding::StructuralU16); let mut encoder = - BlobStructuralEncoder::new(&field, column_idx, &options, compression).unwrap(); + create_test_field_encoder(strategy.as_ref(), &field, &mut column_index, &options) + .unwrap(); // Create test data with larger blobs let large_data = vec![0u8; 1024 * 100]; // 100KB blob @@ -524,7 +510,7 @@ mod tests { // Use the standard test harness check_round_trip_encoding_of_data( vec![array], - &TestCases::default().with_max_file_version(LanceFileVersion::V2_1), + &TestCases::default().with_array_and_u16_encodings(), blob_metadata, ) .await; @@ -628,7 +614,7 @@ mod tests { check_round_trip_encoding_of_data_with_expected( vec![Arc::new(struct_array)], Some(Arc::new(expected_descriptor)), - &TestCases::default().with_min_file_version(LanceFileVersion::V2_2), + &TestCases::default().with_u32_structural_encodings(), blob_metadata, ) .await; @@ -693,7 +679,7 @@ mod tests { check_round_trip_encoding_of_data_with_expected( vec![Arc::new(struct_array)], Some(Arc::new(expected_descriptor)), - &TestCases::default().with_min_file_version(LanceFileVersion::V2_2), + &TestCases::default().with_u32_structural_encodings(), blob_metadata, ) .await; @@ -755,7 +741,7 @@ mod tests { check_round_trip_encoding_of_data_with_expected( vec![Arc::new(struct_array)], Some(Arc::new(expected_descriptor)), - &TestCases::default().with_min_file_version(LanceFileVersion::V2_2), + &TestCases::default().with_u32_structural_encodings(), blob_metadata, ) .await; @@ -817,7 +803,7 @@ mod tests { check_round_trip_encoding_of_data_with_expected( vec![Arc::new(struct_array)], Some(Arc::new(expected_descriptor)), - &TestCases::default().with_min_file_version(LanceFileVersion::V2_2), + &TestCases::default().with_u32_structural_encodings(), blob_metadata, ) .await; diff --git a/rust/lance-encoding/src/encodings/logical/fixed_size_list.rs b/rust/lance-encoding/src/encodings/logical/fixed_size_list.rs index 1308288f226..94d3702fd66 100644 --- a/rust/lance-encoding/src/encodings/logical/fixed_size_list.rs +++ b/rust/lance-encoding/src/encodings/logical/fixed_size_list.rs @@ -534,7 +534,6 @@ mod tests { STRUCTURAL_ENCODING_MINIBLOCK, }, testing::{TestCases, check_specific_random}, - version::LanceFileVersion, }; fn make_fsl_struct_type(struct_fields: Fields, dimension: i32) -> DataType { @@ -688,18 +687,17 @@ mod tests { } #[rstest] - #[case::simple(simple_struct_fields(), 2, LanceFileVersion::V2_2)] - #[case::nested_struct(nested_struct_fields(), 2, LanceFileVersion::V2_2)] - #[case::struct_with_list(struct_with_list_fields(), 2, LanceFileVersion::V2_2)] - #[case::struct_with_large_list(struct_with_large_list_fields(), 2, LanceFileVersion::V2_2)] - #[case::nested_struct_with_list(nested_struct_with_list_fields(), 2, LanceFileVersion::V2_2)] - #[case::struct_with_nested_fsl(struct_with_nested_fsl_fields(), 2, LanceFileVersion::V2_2)] - #[case::struct_with_map(struct_with_map_fields(), 2, LanceFileVersion::V2_2)] + #[case::simple(simple_struct_fields(), 2)] + #[case::nested_struct(nested_struct_fields(), 2)] + #[case::struct_with_list(struct_with_list_fields(), 2)] + #[case::struct_with_large_list(struct_with_large_list_fields(), 2)] + #[case::nested_struct_with_list(nested_struct_with_list_fields(), 2)] + #[case::struct_with_nested_fsl(struct_with_nested_fsl_fields(), 2)] + #[case::struct_with_map(struct_with_map_fields(), 2)] #[test_log::test(tokio::test)] async fn test_fsl_struct_random( #[case] struct_fields: Fields, #[case] dimension: i32, - #[case] min_version: LanceFileVersion, #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)] structural_encoding: &str, ) { @@ -710,7 +708,7 @@ mod tests { structural_encoding.into(), ); let field = Field::new("", data_type, true).with_metadata(field_metadata); - let test_cases = TestCases::basic().with_min_file_version(min_version); + let test_cases = TestCases::basic().with_u32_structural_encodings(); check_specific_random(field, test_cases).await; } diff --git a/rust/lance-encoding/src/encodings/logical/list.rs b/rust/lance-encoding/src/encodings/logical/list.rs index c9bdcc0bc87..b79b651e624 100644 --- a/rust/lance-encoding/src/encodings/logical/list.rs +++ b/rust/lance-encoding/src/encodings/logical/list.rs @@ -251,9 +251,9 @@ mod tests { use arrow_schema::{DataType, Field, Fields}; use rstest::rstest; - use crate::{ - testing::{TestCases, check_basic_random, check_round_trip_encoding_of_data}, - version::LanceFileVersion, + use crate::testing::{ + TestCases, TestEncoding, check_basic_random, check_round_trip_encoding_of_data, + create_test_field_encoder, test_encoding_strategy, }; fn make_list_type(inner_type: DataType) -> DataType { @@ -277,20 +277,16 @@ mod tests { let arrow_field = Field::new("", array.data_type().clone(), true).with_metadata(field_metadata); let lance_field = lance_core::datatypes::Field::try_from(&arrow_field).unwrap(); - let encoding_strategy = crate::encoder::default_encoding_strategy(LanceFileVersion::V2_2); + let encoding_strategy = test_encoding_strategy(TestEncoding::StructuralU32); let mut column_index_seq = crate::encoder::ColumnIndexSequence::default(); - let encoding_options = crate::encoder::EncodingOptions { - version: LanceFileVersion::V2_2, - ..Default::default() - }; - let mut encoder = encoding_strategy - .create_field_encoder( - encoding_strategy.as_ref(), - &lance_field, - &mut column_index_seq, - &encoding_options, - ) - .unwrap(); + let encoding_options = crate::encoder::EncodingOptions::default(); + let mut encoder = create_test_field_encoder( + encoding_strategy.as_ref(), + &lance_field, + &mut column_index_seq, + &encoding_options, + ) + .unwrap(); let mut external_buffers = crate::encoder::OutOfLineBuffers::new(0, crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT); let num_rows = array.len() as u64; @@ -515,7 +511,7 @@ mod tests { structural_encoding.into(), ); - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); + let test_cases = TestCases::default().with_structural_encodings(); check_round_trip_encoding_of_data(chunks, &test_cases, field_metadata).await; } @@ -625,7 +621,7 @@ mod tests { .with_range(5..7) .with_indices(vec![1, 6]) .with_indices(vec![6]) - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(outer_list)], &test_cases, field_metadata) .await; } @@ -656,7 +652,7 @@ mod tests { .with_range(1..3) .with_indices(vec![1, 3]) .with_indices(vec![2]) - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata) .await; } @@ -687,7 +683,7 @@ mod tests { .with_range(1..3) .with_indices(vec![1, 3]) .with_indices(vec![2]) - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata) .await; } @@ -719,7 +715,7 @@ mod tests { .with_range(1..2) .with_indices(vec![0]) .with_indices(vec![1]) - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata) .await; } @@ -771,7 +767,7 @@ mod tests { .with_range(1..2) .with_indices(vec![1]) .with_indices(vec![2]) - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); check_round_trip_encoding_of_data( vec![Arc::new(list_arr)], &test_cases, @@ -810,7 +806,7 @@ mod tests { .with_range(1..2) .with_indices(vec![1]) .with_indices(vec![2]) - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(list_arr)], &test_cases, field_metadata) .await; } @@ -846,7 +842,7 @@ mod tests { ); let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_1) + .with_structural_encodings() .with_page_sizes(vec![100]) .with_range(800..900); check_round_trip_encoding_of_data( @@ -1006,7 +1002,7 @@ mod tests { let test_cases = TestCases::default() .with_range(0..2) .with_indices(vec![1]) - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); check_round_trip_encoding_of_data( vec![struct_array.clone()], &test_cases, @@ -1028,7 +1024,7 @@ mod tests { outer_list_builder.append_null(); let list_array = Arc::new(outer_list_builder.finish()); - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); + let test_cases = TestCases::default().with_structural_encodings(); check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; } @@ -1093,7 +1089,7 @@ mod tests { // This should trigger the assertion failure at primitive.rs:1362 // debug_assert!(rows_avail > 0) - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); + let test_cases = TestCases::default().with_structural_encodings(); // The bug manifests when encoding this specific pattern // Expected: successful round-trip encoding @@ -1147,7 +1143,7 @@ mod tests { .with_range(0..1000) .with_range(0..num_rows as u64) .with_indices(vec![0, (step / 2) as u64, num_rows as u64 - 1]) - .with_max_file_version(LanceFileVersion::V2_2); + .with_dense_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata) .await; } @@ -1187,7 +1183,7 @@ mod tests { .with_range(0..1000) .with_range(0..num_rows as u64) .with_indices(vec![0, (step / 2) as u64, num_rows as u64 - 1]) - .with_max_file_version(LanceFileVersion::V2_2); + .with_dense_encodings(); let list_array = Arc::new(list_array) as ArrayRef; let pages = encode_v22_pages(list_array.clone()).await; assert_split_miniblock_layout(&pages, false); @@ -1223,7 +1219,7 @@ mod tests { let test_cases = TestCases::default() .with_range(0..num_rows as u64) .with_indices(vec![0, empty_prefix_rows as u64, num_rows as u64 - 1]) - .with_max_file_version(LanceFileVersion::V2_2); + .with_dense_encodings(); let list_array = Arc::new(list_array) as ArrayRef; let pages = encode_v22_pages(list_array.clone()).await; assert_split_miniblock_layout(&pages, true); @@ -1260,7 +1256,7 @@ mod tests { let test_cases = TestCases::default() .with_range(0..num_rows as u64) .with_indices(vec![0, null_prefix_rows as u64, num_rows as u64 - 1]) - .with_max_file_version(LanceFileVersion::V2_2); + .with_dense_encodings(); let list_array = Arc::new(list_array) as ArrayRef; let pages = encode_v22_pages(list_array.clone()).await; assert_split_miniblock_layout(&pages, true); @@ -1291,7 +1287,7 @@ mod tests { let test_cases = TestCases::default() .with_range(0..num_rows as u64) .with_indices(vec![0, empty_prefix_rows as u64]) - .with_max_file_version(LanceFileVersion::V2_2); + .with_dense_encodings(); let list_array = Arc::new(list_array) as ArrayRef; let pages = encode_v22_pages(list_array.clone()).await; assert_split_miniblock_layout(&pages, true); @@ -1364,7 +1360,7 @@ mod tests { let test_cases = TestCases::default() .with_range(0..1) .with_indices(vec![0]) - .with_max_file_version(LanceFileVersion::V2_2); + .with_dense_encodings(); check_round_trip_encoding_of_data(vec![outer_list], &test_cases, HashMap::new()).await; } @@ -1412,8 +1408,7 @@ mod tests { .with_range(dense_rows.saturating_sub(8)..(dense_rows + 8)) .with_range(0..total_rows) .with_indices(vec![0, dense_rows - 1, dense_rows, total_rows - 1]) - .with_min_file_version(LanceFileVersion::V2_2) - .with_max_file_version(LanceFileVersion::V2_2); + .with_encoding(TestEncoding::StructuralU32); let list_array = Arc::new(list_array) as ArrayRef; let pages = encode_v22_pages(list_array.clone()).await; assert_split_miniblock_layout(&pages, true); @@ -1438,8 +1433,7 @@ mod tests { let test_cases = TestCases::default() .with_range(0..total_rows) - .with_min_file_version(LanceFileVersion::V2_2) - .with_max_file_version(LanceFileVersion::V2_2); + .with_encoding(TestEncoding::StructuralU32); let list_array = Arc::new(list_array) as ArrayRef; let pages = try_encode_v22_pages_with_metadata(list_array.clone(), field_metadata.clone()) .await diff --git a/rust/lance-encoding/src/encodings/logical/map.rs b/rust/lance-encoding/src/encodings/logical/map.rs index be153030105..d07c267bd2a 100644 --- a/rust/lance-encoding/src/encodings/logical/map.rs +++ b/rust/lance-encoding/src/encodings/logical/map.rs @@ -246,15 +246,14 @@ mod tests { use arrow_schema::{DataType, Field, Fields}; use crate::decoder::{DecodedArray, StructuralDecodeArrayTask}; - use crate::encoder::{ColumnIndexSequence, EncodingOptions, default_encoding_strategy}; + use crate::encoder::{ColumnIndexSequence, EncodingOptions}; use crate::encodings::logical::primitive::sparse::{ SparseCountSet, SparsePositionSet, SparseStructuralLayerPlan, SparseStructuralPlan, SparseValidityMeaning, SparseValiditySet, }; use crate::repdef::{CompositeRepDefUnraveler, RepDefUnraveler}; - use crate::{ - testing::{TestCases, check_round_trip_encoding_of_data}, - version::LanceFileVersion, + use crate::testing::{ + TestCases, TestEncoding, check_round_trip_encoding_of_data, test_encoding_strategy, }; use arrow_schema::Field as ArrowField; use lance_core::datatypes::Field as LanceField; @@ -361,7 +360,7 @@ mod tests { let test_cases = TestCases::default() .with_range(0..2) - .with_min_file_version(LanceFileVersion::V2_2); + .with_u32_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new()) .await; @@ -394,7 +393,7 @@ mod tests { .with_range(0..4) .with_indices(vec![1]) .with_indices(vec![2]) - .with_min_file_version(LanceFileVersion::V2_2); + .with_u32_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new()) .await; @@ -425,7 +424,7 @@ mod tests { .with_range(0..2) .with_indices(vec![0]) .with_indices(vec![1]) - .with_min_file_version(LanceFileVersion::V2_2); + .with_u32_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new()) .await; @@ -474,7 +473,7 @@ mod tests { let test_cases = TestCases::default() .with_range(0..3) .with_indices(vec![0, 2]) - .with_min_file_version(LanceFileVersion::V2_2); + .with_u32_structural_encodings(); check_round_trip_encoding_of_data( vec![Arc::new(struct_array)], @@ -528,7 +527,7 @@ mod tests { let test_cases = TestCases::default() .with_range(0..3) - .with_min_file_version(LanceFileVersion::V2_2); + .with_u32_structural_encodings(); check_round_trip_encoding_of_data( vec![Arc::new(struct_array)], @@ -574,7 +573,7 @@ mod tests { let test_cases = TestCases::default() .with_range(0..3) .with_indices(vec![0, 2]) - .with_min_file_version(LanceFileVersion::V2_2); + .with_u32_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new()) .await; @@ -634,7 +633,7 @@ mod tests { let test_cases = TestCases::default() .with_range(0..1) - .with_min_file_version(LanceFileVersion::V2_2); + .with_u32_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(outer_map)], &test_cases, HashMap::new()) .await; @@ -664,7 +663,7 @@ mod tests { let test_cases = TestCases::default() .with_range(0..2) .with_indices(vec![0, 1]) - .with_min_file_version(LanceFileVersion::V2_2); + .with_u32_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new()) .await; @@ -691,7 +690,7 @@ mod tests { let test_cases = TestCases::default() .with_range(0..2) - .with_min_file_version(LanceFileVersion::V2_2); + .with_u32_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new()) .await; @@ -712,7 +711,7 @@ mod tests { let test_cases = TestCases::default() .with_range(0..2) - .with_min_file_version(LanceFileVersion::V2_2); + .with_u32_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(map_array)], &test_cases, HashMap::new()) .await; @@ -747,7 +746,7 @@ mod tests { let test_cases = TestCases::default() .with_range(0..3) .with_indices(vec![0, 1, 2]) - .with_min_file_version(LanceFileVersion::V2_2); + .with_u32_structural_encodings(); // This test ensures that regardless of the internal keep_original_array setting, // the end-to-end behavior produces equivalent results @@ -766,11 +765,11 @@ mod tests { let map_field = LanceField::try_from(&map_arrow_field).unwrap(); // Test encoder: Try to create encoder with V2_1 version - should fail - let encoder_strategy = default_encoding_strategy(LanceFileVersion::V2_1); + let encoder_strategy = test_encoding_strategy(TestEncoding::StructuralU16); let mut column_index = ColumnIndexSequence::default(); let options = EncodingOptions::default(); - let encoder_result = encoder_strategy.create_field_encoder( + let encoder_result = crate::testing::create_test_field_encoder( encoder_strategy.as_ref(), &map_field, &mut column_index, @@ -787,9 +786,8 @@ mod tests { let encoder_err_msg = format!("{}", encoder_err); assert!( - encoder_err_msg.contains("2.2"), - "Encoder error message should mention version 2.2, got: {}", - encoder_err_msg + encoder_err_msg.contains("not enabled by the selected file format"), + "unexpected encoder error: {encoder_err_msg}" ); assert!( encoder_err_msg.contains("Map data type"), diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 1b16ab73b80..3523dfee116 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -61,7 +61,8 @@ use crate::{ use crate::{ repdef::{ CompositeRepDefUnraveler, ControlWordIterator, ControlWordParser, DefinitionInterpretation, - MiniBlockRepDefBudget, RepDefSlicer, SerializedRepDefs, build_control_word_iterator, + MiniBlockRepDefBudget, NormalizedStructuralPlan, RepDefSlicer, SerializedRepDefs, + build_control_word_iterator, }, utils::accumulation::AccumulationQueue, }; @@ -73,7 +74,6 @@ use crate::constants::{ DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR, DICT_VALUES_COMPRESSION_LEVEL_META_KEY, DICT_VALUES_COMPRESSION_META_KEY, }; -use crate::version::LanceFileVersion; use crate::{ EncodingsIo, buffer::LanceBuffer, @@ -4423,19 +4423,125 @@ const MINIBLOCK_ALIGNMENT: usize = 8; /// TODO: We should concatenate metadata buffers from all pages into a single buffer /// at (roughly) the end of the file so there is, at most, one read per column of /// metadata per file. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum MiniblockChunkSize { + U16, + U32, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ComplexNullEncoding { + RawLevels, + CompressedLevels, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum FixedWidthDictionaryEncoding { + Exclude64Bit, + Include64Bit, +} + +trait PrimitivePageEncodingBehavior: Send + Sync + Debug { + fn validate_field(&self, _field: &Field, _metadata: &HashMap) -> Result<()> { + Ok(()) + } + + fn try_plan_pages( + &self, + _ctx: &PrimitivePlanContext<'_>, + _arrays: &[ArrayRef], + _normalized: &NormalizedStructuralPlan, + _row_number: u64, + _num_rows: u64, + _num_values: u64, + ) -> Result>> { + Ok(None) + } + + fn try_encode_page( + &self, + _ctx: &PrimitiveEncodeContext, + page: PrimitivePageData, + ) -> Result { + Ok(PrimitiveEncodeAttempt::Unhandled(page)) + } +} + +/// One executable primitive-page behavior selected by an exact file +/// composition. +#[derive(Debug, Clone)] +pub struct PrimitivePageEncoding { + behavior: Arc, +} + +impl PrimitivePageEncoding { + /// Reject an explicit request for sparse structural encoding. + pub fn reject_sparse() -> Self { + Self { + behavior: Arc::new(RejectSparsePrimitiveEncoding), + } + } + + /// Encode constant non-null values as a constant page when applicable. + pub fn constant() -> Self { + Self { + behavior: Arc::new(ConstantPrimitiveEncoding), + } + } + + /// Plan and encode sparse structural pages when applicable. + pub fn sparse(compression: Arc) -> Self { + Self { + behavior: Arc::new(SparsePrimitiveEncoding { compression }), + } + } + + /// Encode dense pages with the original u16 miniblock grammar. + pub fn dense_u16(compression: Arc) -> Self { + Self { + behavior: Arc::new(DenseU16PrimitiveEncoding { compression }), + } + } + + /// Encode dense pages with the u32 miniblock grammar. + pub fn dense_u32(compression: Arc) -> Self { + Self { + behavior: Arc::new(DenseU32PrimitiveEncoding { compression }), + } + } +} + +#[derive(Debug)] +struct RejectSparsePrimitiveEncoding; + +#[derive(Debug)] +struct ConstantPrimitiveEncoding; + +#[derive(Debug)] +struct SparsePrimitiveEncoding { + compression: Arc, +} + +#[derive(Debug)] +struct DenseU16PrimitiveEncoding { + compression: Arc, +} + +#[derive(Debug)] +struct DenseU32PrimitiveEncoding { + compression: Arc, +} + pub struct PrimitiveStructuralEncoder { // Accumulates arrays until we have enough data to justify a disk page accumulation_queue: AccumulationQueue, keep_original_array: bool, - support_large_chunk: bool, accumulated_repdefs: Vec, - // The compression strategy we will use to compress the data - compression_strategy: Arc, + page_encodings: Arc<[PrimitivePageEncoding]>, column_index: u32, field: Field, encoding_metadata: Arc>, - version: LanceFileVersion, } struct CompressedLevelsChunk { @@ -4484,6 +4590,17 @@ struct PrimitivePageData { num_rows: u64, } +struct PrimitivePlanContext<'a> { + column_idx: u32, + field: &'a Field, + encoding_metadata: &'a HashMap, +} + +enum PrimitiveEncodeAttempt { + Encoded(EncodedPage), + Unhandled(PrimitivePageData), +} + // Immutable encoder state shared by per-page encode tasks. // // Cloning this only clones Arc-backed configuration and field metadata. Page data @@ -4492,42 +4609,24 @@ struct PrimitivePageData { struct PrimitiveEncodeContext { // Column being encoded. column_idx: u32, - // Logical field metadata for compression/layout selection. field: Field, - // Compression strategy shared across pages. - compression_strategy: Arc, - // Field-level encoding metadata such as structural encoding overrides. encoding_metadata: Arc>, - // Whether miniblock chunks may use the v2.2 large-chunk metadata. - support_large_chunk: bool, - // Lance file version selected by the writer. - version: LanceFileVersion, - // True when the only rep/def information is simple nullable validity. is_simple_validity: bool, - // True when the field has any non-empty rep/def information. has_repdef_info: bool, } impl PrimitiveStructuralEncoder { pub fn try_new( options: &EncodingOptions, - compression_strategy: Arc, + page_encodings: Arc<[PrimitivePageEncoding]>, column_index: u32, field: Field, encoding_metadata: Arc>, ) -> Result { - let requests_sparse = encoding_metadata - .get(STRUCTURAL_ENCODING_META_KEY) - .is_some_and(|requested| requested.eq_ignore_ascii_case(STRUCTURAL_ENCODING_SPARSE)); - if requests_sparse && options.version.resolve() < LanceFileVersion::V2_3 { - return Err(Error::invalid_input_source( - format!( - "Field '{}' requests sparse structural encoding, which requires Lance file format 2.3+; current version is {}", - field.name, - options.version.resolve() - ) - .into(), - )); + for page_encoding in page_encodings.iter() { + page_encoding + .behavior + .validate_field(&field, &encoding_metadata)?; } Ok(Self { accumulation_queue: AccumulationQueue::new( @@ -4535,17 +4634,35 @@ impl PrimitiveStructuralEncoder { column_index, options.keep_original_array, ), - support_large_chunk: options.support_large_chunk(), keep_original_array: options.keep_original_array, accumulated_repdefs: Vec::new(), column_index, - compression_strategy, + page_encodings, field, encoding_metadata, - version: options.version, }) } + fn encode_page( + page_encodings: &[PrimitivePageEncoding], + ctx: &PrimitiveEncodeContext, + mut page: PrimitivePageData, + ) -> Result { + for page_encoding in page_encodings { + match page_encoding.behavior.try_encode_page(ctx, page)? { + PrimitiveEncodeAttempt::Encoded(page) => return Ok(page), + PrimitiveEncodeAttempt::Unhandled(unhandled) => page = unhandled, + } + } + Err(Error::invalid_input_source( + format!( + "No primitive page encoding atom supports field '{}'", + ctx.field.name + ) + .into(), + )) + } + // TODO: This is a heuristic we may need to tune at some point // // As data gets narrow then the "zipping" process gets too expensive @@ -4640,7 +4757,7 @@ impl PrimitiveStructuralEncoder { miniblocks: MiniBlockCompressed, rep: Option>, def: Option>, - support_large_chunk: bool, + miniblock_chunk_size: MiniblockChunkSize, ) -> Result { let bytes_rep = rep .as_ref() @@ -4661,7 +4778,10 @@ impl PrimitiveStructuralEncoder { // 2 bytes for the length of each buffer and up to 7 bytes of padding per buffer let max_extra = 9 * num_buffers; let mut data_buffer = Vec::with_capacity(bytes_rep + bytes_def + bytes_data + max_extra); - let chunk_size_bytes = if support_large_chunk { 4 } else { 2 }; + let chunk_size_bytes = match miniblock_chunk_size { + MiniblockChunkSize::U16 => 2, + MiniblockChunkSize::U32 => 4, + }; let mut meta_buffer = Vec::with_capacity(miniblocks.chunks.len() * chunk_size_bytes); let mut rep_iter = rep.map(|r| r.into_iter()); @@ -4703,7 +4823,7 @@ impl PrimitiveStructuralEncoder { data_buffer.extend_from_slice(&bytes_def.to_le_bytes()); } - if support_large_chunk { + if miniblock_chunk_size == MiniblockChunkSize::U32 { for &buffer_size in &chunk.buffer_sizes { data_buffer.extend_from_slice(&buffer_size.to_le_bytes()); } @@ -4748,10 +4868,9 @@ impl PrimitiveStructuralEncoder { } let chunk_bytes = data_buffer.len() - start_pos; - let max_chunk_size = if support_large_chunk { - 1_u64 << 31 // 28 bits of 8-byte words in u32 metadata - } else { - 32 * 1024 // 32KiB limit with u16 metadata + let max_chunk_size = match miniblock_chunk_size { + MiniblockChunkSize::U16 => 32 * 1024, + MiniblockChunkSize::U32 => 1_u64 << 31, }; if chunk_bytes == 0 || chunk_bytes as u64 > max_chunk_size { return Err(Error::internal(format!( @@ -4778,7 +4897,7 @@ impl PrimitiveStructuralEncoder { let divided_bytes_minus_one = (divided_bytes - 1) as u64; let metadata = (divided_bytes_minus_one << 4) | chunk.log_num_values as u64; - if support_large_chunk { + if miniblock_chunk_size == MiniblockChunkSize::U32 { meta_buffer.extend_from_slice(&(metadata as u32).to_le_bytes()); } else { meta_buffer.extend_from_slice(&(metadata as u16).to_le_bytes()); @@ -4974,10 +5093,10 @@ impl PrimitiveStructuralEncoder { repdef: crate::repdef::SerializedRepDefs, row_number: u64, num_rows: u64, - version: LanceFileVersion, + complex_null_encoding: ComplexNullEncoding, compression_strategy: &dyn CompressionStrategy, ) -> Result { - if version.resolve() < LanceFileVersion::V2_2 { + if complex_null_encoding == ComplexNullEncoding::RawLevels { let rep_bytes = if let Some(rep) = repdef.repetition_levels.as_ref() { LanceBuffer::reinterpret_slice(rep.clone()) } else { @@ -5271,7 +5390,7 @@ impl PrimitiveStructuralEncoder { row_number: u64, dictionary_data: Option, num_rows: u64, - support_large_chunk: bool, + miniblock_chunk_size: MiniblockChunkSize, ) -> Result { if let DataBlock::AllNull(_null_block) = data { // We should not be using mini-block for all-null. There are other structural @@ -5331,7 +5450,8 @@ impl PrimitiveStructuralEncoder { .map(|cd| std::mem::take(&mut cd.data)); let serialized = - Self::serialize_miniblocks(compressed_data, rep_data, def_data, support_large_chunk)?; + Self::serialize_miniblocks(compressed_data, rep_data, def_data, miniblock_chunk_size)?; + let has_large_chunk = miniblock_chunk_size == MiniblockChunkSize::U32; // Metadata, Data, Dictionary, (maybe) Repetition Index let mut data = Vec::with_capacity(4); @@ -5360,7 +5480,7 @@ impl PrimitiveStructuralEncoder { Some((dictionary_encoding, num_dictionary_items)), &repdef.def_meaning, num_items, - support_large_chunk, + has_large_chunk, ); Ok(EncodedPage { num_rows, @@ -5379,7 +5499,7 @@ impl PrimitiveStructuralEncoder { None, &repdef.def_meaning, num_items, - support_large_chunk, + has_large_chunk, ); if let Some(rep_index) = rep_index { @@ -5711,7 +5831,7 @@ impl PrimitiveStructuralEncoder { fn should_dictionary_encode( data_block: &DataBlock, field: &Field, - version: LanceFileVersion, + fixed_width_dictionary_encoding: FixedWidthDictionaryEncoding, ) -> Option { const DEFAULT_SAMPLE_SIZE: usize = 4096; const DEFAULT_SAMPLE_UNIQUE_RATIO: f64 = 0.98; @@ -5720,7 +5840,9 @@ impl PrimitiveStructuralEncoder { // estimating the size for other types. match data_block { DataBlock::FixedWidth(fixed) => { - if fixed.bits_per_value == 64 && version < LanceFileVersion::V2_2 { + if fixed.bits_per_value == 64 + && fixed_width_dictionary_encoding == FixedWidthDictionaryEncoding::Exclude64Bit + { return None; } if fixed.bits_per_value != 64 && fixed.bits_per_value != 128 { @@ -6039,14 +6161,18 @@ impl PrimitiveStructuralEncoder { Ok(pages) } - fn encode_page(ctx: PrimitiveEncodeContext, page: PrimitivePageData) -> Result { + fn encode_dense_page( + ctx: PrimitiveEncodeContext, + page: PrimitivePageData, + compression_strategy: Arc, + miniblock_chunk_size: MiniblockChunkSize, + complex_null_encoding: ComplexNullEncoding, + fixed_width_dictionary_encoding: FixedWidthDictionaryEncoding, + ) -> Result { let PrimitiveEncodeContext { column_idx, field, - compression_strategy, encoding_metadata, - support_large_chunk, - version, is_simple_validity, has_repdef_info, } = ctx; @@ -6063,33 +6189,8 @@ impl PrimitiveStructuralEncoder { repdef, single_row_miniblock_repdef_levels, } => (repdef, single_row_miniblock_repdef_levels), - PrimitivePageStructure::Sparse { - plan, - prepared_values, - } => { - log::debug!( - "Encoding column {} with {} visible items ({} rows) using sparse layout", - column_idx, - num_values, - num_rows - ); - return sparse::writer::encode_page( - column_idx, - &field, - compression_strategy.as_ref(), - prepared_values.map_or_else( - || { - sparse::writer::SparseValueInput::Unprepared(DataBlock::from_arrays( - &arrays, num_values, - )) - }, - sparse::writer::SparseValueInput::Prepared, - ), - plan, - row_number, - num_rows, - support_large_chunk, - ); + PrimitivePageStructure::Sparse { .. } => { + unreachable!("dense atom received sparse page") } }; @@ -6107,7 +6208,7 @@ impl PrimitiveStructuralEncoder { repdef, row_number, num_rows, - version, + complex_null_encoding, compression_strategy.as_ref(), ); } @@ -6139,7 +6240,7 @@ impl PrimitiveStructuralEncoder { repdef, row_number, num_rows, - version, + complex_null_encoding, compression_strategy.as_ref(), ) }; @@ -6158,20 +6259,6 @@ impl PrimitiveStructuralEncoder { let data_block = DataBlock::from_arrays(&arrays, num_values); - if version.resolve() >= LanceFileVersion::V2_2 - && let Some(scalar) = Self::find_constant_scalar(&arrays, leaf_validity.as_ref())? - { - log::debug!( - "Encoding column {} with {} items ({} rows) using constant layout", - column_idx, - num_values, - num_rows - ); - return constant::encode_constant_page( - column_idx, scalar, repdef, row_number, num_rows, - ); - } - if let Some(num_levels) = single_row_miniblock_repdef_levels { let requested_encoding = encoding_metadata .get(STRUCTURAL_ENCODING_META_KEY) @@ -6320,24 +6407,29 @@ impl PrimitiveStructuralEncoder { row_number, Some(dictionary_data_block), num_rows, - support_large_chunk, + miniblock_chunk_size, ); } // Try dictionary encoding first if applicable. If encoding aborts, fall back to the // preferred structural encoding. - let dict_result = Self::should_dictionary_encode(&data_block, &field, version).and_then(|budget| { - log::debug!( - "Encoding column {} with {} items using dictionary encoding (mini-block layout)", - column_idx, - num_values - ); - dict::dictionary_encode( - &data_block, - budget.max_dict_entries, - budget.max_encoded_size, - ) - }); + let dict_result = Self::should_dictionary_encode( + &data_block, + &field, + fixed_width_dictionary_encoding, + ) + .and_then(|budget| { + log::debug!( + "Encoding column {} with {} items using dictionary encoding (mini-block layout)", + column_idx, + num_values + ); + dict::dictionary_encode( + &data_block, + budget.max_dict_entries, + budget.max_encoded_size, + ) + }); if let Some((indices_data_block, dictionary_data_block)) = dict_result { Self::encode_miniblock( @@ -6349,7 +6441,7 @@ impl PrimitiveStructuralEncoder { row_number, Some(dictionary_data_block), num_rows, - support_large_chunk, + miniblock_chunk_size, ) } else if Self::prefers_miniblock(&data_block, encoding_metadata.as_ref()) { log::debug!( @@ -6366,7 +6458,7 @@ impl PrimitiveStructuralEncoder { row_number, None, num_rows, - support_large_chunk, + miniblock_chunk_size, ) } else if Self::prefers_fullzip(encoding_metadata.as_ref()) { log::debug!( @@ -6400,96 +6492,48 @@ impl PrimitiveStructuralEncoder { let is_simple_validity = repdefs.iter().all(|rd| rd.is_simple_validity()); let has_repdef_info = repdefs.iter().any(|rd| !rd.is_empty()); let normalized = RepDefBuilder::normalize(repdefs); - let requested_encoding = self.encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY); - let requests_sparse = requested_encoding - .is_some_and(|requested| requested.eq_ignore_ascii_case(STRUCTURAL_ENCODING_SPARSE)); - let sparse_plan = requests_sparse - .then(|| sparse::writer::plan(&normalized, num_values)) - .transpose()?; - let pages = if let Some(plan) = sparse_plan - && !sparse::writer::uses_constant_layout(&plan, &self.field) - { - vec![PrimitivePageData { - arrays, - structure: PrimitivePageStructure::Sparse { - plan, - prepared_values: None, - }, + let plan_ctx = PrimitivePlanContext { + column_idx: self.column_index, + field: &self.field, + encoding_metadata: &self.encoding_metadata, + }; + let mut pages = None; + for page_encoding in self.page_encodings.iter() { + if let Some(planned) = page_encoding.behavior.try_plan_pages( + &plan_ctx, + &arrays, + &normalized, row_number, num_rows, - }] - } else { - let (repdef, miniblock_repdef_budget) = normalized - .serialize_with_miniblock_repdef_budget( - miniblock::max_repdef_levels_per_chunk, - num_rows, - num_values, - )?; - let automatic_sparse = layout::select_automatic_sparse( - self.version.resolve(), - requested_encoding.map(String::as_str), - &miniblock_repdef_budget, - || { - let data = DataBlock::from_arrays(&arrays, num_values); - if !sparse::writer::supports_value_block(&data) { - return Ok(None); - } - let prepared_values = match sparse::writer::prepare_values( - &self.field, - self.compression_strategy.as_ref(), - data, - self.support_large_chunk, - ) { - Ok(prepared_values) => prepared_values, - Err(error) => { - debug!( - "Keeping column {} on its dense structural path because sparse value preparation is unavailable: {}", - self.column_index, error - ); - return Ok(None); - } - }; - let plan = sparse::writer::plan(&normalized, num_values)?; - if sparse::writer::uses_constant_layout(&plan, &self.field) { - return Ok(None); - } - Ok(Some((plan, prepared_values))) - }, - )?; - match automatic_sparse { - Some((plan, prepared_values)) => vec![PrimitivePageData { - arrays, - structure: PrimitivePageStructure::Sparse { - plan, - prepared_values: Some(prepared_values), - }, - row_number, - num_rows, - }], - None => Self::split_pages_for_miniblock_repdef_budget( - arrays, - repdef, - miniblock_repdef_budget, - row_number, - num_rows, - )?, + num_values, + )? { + pages = Some(planned); + break; } - }; + } + let pages = pages.ok_or_else(|| { + Error::invalid_input_source( + format!( + "No primitive page planner supports field '{}'", + self.field.name + ) + .into(), + ) + })?; let mut tasks = Vec::with_capacity(pages.len()); let ctx = PrimitiveEncodeContext { column_idx: self.column_index, field: self.field.clone(), - compression_strategy: self.compression_strategy.clone(), encoding_metadata: self.encoding_metadata.clone(), - support_large_chunk: self.support_large_chunk, - version: self.version, is_simple_validity, has_repdef_info, }; for page in pages { let ctx = ctx.clone(); - let task = spawn_cpu(move || Self::encode_page(ctx, page)).boxed(); + let page_encodings = self.page_encodings.clone(); + let task = + spawn_cpu(move || Self::encode_page(page_encodings.as_ref(), &ctx, page)).boxed(); tasks.push(task); } Ok(tasks) @@ -6541,6 +6585,291 @@ impl PrimitiveStructuralEncoder { } } +impl PrimitivePageEncodingBehavior for RejectSparsePrimitiveEncoding { + fn validate_field(&self, field: &Field, metadata: &HashMap) -> Result<()> { + if metadata + .get(STRUCTURAL_ENCODING_META_KEY) + .is_some_and(|requested| requested.eq_ignore_ascii_case(STRUCTURAL_ENCODING_SPARSE)) + { + return Err(Error::invalid_input_source( + format!( + "Field '{}' requests sparse structural encoding, which is not enabled by the selected file format", + field.name + ) + .into(), + )); + } + Ok(()) + } +} + +fn plan_dense_primitive_pages( + arrays: &[ArrayRef], + normalized: &NormalizedStructuralPlan, + row_number: u64, + num_rows: u64, + num_values: u64, +) -> Result> { + let (repdef, miniblock_repdef_budget) = normalized.serialize_with_miniblock_repdef_budget( + miniblock::max_repdef_levels_per_chunk, + num_rows, + num_values, + )?; + PrimitiveStructuralEncoder::split_pages_for_miniblock_repdef_budget( + arrays.to_vec(), + repdef, + miniblock_repdef_budget, + row_number, + num_rows, + ) +} + +impl PrimitivePageEncodingBehavior for DenseU16PrimitiveEncoding { + fn try_plan_pages( + &self, + _ctx: &PrimitivePlanContext<'_>, + arrays: &[ArrayRef], + normalized: &NormalizedStructuralPlan, + row_number: u64, + num_rows: u64, + num_values: u64, + ) -> Result>> { + Ok(Some(plan_dense_primitive_pages( + arrays, normalized, row_number, num_rows, num_values, + )?)) + } + + fn try_encode_page( + &self, + ctx: &PrimitiveEncodeContext, + page: PrimitivePageData, + ) -> Result { + if !matches!(&page.structure, PrimitivePageStructure::Dense { .. }) { + return Ok(PrimitiveEncodeAttempt::Unhandled(page)); + } + Ok(PrimitiveEncodeAttempt::Encoded( + PrimitiveStructuralEncoder::encode_dense_page( + ctx.clone(), + page, + self.compression.clone(), + MiniblockChunkSize::U16, + ComplexNullEncoding::RawLevels, + FixedWidthDictionaryEncoding::Exclude64Bit, + )?, + )) + } +} + +impl PrimitivePageEncodingBehavior for DenseU32PrimitiveEncoding { + fn try_plan_pages( + &self, + _ctx: &PrimitivePlanContext<'_>, + arrays: &[ArrayRef], + normalized: &NormalizedStructuralPlan, + row_number: u64, + num_rows: u64, + num_values: u64, + ) -> Result>> { + Ok(Some(plan_dense_primitive_pages( + arrays, normalized, row_number, num_rows, num_values, + )?)) + } + + fn try_encode_page( + &self, + ctx: &PrimitiveEncodeContext, + page: PrimitivePageData, + ) -> Result { + if !matches!(&page.structure, PrimitivePageStructure::Dense { .. }) { + return Ok(PrimitiveEncodeAttempt::Unhandled(page)); + } + Ok(PrimitiveEncodeAttempt::Encoded( + PrimitiveStructuralEncoder::encode_dense_page( + ctx.clone(), + page, + self.compression.clone(), + MiniblockChunkSize::U32, + ComplexNullEncoding::CompressedLevels, + FixedWidthDictionaryEncoding::Include64Bit, + )?, + )) + } +} + +impl PrimitivePageEncodingBehavior for SparsePrimitiveEncoding { + fn try_plan_pages( + &self, + ctx: &PrimitivePlanContext<'_>, + arrays: &[ArrayRef], + normalized: &NormalizedStructuralPlan, + row_number: u64, + num_rows: u64, + num_values: u64, + ) -> Result>> { + let requested_encoding = ctx.encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY); + let requests_sparse = requested_encoding + .is_some_and(|requested| requested.eq_ignore_ascii_case(STRUCTURAL_ENCODING_SPARSE)); + if requests_sparse { + let plan = sparse::writer::plan(normalized, num_values)?; + if sparse::writer::uses_constant_layout(&plan, ctx.field) { + return Ok(None); + } + return Ok(Some(vec![PrimitivePageData { + arrays: arrays.to_vec(), + structure: PrimitivePageStructure::Sparse { + plan, + prepared_values: None, + }, + row_number, + num_rows, + }])); + } + + let (_, miniblock_repdef_budget) = normalized.serialize_with_miniblock_repdef_budget( + miniblock::max_repdef_levels_per_chunk, + num_rows, + num_values, + )?; + let automatic_sparse = layout::select_automatic_sparse( + requested_encoding.map(String::as_str), + &miniblock_repdef_budget, + || { + let data = DataBlock::from_arrays(arrays, num_values); + if !sparse::writer::supports_value_block(&data) { + return Ok(None); + } + let prepared_values = match sparse::writer::prepare_values( + ctx.field, + self.compression.as_ref(), + data, + MiniblockChunkSize::U32, + ) { + Ok(prepared_values) => prepared_values, + Err(error) => { + debug!( + "Keeping column {} on its dense structural path because sparse value preparation is unavailable: {}", + ctx.column_idx, error + ); + return Ok(None); + } + }; + let plan = sparse::writer::plan(normalized, num_values)?; + if sparse::writer::uses_constant_layout(&plan, ctx.field) { + return Ok(None); + } + Ok(Some((plan, prepared_values))) + }, + )?; + Ok(automatic_sparse.map(|(plan, prepared_values)| { + vec![PrimitivePageData { + arrays: arrays.to_vec(), + structure: PrimitivePageStructure::Sparse { + plan, + prepared_values: Some(prepared_values), + }, + row_number, + num_rows, + }] + })) + } + + fn try_encode_page( + &self, + ctx: &PrimitiveEncodeContext, + page: PrimitivePageData, + ) -> Result { + if !matches!(&page.structure, PrimitivePageStructure::Sparse { .. }) { + return Ok(PrimitiveEncodeAttempt::Unhandled(page)); + } + let PrimitivePageData { + arrays, + structure: + PrimitivePageStructure::Sparse { + plan, + prepared_values, + }, + row_number, + num_rows, + } = page + else { + unreachable!() + }; + let num_values = arrays.iter().map(|array| array.len() as u64).sum(); + log::debug!( + "Encoding column {} with {} visible items ({} rows) using sparse layout", + ctx.column_idx, + num_values, + num_rows + ); + Ok(PrimitiveEncodeAttempt::Encoded( + sparse::writer::encode_page( + ctx.column_idx, + &ctx.field, + self.compression.as_ref(), + prepared_values.map_or_else( + || { + sparse::writer::SparseValueInput::Unprepared(DataBlock::from_arrays( + &arrays, num_values, + )) + }, + sparse::writer::SparseValueInput::Prepared, + ), + plan, + row_number, + num_rows, + MiniblockChunkSize::U32, + )?, + )) + } +} + +impl PrimitivePageEncodingBehavior for ConstantPrimitiveEncoding { + fn try_encode_page( + &self, + ctx: &PrimitiveEncodeContext, + page: PrimitivePageData, + ) -> Result { + let PrimitivePageStructure::Dense { repdef, .. } = &page.structure else { + return Ok(PrimitiveEncodeAttempt::Unhandled(page)); + }; + let num_values: u64 = page.arrays.iter().map(|array| array.len() as u64).sum(); + if num_values == 0 { + return Ok(PrimitiveEncodeAttempt::Unhandled(page)); + } + let leaf_validity = PrimitiveStructuralEncoder::leaf_validity(repdef, num_values as usize)?; + if leaf_validity + .as_ref() + .is_some_and(|validity| validity.count_set_bits() == 0) + || matches!(ctx.field.data_type(), DataType::Struct(fields) if fields.is_empty()) + { + return Ok(PrimitiveEncodeAttempt::Unhandled(page)); + } + let Some(scalar) = + PrimitiveStructuralEncoder::find_constant_scalar(&page.arrays, leaf_validity.as_ref())? + else { + return Ok(PrimitiveEncodeAttempt::Unhandled(page)); + }; + let PrimitivePageData { + structure: PrimitivePageStructure::Dense { repdef, .. }, + row_number, + num_rows, + .. + } = page + else { + unreachable!() + }; + log::debug!( + "Encoding column {} with {} items ({} rows) using constant layout", + ctx.column_idx, + num_values, + num_rows + ); + Ok(PrimitiveEncodeAttempt::Encoded( + constant::encode_constant_page(ctx.column_idx, scalar, repdef, row_number, num_rows)?, + )) + } +} + impl FieldEncoder for PrimitiveStructuralEncoder { // Buffers data, if there is enough to write a page then we create an encode task fn maybe_encode( @@ -6591,11 +6920,12 @@ impl FieldEncoder for PrimitiveStructuralEncoder { mod tests { use super::{ ChunkInstructions, DataBlock, DecodeMiniBlockTask, FixedPerValueDecompressor, - FixedWidthDataBlock, FullZipCacheableState, FullZipDecodeDetails, FullZipReadSource, - FullZipRepIndexDetails, FullZipScheduler, LazyLevels, LevelCodec, LevelCursor, LevelPlan, - MiniBlockChunk, MiniBlockChunkIndex, MiniBlockCompressed, PerValueDecompressor, - PreambleAction, RunEndsBuilder, RunPosition, RunStorage, StructuralPageScheduler, - VariableFullZipDecoder, dense_levels_from_block, validate_complex_all_null_levels, + FixedWidthDataBlock, FixedWidthDictionaryEncoding, FullZipCacheableState, + FullZipDecodeDetails, FullZipReadSource, FullZipRepIndexDetails, FullZipScheduler, + LazyLevels, LevelCodec, LevelCursor, LevelPlan, MiniBlockChunk, MiniBlockChunkIndex, + MiniBlockCompressed, MiniblockChunkSize, PerValueDecompressor, PreambleAction, + RunEndsBuilder, RunPosition, RunStorage, StructuralPageScheduler, VariableFullZipDecoder, + dense_levels_from_block, validate_complex_all_null_levels, }; use crate::buffer::LanceBuffer; use crate::compression::{BlockCompressor, DefaultDecompressionStrategy}; @@ -6614,8 +6944,8 @@ mod tests { use crate::format::pb21; use crate::format::pb21::compressive_encoding::Compression; use crate::repdef::build_control_word_iterator; + use crate::testing::TestEncoding; use crate::testing::{TestCases, check_round_trip_encoding_of_data}; - use crate::version::LanceFileVersion; use arrow_array::{ArrayRef, Int8Array, StringArray}; use arrow_buffer::ScalarBuffer; use arrow_schema::{DataType, Field as ArrowField}; @@ -7934,7 +8264,7 @@ mod tests { ); let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_1) + .with_structural_encodings() .with_batch_size(100) .with_range(0..num_rows.min(500) as u64) .with_indices(vec![0, num_rows as u64 / 2, (num_rows - 1) as u64]); @@ -7945,7 +8275,7 @@ mod tests { async fn test_minichunk_size_helper( string_data: Vec>, minichunk_size: u64, - file_version: LanceFileVersion, + encodings: &[TestEncoding], ) { use crate::constants::MINICHUNK_SIZE_META_KEY; use crate::testing::{TestCases, check_round_trip_encoding_of_data}; @@ -7965,7 +8295,7 @@ mod tests { ); let test_cases = TestCases::default() - .with_min_file_version(file_version) + .with_encodings(encodings.iter().copied()) .with_batch_size(1000); check_round_trip_encoding_of_data(vec![string_array], &test_cases, metadata).await; @@ -7979,7 +8309,16 @@ mod tests { string_data.push(Some(format!("test_string_{}", i).repeat(50))); } // configure minichunk size to 64 bytes (smaller than the default 4kb) for Lance 2.1 - test_minichunk_size_helper(string_data, 64, LanceFileVersion::V2_1).await; + test_minichunk_size_helper( + string_data, + 64, + &[ + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse, + ], + ) + .await; } #[tokio::test] @@ -7990,7 +8329,12 @@ mod tests { for i in 0..10000 { string_data.push(Some(format!("test_string_{}", i).repeat(50))); } - test_minichunk_size_helper(string_data, 128 * 1024, LanceFileVersion::V2_2).await; + test_minichunk_size_helper( + string_data, + 128 * 1024, + &[TestEncoding::StructuralU32, TestEncoding::StructuralSparse], + ) + .await; } #[tokio::test] @@ -8000,7 +8344,12 @@ mod tests { for i in 0..10000 { string_data.push(Some(format!("t_{}", i))); } - test_minichunk_size_helper(string_data, 128 * 1024, LanceFileVersion::V2_2).await; + test_minichunk_size_helper( + string_data, + 128 * 1024, + &[TestEncoding::StructuralU32, TestEncoding::StructuralSparse], + ) + .await; } #[tokio::test] @@ -8027,7 +8376,7 @@ mod tests { // Configure test to use V2_2 and verify encoding let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_2) + .with_u32_structural_encodings() .with_verify_encoding(Arc::new(|cols: &[crate::encoder::EncodedColumn], _| { assert_eq!(cols.len(), 1); let col = &cols[0]; @@ -8097,7 +8446,7 @@ mod tests { ) .with_metadata(metadata); - encode_first_page(field, dict_array, LanceFileVersion::V2_2).await + encode_first_page(field, dict_array, TestEncoding::StructuralU32).await } async fn encode_auto_fixed_dict_page( @@ -8127,7 +8476,7 @@ mod tests { let field = arrow_schema::Field::new("fixed_col", DataType::Decimal128(38, 0), false) .with_metadata(field_metadata); - encode_first_page(field, decimal, LanceFileVersion::V2_2).await + encode_first_page(field, decimal, TestEncoding::StructuralU32).await } #[tokio::test] @@ -8254,7 +8603,6 @@ mod tests { async fn test_dictionary_encode_int64() { use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY}; use crate::testing::{TestCases, check_round_trip_encoding_of_data}; - use crate::version::LanceFileVersion; use arrow_array::{ArrayRef, Int64Array}; use std::collections::HashMap; use std::sync::Arc; @@ -8277,7 +8625,7 @@ mod tests { metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string()); let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_2) + .with_u32_structural_encodings() .with_batch_size(1000) .with_range(0..1000) .with_indices(vec![0, 1, 10, 999]) @@ -8290,7 +8638,6 @@ mod tests { async fn test_dictionary_encode_float64() { use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY}; use crate::testing::{TestCases, check_round_trip_encoding_of_data}; - use crate::version::LanceFileVersion; use arrow_array::{ArrayRef, Float64Array}; use std::collections::HashMap; use std::sync::Arc; @@ -8313,7 +8660,7 @@ mod tests { metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string()); let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_2) + .with_u32_structural_encodings() .with_batch_size(1000) .with_range(0..1000) .with_indices(vec![0, 1, 10, 999]) @@ -8448,7 +8795,7 @@ mod tests { let result = PrimitiveStructuralEncoder::should_dictionary_encode( &block, &field, - LanceFileVersion::V2_1, + FixedWidthDictionaryEncoding::Exclude64Bit, ); assert!( @@ -8495,7 +8842,7 @@ mod tests { let result = PrimitiveStructuralEncoder::should_dictionary_encode( &block, &field, - LanceFileVersion::V2_2, + FixedWidthDictionaryEncoding::Include64Bit, ); assert!( @@ -8522,7 +8869,7 @@ mod tests { let result = PrimitiveStructuralEncoder::should_dictionary_encode( &block, &field, - LanceFileVersion::V2_2, + FixedWidthDictionaryEncoding::Include64Bit, ); assert!( @@ -8540,7 +8887,7 @@ mod tests { let field = arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata); let array = create_sorted_string_array(200_000, 8_000); - let page = encode_first_page(field, array, LanceFileVersion::V2_2).await; + let page = encode_first_page(field, array, TestEncoding::StructuralU32).await; let _ = dictionary_encoding_from_page(&page); } @@ -8560,7 +8907,7 @@ mod tests { let result = PrimitiveStructuralEncoder::should_dictionary_encode( &block, &field, - LanceFileVersion::V2_1, + FixedWidthDictionaryEncoding::Exclude64Bit, ); assert!( @@ -8586,7 +8933,7 @@ mod tests { let result = PrimitiveStructuralEncoder::should_dictionary_encode( &block, &field, - LanceFileVersion::V2_1, + FixedWidthDictionaryEncoding::Exclude64Bit, ); assert!( @@ -8612,9 +8959,13 @@ mod tests { num_values: 32_769, }; - let serialized = - PrimitiveStructuralEncoder::serialize_miniblocks(miniblocks, None, None, false) - .unwrap(); + let serialized = PrimitiveStructuralEncoder::serialize_miniblocks( + miniblocks, + None, + None, + MiniblockChunkSize::U16, + ) + .unwrap(); let chunk_metadata = serialized.metadata.borrow_to_typed_slice::(); assert_eq!(chunk_metadata.len(), 2); @@ -8628,33 +8979,33 @@ mod tests { async fn encode_first_page( field: arrow_schema::Field, array: ArrayRef, - version: LanceFileVersion, + version: TestEncoding, ) -> crate::encoder::EncodedPage { - use crate::encoder::{ - ColumnIndexSequence, EncodingOptions, MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers, - default_encoding_strategy, - }; use crate::repdef::RepDefBuilder; + use crate::{ + encoder::{ + ColumnIndexSequence, EncodingOptions, MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers, + }, + testing::{create_test_field_encoder, test_encoding_strategy}, + }; let lance_field = lance_core::datatypes::Field::try_from(&field).unwrap(); - let encoding_strategy = default_encoding_strategy(version); + let encoding_strategy = test_encoding_strategy(version); let mut column_index_seq = ColumnIndexSequence::default(); let encoding_options = EncodingOptions { cache_bytes_per_column: 1, max_page_bytes: 32 * 1024 * 1024, keep_original_array: true, buffer_alignment: MIN_PAGE_BUFFER_ALIGNMENT, - version, }; - let mut encoder = encoding_strategy - .create_field_encoder( - encoding_strategy.as_ref(), - &lance_field, - &mut column_index_seq, - &encoding_options, - ) - .unwrap(); + let mut encoder = create_test_field_encoder( + encoding_strategy.as_ref(), + &lance_field, + &mut column_index_seq, + &encoding_options, + ) + .unwrap(); let mut external_buffers = OutOfLineBuffers::new(0, MIN_PAGE_BUFFER_ALIGNMENT); let repdef = RepDefBuilder::default(); @@ -8685,7 +9036,7 @@ mod tests { .unwrap(), ); let field = arrow_schema::Field::new("c", DataType::FixedSizeBinary(33), true); - let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await; + let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await; let PageEncoding::Structural(layout) = &page.description else { panic!("Expected structural encoding"); @@ -8697,8 +9048,7 @@ mod tests { assert_eq!(page.data.len(), 1); let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_2) - .with_max_file_version(LanceFileVersion::V2_2) + .with_encoding(TestEncoding::StructuralU32) .with_page_sizes(vec![4096]); check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; } @@ -8711,7 +9061,7 @@ mod tests { std::iter::repeat_n("hello", 512), )); let field = arrow_schema::Field::new("c", DataType::Utf8, true); - let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await; + let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await; let PageEncoding::Structural(layout) = &page.description else { panic!("Expected structural encoding"); @@ -8723,8 +9073,7 @@ mod tests { assert_eq!(page.data.len(), 1); let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_2) - .with_max_file_version(LanceFileVersion::V2_2) + .with_encoding(TestEncoding::StructuralU32) .with_page_sizes(vec![4096]); check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; } @@ -8741,7 +9090,7 @@ mod tests { Some(7), ])); let field = arrow_schema::Field::new("c", DataType::Int32, true); - let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await; + let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await; let PageEncoding::Structural(layout) = &page.description else { panic!("Expected structural encoding"); @@ -8753,8 +9102,7 @@ mod tests { assert_eq!(page.data.len(), 2); let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_2) - .with_max_file_version(LanceFileVersion::V2_2) + .with_encoding(TestEncoding::StructuralU32) .with_page_sizes(vec![4096]); check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; } @@ -8787,7 +9135,7 @@ mod tests { ))), true, ); - let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await; + let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await; let PageEncoding::Structural(layout) = &page.description else { panic!("Expected structural encoding"); @@ -8799,8 +9147,7 @@ mod tests { assert_eq!(page.data.len(), 2); let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_2) - .with_max_file_version(LanceFileVersion::V2_2) + .with_encoding(TestEncoding::StructuralU32) .with_page_sizes(vec![4096]); check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; } @@ -8826,7 +9173,7 @@ mod tests { ), true, ); - let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await; + let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await; if let PageEncoding::Structural(layout) = &page.description { assert!( @@ -8836,8 +9183,7 @@ mod tests { } let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_2) - .with_max_file_version(LanceFileVersion::V2_2) + .with_encoding(TestEncoding::StructuralU32) .with_page_sizes(vec![4096]); check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; } @@ -8848,7 +9194,7 @@ mod tests { let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![7; 1024])); let field = arrow_schema::Field::new("c", DataType::Int32, true); - let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_1).await; + let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU16).await; let PageEncoding::Structural(layout) = &page.description else { return; @@ -8859,8 +9205,7 @@ mod tests { ); let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_1) - .with_max_file_version(LanceFileVersion::V2_1) + .with_encoding(TestEncoding::StructuralU16) .with_page_sizes(vec![4096]); check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; } @@ -8871,7 +9216,7 @@ mod tests { let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![None, None, None])); let field = arrow_schema::Field::new("c", DataType::Int32, true); - let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await; + let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await; let PageEncoding::Structural(layout) = &page.description else { panic!("Expected structural encoding"); @@ -8883,26 +9228,26 @@ mod tests { assert_eq!(page.data.len(), 0); let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_2) - .with_max_file_version(LanceFileVersion::V2_2) + .with_encoding(TestEncoding::StructuralU32) .with_page_sizes(vec![4096]); check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; } #[test] fn test_encode_decode_complex_all_null_vals_roundtrip() { - use crate::compression::{ - DecompressionStrategy, DefaultCompressionStrategy, DefaultDecompressionStrategy, - }; + use crate::compression::{DecompressionStrategy, DefaultDecompressionStrategy}; let values: Arc<[u16]> = Arc::from((0..2048).map(|i| (i % 5) as u16).collect::>()); - let compression_strategy = DefaultCompressionStrategy::default(); + let compression_strategy = crate::testing::test_compression_strategy( + TestEncoding::StructuralU16, + crate::compression_config::CompressionParams::default(), + ); let decompression_strategy = DefaultDecompressionStrategy::default(); let (compressed_buf, encoding) = PrimitiveStructuralEncoder::encode_complex_all_null_vals( &values, - &compression_strategy, + compression_strategy.as_ref(), ) .unwrap(); @@ -8938,7 +9283,8 @@ mod tests { true, ); - let page_v21 = encode_first_page(field.clone(), arr.clone(), LanceFileVersion::V2_1).await; + let page_v21 = + encode_first_page(field.clone(), arr.clone(), TestEncoding::StructuralU16).await; let PageEncoding::Structural(layout_v21) = &page_v21.description else { panic!("Expected structural encoding"); }; @@ -8950,7 +9296,7 @@ mod tests { assert_eq!(layout_v21.num_rep_values, 0); assert_eq!(layout_v21.num_def_values, 0); - let page_v22 = encode_first_page(field, arr, LanceFileVersion::V2_2).await; + let page_v22 = encode_first_page(field, arr, TestEncoding::StructuralU32).await; let PageEncoding::Structural(layout_v22) = &page_v22.description else { panic!("Expected structural encoding"); }; @@ -8969,7 +9315,7 @@ mod tests { (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }), ); - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_2); + let test_cases = TestCases::default().with_u32_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new()) .await; } @@ -8984,7 +9330,7 @@ mod tests { (0..5000).map(|_| None::>>), ); - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_2); + let test_cases = TestCases::default().with_u32_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new()) .await; } @@ -9492,7 +9838,7 @@ mod tests { } let list_array = Arc::new(list_builder.finish()); - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); + let test_cases = TestCases::default().with_structural_encodings(); check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; } } diff --git a/rust/lance-encoding/src/encodings/logical/primitive/layout.rs b/rust/lance-encoding/src/encodings/logical/primitive/layout.rs index f9749219ea6..95df12ede6b 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/layout.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/layout.rs @@ -3,17 +3,15 @@ use lance_core::Result; -use crate::{repdef::MiniBlockRepDefBudget, version::LanceFileVersion}; +use crate::repdef::MiniBlockRepDefBudget; /// Runs automatic sparse planning only after the dense mini-block budget makes it useful. pub(super) fn select_automatic_sparse( - version: LanceFileVersion, requested_encoding: Option<&str>, dense_budget: &MiniBlockRepDefBudget, candidate: impl FnOnce() -> Result>, ) -> Result> { - if version < LanceFileVersion::V2_3 - || requested_encoding.is_some() + if requested_encoding.is_some() || !matches!( dense_budget, MiniBlockRepDefBudget::RequiresPageSplit(_) @@ -34,32 +32,25 @@ mod tests { #[test] fn within_budget_does_not_construct_sparse_candidate() { - let selected = select_automatic_sparse::<()>( - LanceFileVersion::V2_3, - None, - &MiniBlockRepDefBudget::WithinBudget, - || panic!("within-budget pages must not construct sparse candidates"), - ) - .unwrap(); + let selected = + select_automatic_sparse::<()>(None, &MiniBlockRepDefBudget::WithinBudget, || { + panic!("within-budget pages must not construct sparse candidates") + }) + .unwrap(); assert!(selected.is_none()); } #[test] fn over_budget_selects_only_eligible_candidates() { let split = MiniBlockRepDefBudget::RequiresPageSplit(Vec::new()); - let selected = - select_automatic_sparse(LanceFileVersion::V2_3, None, &split, || Ok(Some(42))).unwrap(); + let selected = select_automatic_sparse(None, &split, || Ok(Some(42))).unwrap(); assert_eq!(selected, Some(42)); - let ineligible = - select_automatic_sparse::<()>(LanceFileVersion::V2_3, None, &split, || Ok(None)) - .unwrap(); + let ineligible = select_automatic_sparse::<()>(None, &split, || Ok(None)).unwrap(); assert!(ineligible.is_none()); let unsplittable = MiniBlockRepDefBudget::SingleRowOverBudget(70_000); - let selected = - select_automatic_sparse(LanceFileVersion::V2_3, None, &unsplittable, || Ok(Some(7))) - .unwrap(); + let selected = select_automatic_sparse(None, &unsplittable, || Ok(Some(7))).unwrap(); assert_eq!(selected, Some(7)); } @@ -71,20 +62,11 @@ mod tests { STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_SPARSE, ] { - let selected = select_automatic_sparse::<()>( - LanceFileVersion::V2_3, - Some(requested), - &split, - || panic!("explicit modes must not invoke automatic sparse planning"), - ) + let selected = select_automatic_sparse::<()>(Some(requested), &split, || { + panic!("explicit modes must not invoke automatic sparse planning") + }) .unwrap(); assert!(selected.is_none()); } - - let selected = select_automatic_sparse::<()>(LanceFileVersion::V2_2, None, &split, || { - panic!("Lance 2.2 must not invoke automatic sparse planning") - }) - .unwrap(); - assert!(selected.is_none()); } } diff --git a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs index 3d82b4a5ef6..4dc13fa8aac 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs @@ -19,6 +19,8 @@ use crate::{ statistics::ComputeStat, }; +use super::super::MiniblockChunkSize; + use super::{ SparseCountSet, SparsePositionSet, SparseStructuralLayerPlan, SparseStructuralPlan, SparseValidityMeaning, SparseValiditySet, @@ -438,7 +440,7 @@ fn with_explicit_value_counts( fn serialize_value_chunks( compressed: SparseMiniBlockCompressed, - support_large_chunk: bool, + miniblock_chunk_size: MiniblockChunkSize, ) -> Result { let bytes_data = compressed.data.iter().map(LanceBuffer::len).sum::(); let num_buffers = compressed.data.len(); @@ -457,7 +459,7 @@ fn serialize_value_chunks( let chunk_start = data_buffer.len(); debug_assert_eq!(chunk_start % MINIBLOCK_ALIGNMENT, 0); data_buffer.extend_from_slice(&0_u16.to_le_bytes()); - if support_large_chunk { + if miniblock_chunk_size == MiniblockChunkSize::U32 { for buffer_size in &chunk.buffer_sizes { data_buffer.extend_from_slice(&buffer_size.to_le_bytes()); } @@ -540,7 +542,7 @@ pub fn prepare_values( field: &Field, compression_strategy: &dyn CompressionStrategy, data: DataBlock, - support_large_chunk: bool, + miniblock_chunk_size: MiniblockChunkSize, ) -> Result { match &data { DataBlock::AllNull(_) => { @@ -564,8 +566,10 @@ pub fn prepare_values( let num_values = data.num_values(); let compressor = compression_strategy.create_miniblock_compressor(field, &data)?; let (compressed, value_compression) = compressor.compress(data)?; - let values = - serialize_value_chunks(with_explicit_value_counts(compressed)?, support_large_chunk)?; + let values = serialize_value_chunks( + with_explicit_value_counts(compressed)?, + miniblock_chunk_size, + )?; Ok(PreparedSparseValues { num_values, value_compression, @@ -825,7 +829,7 @@ pub(in crate::encodings::logical::primitive) fn encode_page( plan: SparseStructuralPlan, row_number: u64, num_rows: u64, - support_large_chunk: bool, + miniblock_chunk_size: MiniblockChunkSize, ) -> Result { let PreparedSparseValues { num_values, @@ -833,7 +837,7 @@ pub(in crate::encodings::logical::primitive) fn encode_page( values, } = match values { SparseValueInput::Unprepared(data) => { - prepare_values(field, compression_strategy, data, support_large_chunk)? + prepare_values(field, compression_strategy, data, miniblock_chunk_size)? } SparseValueInput::Prepared(prepared) => prepared, }; @@ -851,7 +855,7 @@ pub(in crate::encodings::logical::primitive) fn encode_page( num_buffers: values.num_buffers, num_items: plan.num_items, num_visible_items: plan.num_visible_items, - has_large_chunk: support_large_chunk, + has_large_chunk: miniblock_chunk_size == MiniblockChunkSize::U32, structural_layers: structural.layers, }, )), @@ -891,10 +895,11 @@ mod tests { data::FixedSizeListBlock, encoder::{ ColumnIndexSequence, EncodingOptions, FieldEncoder, MIN_PAGE_BUFFER_ALIGNMENT, - OutOfLineBuffers, default_encoding_strategy, + OutOfLineBuffers, + }, + testing::{ + TestCases, TestEncoding, check_round_trip_encoding_of_data, test_encoding_strategy, }, - testing::{TestCases, check_round_trip_encoding_of_data}, - version::LanceFileVersion, }; use super::*; @@ -1196,19 +1201,18 @@ mod tests { fn create_encoder( array: &ArrayRef, - version: LanceFileVersion, + version: TestEncoding, metadata: HashMap, ) -> Result> { let arrow_field = ArrowField::new("values", array.data_type().clone(), true).with_metadata(metadata); let field = Field::try_from(&arrow_field)?; - let strategy = default_encoding_strategy(version); + let strategy = test_encoding_strategy(version); let options = EncodingOptions { cache_bytes_per_column: 1, - version, ..Default::default() }; - strategy.create_field_encoder( + crate::testing::create_test_field_encoder( strategy.as_ref(), &field, &mut ColumnIndexSequence::default(), @@ -1218,7 +1222,7 @@ mod tests { async fn encode_pages( array: ArrayRef, - version: LanceFileVersion, + version: TestEncoding, metadata: HashMap, ) -> Result> { encode_chunks(vec![array], version, metadata).await @@ -1226,7 +1230,7 @@ mod tests { async fn encode_chunks( arrays: Vec, - version: LanceFileVersion, + version: TestEncoding, metadata: HashMap, ) -> Result> { let first = arrays @@ -1474,9 +1478,13 @@ mod tests { None, Some(40), ])) as ArrayRef; - let pages = encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) - .await - .unwrap(); + let pages = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); assert_eq!(pages.len(), 1); let sparse = sparse_layout(&pages[0]); assert_eq!(sparse.num_items, 6); @@ -1497,8 +1505,7 @@ mod tests { )); let cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_page_sizes(vec![1]) .with_range(1..5) .with_indices(vec![0, 2, 5]); @@ -1508,9 +1515,13 @@ mod tests { #[tokio::test] async fn test_explicit_sparse_nullable_struct_roundtrip() { let array = nullable_struct(); - let pages = encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) - .await - .unwrap(); + let pages = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); assert_eq!(pages.len(), 1); let sparse = sparse_layout(&pages[0]); assert_eq!(sparse.structural_layers.len(), 2); @@ -1535,8 +1546,7 @@ mod tests { )); let cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_page_sizes(vec![1]) .with_range(1..5) .with_indices(vec![0, 2, 4]); @@ -1563,9 +1573,13 @@ mod tests { ], Some(null_buffer([true, false, true, true, false])), )) as ArrayRef; - let pages = encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) - .await - .unwrap(); + let pages = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); assert!(pages.iter().any(|page| matches!( page_layout(page), pb21::page_layout::Layout::ConstantLayout(_) @@ -1576,8 +1590,7 @@ mod tests { ))); let cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_page_sizes(vec![1]) .with_range(1..5) .with_indices(vec![0, 2, 4]); @@ -1605,7 +1618,7 @@ mod tests { let null_positions = encode_pages( mostly_valid.clone(), - LanceFileVersion::V2_3, + TestEncoding::StructuralSparse, sparse_metadata(), ) .await @@ -1622,7 +1635,7 @@ mod tests { let valid_positions = encode_pages( mostly_null.clone(), - LanceFileVersion::V2_3, + TestEncoding::StructuralSparse, sparse_metadata(), ) .await @@ -1638,8 +1651,7 @@ mod tests { ); let cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_page_sizes(vec![1]) .with_range(1..5) .with_indices(vec![0, 2, 5]); @@ -1652,9 +1664,13 @@ mod tests { async fn test_explicit_sparse_nested_page_boundaries_range_and_take() { let nested = deeply_nested(); let chunks = vec![nested.slice(0, 2), nested.slice(2, 3)]; - let pages = encode_chunks(chunks.clone(), LanceFileVersion::V2_3, sparse_metadata()) - .await - .unwrap(); + let pages = encode_chunks( + chunks.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); assert!(pages.len() >= 2, "expected multiple sparse pages"); let sparse = pages .iter() @@ -1681,8 +1697,7 @@ mod tests { })); let cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_page_sizes(vec![1]) .with_batch_size(2) .with_range(1..5) @@ -1703,8 +1718,7 @@ mod tests { Some(vec![false, true, true, true, true]), ); let cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_page_sizes(vec![1]) .with_range(0..5) .with_range(1..4) @@ -1712,9 +1726,13 @@ mod tests { .with_indices(vec![2, 3]); for array in [list, large_list] { - let pages = encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) - .await - .unwrap(); + let pages = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); assert!(pages.iter().map(sparse_layout).all(|layout| { layout.structural_layers.iter().any(|layer| { matches!( @@ -1730,9 +1748,13 @@ mod tests { #[tokio::test] async fn test_explicit_sparse_map_and_fixed_size_list_roundtrip() { let map = map_i32(); - let map_pages = encode_pages(map.clone(), LanceFileVersion::V2_3, sparse_metadata()) - .await - .unwrap(); + let map_pages = encode_pages( + map.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); assert!(map_pages.iter().map(sparse_layout).any(|layout| { layout.structural_layers.iter().any(|layer| { matches!( @@ -1742,17 +1764,20 @@ mod tests { }) })); let map_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_page_sizes(vec![1]) .with_range(1..4) .with_indices(vec![0, 2, 3]); check_round_trip_encoding_of_data(vec![map], &map_cases, sparse_metadata()).await; let fsl = fixed_size_list_struct(); - let fsl_pages = encode_pages(fsl.clone(), LanceFileVersion::V2_3, sparse_metadata()) - .await - .unwrap(); + let fsl_pages = encode_pages( + fsl.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); for page in &fsl_pages { let layout = sparse_layout(page); let outer_slots = layer_num_slots(layout.structural_layers.first().unwrap()); @@ -1770,8 +1795,7 @@ mod tests { .any(|layer| fixed_size_list_dimension(layer) == Some(2)) })); let fsl_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_page_sizes(vec![1]) .with_range(1..6) .with_indices(vec![0, 3, 5]); @@ -1781,9 +1805,13 @@ mod tests { #[tokio::test] async fn test_explicit_sparse_list_fixed_size_list_struct_roundtrip() { let array = list_fixed_size_list_struct(); - let pages = encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) - .await - .unwrap(); + let pages = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); assert!(pages.iter().map(sparse_layout).any(|layout| { let kinds = layout .structural_layers @@ -1803,8 +1831,7 @@ mod tests { } let cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_page_sizes(vec![1]) .with_range(1..4) .with_indices(vec![0, 2, 3]) @@ -1815,9 +1842,13 @@ mod tests { #[tokio::test] async fn test_explicit_sparse_serializes_semantic_list_forms() { let all_array = list_i32(vec![0, 2, 4, 6], None); - let all = encode_pages(all_array.clone(), LanceFileVersion::V2_3, sparse_metadata()) - .await - .unwrap(); + let all = encode_pages( + all_array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); let all_layer = list_layer(sparse_layout(&all[0])); assert!(matches!( all_layer.non_empty_positions.as_ref().unwrap().positions, @@ -1834,7 +1865,7 @@ mod tests { let range_array = list_i32(vec![0, 2, 4, 4, 4], None); let range = encode_pages( range_array.clone(), - LanceFileVersion::V2_3, + TestEncoding::StructuralSparse, sparse_metadata(), ) .await @@ -1860,7 +1891,7 @@ mod tests { let explicit_array = list_i32(vec![0, 1, 1, 4, 4, 6], None); let explicit = encode_pages( explicit_array.clone(), - LanceFileVersion::V2_3, + TestEncoding::StructuralSparse, sparse_metadata(), ) .await @@ -1881,8 +1912,7 @@ mod tests { assert_eq!(explicit[0].data.len(), 4); let cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_page_sizes(vec![1]) .with_range(1..3) .with_indices(vec![0, 2]); @@ -1895,7 +1925,7 @@ mod tests { async fn test_constant_layout_boundary_is_explicit() { let structural_only = encode_pages( list_i32(vec![0, 0, 0, 0], None), - LanceFileVersion::V2_3, + TestEncoding::StructuralSparse, sparse_metadata(), ) .await @@ -1906,28 +1936,33 @@ mod tests { )); let empty_struct = Arc::new(StructArray::new_empty_fields(3, None)) as ArrayRef; - let empty_struct_pages = - encode_pages(empty_struct, LanceFileVersion::V2_3, sparse_metadata()) - .await - .unwrap(); + let empty_struct_pages = encode_pages( + empty_struct, + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); assert!(matches!( page_layout(&empty_struct_pages[0]), pb21::page_layout::Layout::ConstantLayout(_) )); let all_null = Arc::new(Int32Array::from(vec![None, None, None])) as ArrayRef; - let all_null_pages = encode_pages(all_null, LanceFileVersion::V2_3, sparse_metadata()) - .await - .unwrap(); + let all_null_pages = + encode_pages(all_null, TestEncoding::StructuralSparse, sparse_metadata()) + .await + .unwrap(); assert!(matches!( page_layout(&all_null_pages[0]), pb21::page_layout::Layout::ConstantLayout(_) )); let constant = Arc::new(Int32Array::from(vec![7, 7, 7])) as ArrayRef; - let constant_pages = encode_pages(constant, LanceFileVersion::V2_3, sparse_metadata()) - .await - .unwrap(); + let constant_pages = + encode_pages(constant, TestEncoding::StructuralSparse, sparse_metadata()) + .await + .unwrap(); assert!(matches!( page_layout(&constant_pages[0]), pb21::page_layout::Layout::SparseLayout(_) @@ -1937,12 +1972,12 @@ mod tests { #[tokio::test] async fn test_default_and_explicit_dense_layouts_are_unchanged() { let array = Arc::new(Int32Array::from_iter_values(0..16)) as ArrayRef; - let v2_2_default = encode_pages(array.clone(), LanceFileVersion::V2_2, HashMap::new()) + let v2_2_default = encode_pages(array.clone(), TestEncoding::StructuralU32, HashMap::new()) .await .unwrap(); let v2_2_miniblock = encode_pages( array.clone(), - LanceFileVersion::V2_2, + TestEncoding::StructuralU32, structural_metadata(STRUCTURAL_ENCODING_MINIBLOCK), ) .await @@ -1957,9 +1992,13 @@ mod tests { assert_eq!(default.data, explicit.data); } - let v2_3_default = encode_pages(array.clone(), LanceFileVersion::V2_3, HashMap::new()) - .await - .unwrap(); + let v2_3_default = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + HashMap::new(), + ) + .await + .unwrap(); assert!(matches!( page_layout(&v2_3_default[0]), pb21::page_layout::Layout::MiniBlockLayout(_) @@ -1967,7 +2006,7 @@ mod tests { let v2_3_miniblock = encode_pages( array.clone(), - LanceFileVersion::V2_3, + TestEncoding::StructuralSparse, structural_metadata(STRUCTURAL_ENCODING_MINIBLOCK), ) .await @@ -1979,7 +2018,7 @@ mod tests { let v2_3_fullzip = encode_pages( array.clone(), - LanceFileVersion::V2_3, + TestEncoding::StructuralSparse, structural_metadata(STRUCTURAL_ENCODING_FULLZIP), ) .await @@ -1989,7 +2028,7 @@ mod tests { pb21::page_layout::Layout::FullZipLayout(_) )); - let v2_3_sparse = encode_pages(array, LanceFileVersion::V2_3, sparse_metadata()) + let v2_3_sparse = encode_pages(array, TestEncoding::StructuralSparse, sparse_metadata()) .await .unwrap(); assert!(matches!( @@ -2005,7 +2044,7 @@ mod tests { let within_budget = encode_pages( sparse_i32_list(4_096, 1_024), - LanceFileVersion::V2_3, + TestEncoding::StructuralSparse, HashMap::new(), ) .await @@ -2016,9 +2055,13 @@ mod tests { pb21::page_layout::Layout::MiniBlockLayout(_) )); - let automatic = encode_pages(array.clone(), LanceFileVersion::V2_3, HashMap::new()) - .await - .unwrap(); + let automatic = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + HashMap::new(), + ) + .await + .unwrap(); assert_eq!(automatic.len(), 1); assert!(matches!( page_layout(&automatic[0]), @@ -2027,7 +2070,7 @@ mod tests { let miniblock = encode_pages( array.clone(), - LanceFileVersion::V2_3, + TestEncoding::StructuralSparse, structural_metadata(STRUCTURAL_ENCODING_MINIBLOCK), ) .await @@ -2040,7 +2083,7 @@ mod tests { let fullzip = encode_pages( array.clone(), - LanceFileVersion::V2_3, + TestEncoding::StructuralSparse, structural_metadata(STRUCTURAL_ENCODING_FULLZIP), ) .await @@ -2051,21 +2094,25 @@ mod tests { pb21::page_layout::Layout::FullZipLayout(_) ))); - let sparse = encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) - .await - .unwrap(); + let sparse = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); assert_eq!(sparse.len(), 1); assert!(matches!( page_layout(&sparse[0]), pb21::page_layout::Layout::SparseLayout(_) )); - let v2_2_default = encode_pages(array.clone(), LanceFileVersion::V2_2, HashMap::new()) + let v2_2_default = encode_pages(array.clone(), TestEncoding::StructuralU32, HashMap::new()) .await .unwrap(); let v2_2_miniblock = encode_pages( array.clone(), - LanceFileVersion::V2_2, + TestEncoding::StructuralU32, structural_metadata(STRUCTURAL_ENCODING_MINIBLOCK), ) .await @@ -2077,8 +2124,7 @@ mod tests { } let cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_page_sizes(vec![1024 * 1024]) .with_batch_size(NUM_ROWS as u32) .with_range(1_999..2_002) @@ -2093,16 +2139,20 @@ mod tests { Arc::new(ArrowField::new("item", DataType::Int32, true)), ); - let automatic = encode_pages(array.clone(), LanceFileVersion::V2_3, HashMap::new()) - .await - .unwrap(); + let automatic = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + HashMap::new(), + ) + .await + .unwrap(); assert_eq!(automatic.len(), 1); assert!(matches!( page_layout(&automatic[0]), pb21::page_layout::Layout::SparseLayout(_) )); - let v2_2 = encode_pages(array.clone(), LanceFileVersion::V2_2, HashMap::new()) + let v2_2 = encode_pages(array.clone(), TestEncoding::StructuralU32, HashMap::new()) .await .unwrap(); assert_eq!(v2_2.len(), 1); @@ -2113,7 +2163,7 @@ mod tests { let miniblock_error = encode_pages( array.clone(), - LanceFileVersion::V2_3, + TestEncoding::StructuralSparse, structural_metadata(STRUCTURAL_ENCODING_MINIBLOCK), ) .await @@ -2126,7 +2176,7 @@ mod tests { let fullzip = encode_pages( array.clone(), - LanceFileVersion::V2_3, + TestEncoding::StructuralSparse, structural_metadata(STRUCTURAL_ENCODING_FULLZIP), ) .await @@ -2136,18 +2186,20 @@ mod tests { pb21::page_layout::Layout::FullZipLayout(_) )); - let explicit_sparse = - encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) - .await - .unwrap(); + let explicit_sparse = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap(); assert!(matches!( page_layout(&explicit_sparse[0]), pb21::page_layout::Layout::SparseLayout(_) )); let cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_range(0..1) .with_indices(vec![0]); check_round_trip_encoding_of_data(vec![array], &cases, HashMap::new()).await; @@ -2162,7 +2214,7 @@ mod tests { let dictionary_array = sparse_list_values(NUM_ROWS, 2_000, dictionary, dictionary_field); let dictionary_pages = encode_pages( dictionary_array.clone(), - LanceFileVersion::V2_3, + TestEncoding::StructuralSparse, HashMap::new(), ) .await @@ -2175,10 +2227,13 @@ mod tests { let (packed_values, packed_field) = variable_packed_struct_values(num_values); let packed_array = sparse_list_values(NUM_ROWS, 2_000, packed_values, packed_field); - let packed_pages = - encode_pages(packed_array.clone(), LanceFileVersion::V2_3, HashMap::new()) - .await - .unwrap(); + let packed_pages = encode_pages( + packed_array.clone(), + TestEncoding::StructuralSparse, + HashMap::new(), + ) + .await + .unwrap(); assert!(packed_pages.len() > 1); assert!(packed_pages.iter().all(|page| matches!( page_layout(page), @@ -2189,7 +2244,7 @@ mod tests { let unsplittable_packed = unsplittable_nested_list(packed_values, packed_field); let packed_fallback = encode_pages( unsplittable_packed.clone(), - LanceFileVersion::V2_3, + TestEncoding::StructuralSparse, HashMap::new(), ) .await @@ -2204,14 +2259,14 @@ mod tests { let unsplittable_dictionary = unsplittable_nested_list(dictionary, dictionary_field); let v2_2_error = encode_pages( unsplittable_dictionary.clone(), - LanceFileVersion::V2_2, + TestEncoding::StructuralU32, HashMap::new(), ) .await .unwrap_err(); let v2_3_error = encode_pages( unsplittable_dictionary, - LanceFileVersion::V2_3, + TestEncoding::StructuralSparse, HashMap::new(), ) .await @@ -2224,23 +2279,20 @@ mod tests { ); let cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_page_sizes(vec![1024 * 1024]) .with_batch_size(NUM_ROWS as u32) .with_range(1_999..2_002) .with_indices(vec![0, 2_000, NUM_ROWS as u64 - 1]); check_round_trip_encoding_of_data(vec![dictionary_array], &cases, HashMap::new()).await; let packed_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_page_sizes(vec![1024 * 1024]) .with_batch_size(NUM_ROWS as u32); check_round_trip_encoding_of_data(vec![packed_array], &packed_cases, HashMap::new()).await; let single_row_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_page_sizes(vec![1024 * 1024]) .with_range(0..1) .with_indices(vec![0]); @@ -2282,12 +2334,16 @@ mod tests { let item_field = Arc::new(ArrowField::new("item", values.data_type().clone(), true)); let array = unsplittable_nested_list(values, item_field); - let v2_2 = encode_pages(array.clone(), LanceFileVersion::V2_2, HashMap::new()) - .await - .unwrap(); - let v2_3 = encode_pages(array.clone(), LanceFileVersion::V2_3, HashMap::new()) + let v2_2 = encode_pages(array.clone(), TestEncoding::StructuralU32, HashMap::new()) .await .unwrap(); + let v2_3 = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + HashMap::new(), + ) + .await + .unwrap(); for pages in [&v2_2, &v2_3] { assert_eq!(pages.len(), 1, "unexpected {label} page count"); assert!( @@ -2299,18 +2355,20 @@ mod tests { ); } - let explicit_error = - encode_pages(array.clone(), LanceFileVersion::V2_3, sparse_metadata()) - .await - .unwrap_err(); + let explicit_error = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + sparse_metadata(), + ) + .await + .unwrap_err(); assert!( explicit_error.to_string().contains("too wide"), "explicit sparse should preserve the {label} value error: {explicit_error}" ); let cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_2) - .with_max_file_version(LanceFileVersion::V2_3) + .with_u32_structural_encodings() .with_range(0..1) .with_indices(vec![0]); check_round_trip_encoding_of_data(vec![array], &cases, HashMap::new()).await; @@ -2320,25 +2378,28 @@ mod tests { #[test] fn test_explicit_sparse_rejects_lance_2_2() { let array = Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef; - let Err(error) = create_encoder(&array, LanceFileVersion::V2_2, sparse_metadata()) else { + let Err(error) = create_encoder(&array, TestEncoding::StructuralU32, sparse_metadata()) + else { panic!("expected Lance 2.2 to reject explicit sparse encoding"); }; assert!( error .to_string() - .contains("requires Lance file format 2.3+") + .contains("not enabled by the selected file format") ); let structural_only = list_i32(vec![0, 0, 0], None); - let Err(error) = - create_encoder(&structural_only, LanceFileVersion::V2_2, sparse_metadata()) - else { + let Err(error) = create_encoder( + &structural_only, + TestEncoding::StructuralU32, + sparse_metadata(), + ) else { panic!("expected Lance 2.2 structural-only input to reject explicit sparse encoding"); }; assert!( error .to_string() - .contains("requires Lance file format 2.3+") + .contains("not enabled by the selected file format") ); } } diff --git a/rust/lance-encoding/src/encodings/logical/struct.rs b/rust/lance-encoding/src/encodings/logical/struct.rs index dc734da72b5..a038bcd2f48 100644 --- a/rust/lance-encoding/src/encodings/logical/struct.rs +++ b/rust/lance-encoding/src/encodings/logical/struct.rs @@ -600,7 +600,7 @@ impl FieldEncoder for StructFieldEncoder { let mut header = EncodedColumn::default(); header.final_pages.push(EncodedPage { data: Vec::new(), - description: PageEncoding::Legacy(pb::ArrayEncoding { + description: PageEncoding::Array(pb::ArrayEncoding { array_encoding: Some(pb::array_encoding::ArrayEncoding::Struct( pb::SimpleStruct {}, )), @@ -632,10 +632,7 @@ mod tests { use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow_schema::{DataType, Field, Fields}; - use crate::{ - testing::{TestCases, check_basic_random, check_round_trip_encoding_of_data}, - version::LanceFileVersion, - }; + use crate::testing::{TestCases, check_basic_random, check_round_trip_encoding_of_data}; #[test_log::test(tokio::test)] async fn test_simple_struct() { @@ -694,7 +691,7 @@ mod tests { Some(rows_validity), ); - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); + let test_cases = TestCases::default().with_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(rows)], &test_cases, HashMap::new()).await; } @@ -724,7 +721,7 @@ mod tests { ); check_round_trip_encoding_of_data( vec![Arc::new(struct_array)], - &TestCases::default().with_min_file_version(LanceFileVersion::V2_1), + &TestCases::default().with_structural_encodings(), HashMap::new(), ) .await; @@ -755,7 +752,7 @@ mod tests { ); check_round_trip_encoding_of_data( vec![Arc::new(struct_array)], - &TestCases::default().with_min_file_version(LanceFileVersion::V2_1), + &TestCases::default().with_structural_encodings(), HashMap::new(), ) .await; @@ -834,7 +831,7 @@ mod tests { check_round_trip_encoding_of_data( vec![Arc::new(list_array)], - &TestCases::default().with_min_file_version(LanceFileVersion::V2_2), + &TestCases::default().with_u32_structural_encodings(), HashMap::new(), ) .await; @@ -868,7 +865,7 @@ mod tests { .with_range(1..2) .with_indices(vec![0]) .with_indices(vec![1]) - .with_min_file_version(LanceFileVersion::V2_1), + .with_structural_encodings(), HashMap::new(), ) .await; @@ -916,7 +913,7 @@ mod tests { check_round_trip_encoding_of_data( vec![Arc::new(row_array)], - &TestCases::default().with_min_file_version(LanceFileVersion::V2_1), + &TestCases::default().with_structural_encodings(), HashMap::new(), ) .await; diff --git a/rust/lance-encoding/src/encodings/physical/binary.rs b/rust/lance-encoding/src/encodings/physical/binary.rs index d02cf2da693..86ac2f69934 100644 --- a/rust/lance-encoding/src/encodings/physical/binary.rs +++ b/rust/lance-encoding/src/encodings/physical/binary.rs @@ -542,22 +542,14 @@ mod tests { use rstest::rstest; use std::{collections::HashMap, sync::Arc, vec}; - use crate::{ - testing::{ - FnArrayGeneratorProvider, TestCases, check_basic_random, - check_round_trip_encoding_of_data, - }, - version::LanceFileVersion, + use crate::testing::{ + FnArrayGeneratorProvider, TestCases, check_basic_random, check_round_trip_encoding_of_data, }; #[test_log::test(tokio::test)] async fn test_utf8_binary() { let field = Field::new("", DataType::Utf8, false); - check_specific_random( - field, - TestCases::basic().with_min_file_version(LanceFileVersion::V2_1), - ) - .await; + check_specific_random(field, TestCases::basic().with_structural_encodings()).await; } #[rstest] @@ -592,7 +584,7 @@ mod tests { field_metadata.insert(COMPRESSION_META_KEY.to_string(), "fsst".into()); let field = Field::new("", data_type, true).with_metadata(field_metadata); // TODO (https://github.com/lance-format/lance/issues/4783) - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); + let test_cases = TestCases::default().with_structural_encodings(); check_specific_random(field, test_cases).await; } @@ -610,11 +602,7 @@ mod tests { ); field_metadata.insert(COMPRESSION_META_KEY.to_string(), "fsst".into()); let field = Field::new("", data_type, true).with_metadata(field_metadata); - check_specific_random( - field, - TestCases::basic().with_min_file_version(LanceFileVersion::V2_1), - ) - .await; + check_specific_random(field, TestCases::basic().with_structural_encodings()).await; } #[test_log::test(tokio::test)] @@ -801,7 +789,7 @@ mod tests { #[values(true, false)] with_nulls: bool, #[values(100, 500, 35000)] dict_size: u32, ) { - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); + let test_cases = TestCases::default().with_structural_encodings(); let strings = (0..dict_size) .map(|i| i.to_string()) .collect::>(); @@ -829,7 +817,7 @@ mod tests { let test_cases = TestCases::default() .with_expected_encoding("variable") - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); // Test both automatic selection and explicit configuration // 1. Test automatic binary encoding selection (small strings that won't trigger FSST) diff --git a/rust/lance-encoding/src/encodings/physical/bitpacking.rs b/rust/lance-encoding/src/encodings/physical/bitpacking.rs index be0b747e7dc..e5406b3f6fa 100644 --- a/rust/lance-encoding/src/encodings/physical/bitpacking.rs +++ b/rust/lance-encoding/src/encodings/physical/bitpacking.rs @@ -545,7 +545,6 @@ mod test { compression::MiniBlockDecompressor, data::{BlockInfo, DataBlock, FixedWidthDataBlock}, testing::{TestCases, check_round_trip_encoding_of_data}, - version::LanceFileVersion, }; #[rstest] @@ -569,7 +568,7 @@ mod test { #[test_log::test(tokio::test)] async fn test_miniblock_bitpack() { - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); + let test_cases = TestCases::default().with_structural_encodings(); let arrays = vec![ Arc::new(Int8Array::from(vec![100; 1024])) as Arc, @@ -608,7 +607,7 @@ mod test { // Test bitpacking encoding verification with varied small values that should trigger bitpacking let test_cases = TestCases::default() .with_expected_encoding("inline_bitpacking") - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); // Generate data with varied small values to avoid RLE // Mix different values but keep them small to trigger bitpacking @@ -632,7 +631,7 @@ mod test { let test_cases = TestCases::default() .with_expected_encoding("inline_bitpacking") - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); // Build 2048 values: first 1024 all zeros (bit_width=0), // next 1024 small varied values to avoid RLE and trigger bitpacking. diff --git a/rust/lance-encoding/src/encodings/physical/block.rs b/rust/lance-encoding/src/encodings/physical/block.rs index a1f5bdb3fdd..188af50c8a0 100644 --- a/rust/lance-encoding/src/encodings/physical/block.rs +++ b/rust/lance-encoding/src/encodings/physical/block.rs @@ -46,9 +46,20 @@ pub struct CompressionConfig { } impl CompressionConfig { - pub(crate) fn new(scheme: CompressionScheme, level: Option) -> Self { + /// Create a compression configuration for an encoding mechanism. + pub fn new(scheme: CompressionScheme, level: Option) -> Self { Self { scheme, level } } + + /// Return the selected compression scheme. + pub fn scheme(&self) -> CompressionScheme { + self.scheme + } + + /// Return the optional compression level. + pub fn level(&self) -> Option { + self.level + } } impl Default for CompressionConfig { @@ -765,7 +776,6 @@ mod tests { }, encodings::physical::block::lz4::Lz4BufferCompressor, testing::{FnArrayGeneratorProvider, TestCases, check_round_trip_encoding_generated}, - version::LanceFileVersion, }; #[test] @@ -809,7 +819,7 @@ mod tests { // Need to use large pages as small pages might be too small to compress .with_page_sizes(vec![1024 * 1024]) .with_expected_encoding("zstd") - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); // Can't use the default random provider because random data isn't compressible // and we will fallback to uncompressed encoding diff --git a/rust/lance-encoding/src/encodings/physical/fsst.rs b/rust/lance-encoding/src/encodings/physical/fsst.rs index 8c1fe4141df..14e7e0c0ccc 100644 --- a/rust/lance-encoding/src/encodings/physical/fsst.rs +++ b/rust/lance-encoding/src/encodings/physical/fsst.rs @@ -377,16 +377,13 @@ mod tests { use lance_datagen::{ByteCount, RowCount}; - use crate::{ - testing::{TestCases, check_round_trip_encoding_of_data}, - version::LanceFileVersion, - }; + use crate::testing::{TestCases, check_round_trip_encoding_of_data}; #[test_log::test(tokio::test)] async fn test_fsst() { let test_cases = TestCases::default() .with_expected_encoding("fsst") - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); // Generate data suitable for FSST (large strings, total size > 32KB) let arr = lance_datagen::gen_batch() diff --git a/rust/lance-encoding/src/encodings/physical/packed.rs b/rust/lance-encoding/src/encodings/physical/packed.rs index ad2221dffed..6b6fa27ba2b 100644 --- a/rust/lance-encoding/src/encodings/physical/packed.rs +++ b/rust/lance-encoding/src/encodings/physical/packed.rs @@ -18,7 +18,7 @@ use lance_core::{Error, Result, datatypes::Field}; use crate::{ buffer::LanceBuffer, compression::{ - DefaultCompressionStrategy, FixedPerValueDecompressor, MiniBlockDecompressor, + CompressionStrategy, FixedPerValueDecompressor, MiniBlockDecompressor, VariablePerValueDecompressor, }, data::{ @@ -282,12 +282,12 @@ impl VariablePackedFieldData { #[derive(Debug)] pub struct PackedStructVariablePerValueEncoder { - strategy: DefaultCompressionStrategy, + strategy: Arc, fields: Vec, } impl PackedStructVariablePerValueEncoder { - pub fn new(strategy: DefaultCompressionStrategy, fields: Vec) -> Self { + pub fn new(strategy: Arc, fields: Vec) -> Self { Self { strategy, fields } } } @@ -324,11 +324,7 @@ impl PerValueCompressor for PackedStructVariablePerValueEncoder { let mut field_metadata = Vec::with_capacity(self.fields.len()); for (field, child_block) in self.fields.iter().zip(struct_block.children) { - let compressor = crate::compression::CompressionStrategy::create_per_value( - &self.strategy, - field, - &child_block, - )?; + let compressor = self.strategy.create_per_value(field, &child_block)?; let (compressed, encoding) = compressor.compress(child_block)?; match compressed { PerValueDataBlock::Fixed(block) => { @@ -758,12 +754,13 @@ impl VariablePerValueDecompressor for PackedStructVariablePerValueDecompressor { mod tests { use super::*; use crate::{ - compression::CompressionStrategy, - compression::{DefaultCompressionStrategy, DefaultDecompressionStrategy}, + compression::DefaultDecompressionStrategy, + compression_config::CompressionParams, constants::PACKED_STRUCT_META_KEY, statistics::ComputeStat, - testing::{TestCases, check_round_trip_encoding_of_data}, - version::LanceFileVersion, + testing::{ + TestCases, TestEncoding, check_round_trip_encoding_of_data, test_compression_strategy, + }, }; use arrow_array::{ Array, ArrayRef, BinaryArray, Int32Array, Int64Array, LargeStringArray, StringArray, @@ -863,9 +860,9 @@ mod tests { let data_block = DataBlock::Struct(struct_block); let compression_strategy = - DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_2); + test_compression_strategy(TestEncoding::StructuralU32, CompressionParams::default()); let compressor = crate::compression::CompressionStrategy::create_per_value( - &compression_strategy, + compression_strategy.as_ref(), &struct_field, &data_block, )?; @@ -931,9 +928,9 @@ mod tests { let data_block = DataBlock::Struct(struct_block); let compression_strategy = - DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_2); + test_compression_strategy(TestEncoding::StructuralU32, CompressionParams::default()); let compressor = crate::compression::CompressionStrategy::create_per_value( - &compression_strategy, + compression_strategy.as_ref(), &struct_field, &data_block, )?; @@ -1013,7 +1010,7 @@ mod tests { ])); let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_2) + .with_u32_structural_encodings() .with_expected_encoding("variable_packed_struct"); check_round_trip_encoding_of_data(vec![array], &test_cases, meta).await; @@ -1052,9 +1049,9 @@ mod tests { let data_block = DataBlock::Struct(struct_block); let compression_strategy = - DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_2); + test_compression_strategy(TestEncoding::StructuralU32, CompressionParams::default()); let compressor = crate::compression::CompressionStrategy::create_per_value( - &compression_strategy, + compression_strategy.as_ref(), &struct_field, &data_block, )?; @@ -1135,7 +1132,7 @@ mod tests { }; let compression_strategy = - DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_1); + test_compression_strategy(TestEncoding::StructuralU16, CompressionParams::default()); let result = compression_strategy.create_per_value(&struct_field, &DataBlock::Struct(struct_block)); diff --git a/rust/lance-encoding/src/encodings/physical/rle.rs b/rust/lance-encoding/src/encodings/physical/rle.rs index b04716c2b44..eb71e82d9eb 100644 --- a/rust/lance-encoding/src/encodings/physical/rle.rs +++ b/rust/lance-encoding/src/encodings/physical/rle.rs @@ -3005,7 +3005,6 @@ mod tests { #[test_log::test(tokio::test)] async fn test_rle_encoding_verification() { use crate::testing::{TestCases, check_round_trip_encoding_of_data}; - use crate::version::LanceFileVersion; use arrow_array::{Array, Int32Array}; use lance_datagen::{ArrayGenerator, RowCount}; use std::collections::HashMap; @@ -3013,7 +3012,7 @@ mod tests { let test_cases = TestCases::default() .with_expected_encoding("rle") - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); // Test both explicit metadata and automatic selection // 1. Test with explicit RLE threshold metadata (also disable BSS) diff --git a/rust/lance-encoding/src/encodings/physical/value.rs b/rust/lance-encoding/src/encodings/physical/value.rs index 606f49b699a..412ce236dfe 100644 --- a/rust/lance-encoding/src/encodings/physical/value.rs +++ b/rust/lance-encoding/src/encodings/physical/value.rs @@ -784,7 +784,6 @@ mod tests { FnArrayGeneratorProvider, TestCases, check_basic_random, check_round_trip_encoding_generated, check_round_trip_encoding_of_data, }, - version::LanceFileVersion, }; use super::ValueEncoder; @@ -834,7 +833,7 @@ mod tests { .with_indices(vec![0, 1, 2]) .with_indices(vec![1]) .with_indices(vec![2]) - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); check_round_trip_encoding_of_data(vec![items], &test_cases, HashMap::default()).await; } @@ -845,7 +844,7 @@ mod tests { (0..5000).map(|i| if i % 2 == 0 { Some(i) } else { None }), )); - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); + let test_cases = TestCases::default().with_structural_encodings(); check_round_trip_encoding_of_data(vec![items], &test_cases, HashMap::default()).await; } @@ -877,7 +876,7 @@ mod tests { #[test_log::test(tokio::test)] async fn test_decimal128_dictionary_encoding() { - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); + let test_cases = TestCases::default().with_structural_encodings(); let decimals: Vec = (0..100).collect(); let repeated_strings: Vec<_> = decimals .iter() @@ -923,7 +922,7 @@ mod tests { let test_cases = TestCases::default() .with_page_sizes(vec![1000, 2000, 3000, 60000]) .with_batch_size(batch_size) - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); check_round_trip_encoding_of_data(data.clone(), &test_cases, HashMap::new()).await; } @@ -1092,7 +1091,7 @@ mod tests { let list_array = FixedSizeListArray::new(items_field, 2, items, Some(NullBuffer::new(list_nulls))); - let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); + let test_cases = TestCases::default().with_structural_encodings(); check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new()) .await; @@ -1110,7 +1109,7 @@ mod tests { let list_arr = ListArray::new(list_field, OffsetBuffer::new(offsets), Arc::new(fsl), None); let test_cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_1) + .with_structural_encodings() .with_batch_size(1); check_round_trip_encoding_of_data(vec![Arc::new(list_arr)], &test_cases, HashMap::new()) @@ -1217,7 +1216,7 @@ mod tests { let test_cases = TestCases::default() .with_expected_encoding("flat") - .with_min_file_version(LanceFileVersion::V2_1); + .with_structural_encodings(); // Test both explicit configuration and automatic fallback scenarios // 1. Test explicit "none" compression to force flat encoding diff --git a/rust/lance-encoding/src/lib.rs b/rust/lance-encoding/src/lib.rs index a58e0a14c59..7a0862ee3e3 100644 --- a/rust/lance-encoding/src/lib.rs +++ b/rust/lance-encoding/src/lib.rs @@ -8,6 +8,7 @@ use futures::{FutureExt, TryFutureExt, future::BoxFuture}; use lance_core::Result; +pub mod array_encoding; pub mod buffer; pub mod compression; pub mod compression_config; @@ -17,7 +18,6 @@ pub mod decoder; pub mod encoder; pub mod encodings; pub mod format; -pub mod previous; pub mod repdef; pub mod statistics; #[cfg(test)] diff --git a/rust/lance-encoding/src/previous.rs b/rust/lance-encoding/src/previous.rs deleted file mode 100644 index eb4e1bcdb0d..00000000000 --- a/rust/lance-encoding/src/previous.rs +++ /dev/null @@ -1,8 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The Lance Authors - -//! Legacy code for the 2.0 format that is no longer used in 2.1+ - -pub mod decoder; -pub mod encoder; -pub mod encodings; diff --git a/rust/lance-encoding/src/previous/decoder.rs b/rust/lance-encoding/src/previous/decoder.rs deleted file mode 100644 index bf32bea3d7c..00000000000 --- a/rust/lance-encoding/src/previous/decoder.rs +++ /dev/null @@ -1,132 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The Lance Authors - -use std::{collections::VecDeque, ops::Range}; - -use crate::decoder::{ - FilterExpression, NextDecodeTask, PriorityRange, ScheduledScanLine, SchedulerContext, -}; - -use arrow_schema::DataType; -use futures::future::BoxFuture; -use lance_core::{Error, Result}; - -pub trait SchedulingJob: std::fmt::Debug { - fn schedule_next( - &mut self, - context: &mut SchedulerContext, - priority: &dyn PriorityRange, - ) -> Result; - - fn num_rows(&self) -> u64; -} - -/// A scheduler for a field's worth of data -/// -/// Each field in a reader's output schema maps to one field scheduler. This scheduler may -/// map to more than one column. For example, one field of struct data may -/// cover many columns of child data. In fact, the entire file is treated as one -/// top-level struct field. -/// -/// The scheduler is responsible for calculating the necessary I/O. One schedule_range -/// request could trigger multiple batches of I/O across multiple columns. The scheduler -/// should emit decoders into the sink as quickly as possible. -/// -/// As soon as the scheduler encounters a batch of data that can decoded then the scheduler -/// should emit a decoder in the "unloaded" state. The decode stream will pull the decoder -/// and start decoding. -/// -/// The order in which decoders are emitted is important. Pages should be emitted in -/// row-major order allowing decode of complete rows as quickly as possible. -/// -/// The `FieldScheduler` should be stateless and `Send` and `Sync`. This is -/// because it might need to be shared. For example, a list page has a reference to -/// the field schedulers for its items column. This is shared with the follow-up I/O -/// task created when the offsets are loaded. -/// -/// See [`crate::decoder`] for more information -pub trait FieldScheduler: Send + Sync + std::fmt::Debug { - /// Called at the beginning of scheduling to initialize the scheduler - fn initialize<'a>( - &'a self, - filter: &'a FilterExpression, - context: &'a SchedulerContext, - ) -> BoxFuture<'a, Result<()>>; - /// Schedules I/O for the requested portions of the field. - /// - /// Note: `ranges` must be ordered and non-overlapping - /// TODO: Support unordered or overlapping ranges in file scheduler - fn schedule_ranges<'a>( - &'a self, - ranges: &[Range], - filter: &FilterExpression, - ) -> Result>; - /// The number of rows in this field - fn num_rows(&self) -> u64; -} - -#[derive(Debug)] -pub struct DecoderReady { - // The decoder that is ready to be decoded - pub decoder: Box, - // The path to the decoder, the first value is the column index - // following values, if present, are nested child indices - // - // For example, a path of [1, 1, 0] would mean to grab the second - // column, then the second child, and then the first child. - // - // It could represent x in the following schema: - // - // score: float64 - // points: struct - // color: string - // location: struct - // x: float64 - // - // Currently, only struct decoders have "children" although other - // decoders may at some point as well. List children are only - // handled through indirect I/O at the moment and so they don't - // need to be represented (yet) - pub path: VecDeque, -} - -/// A decoder for a field's worth of data -/// -/// The decoder is initially "unloaded" (doesn't have all its data). The [`Self::wait_for_loaded`] -/// method should be called to wait for the needed I/O data before attempting to decode -/// any further. -/// -/// Unlike the other decoder types it is assumed that `LogicalPageDecoder` is stateful -/// and only `Send`. This is why we don't need a `rows_to_skip` argument in [`Self::drain`] -pub trait LogicalPageDecoder: std::fmt::Debug + Send { - /// Add a newly scheduled child decoder - /// - /// The default implementation does not expect children and returns - /// an error. - fn accept_child(&mut self, _child: DecoderReady) -> Result<()> { - Err(Error::internal(format!( - "The decoder {:?} does not expect children but received a child", - self - ))) - } - /// Waits until at least `num_rows` have been loaded - fn wait_for_loaded(&'_ mut self, loaded_need: u64) -> BoxFuture<'_, Result<()>>; - /// The number of rows loaded so far - fn rows_loaded(&self) -> u64; - /// The number of rows that still need loading - fn rows_unloaded(&self) -> u64 { - self.num_rows() - self.rows_loaded() - } - /// The total number of rows in the field - fn num_rows(&self) -> u64; - /// The number of rows that have been drained so far - fn rows_drained(&self) -> u64; - /// The number of rows that are still available to drain - fn rows_left(&self) -> u64 { - self.num_rows() - self.rows_drained() - } - /// Creates a task to decode `num_rows` of data into an array - fn drain(&mut self, num_rows: u64) -> Result; - /// The data type of the decoded data - fn data_type(&self) -> &DataType; -} diff --git a/rust/lance-encoding/src/previous/encodings.rs b/rust/lance-encoding/src/previous/encodings.rs deleted file mode 100644 index 67a43993508..00000000000 --- a/rust/lance-encoding/src/previous/encodings.rs +++ /dev/null @@ -1,7 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The Lance Authors - -//! Legacy code for the 2.0 format that is no longer used in 2.1+ - -pub mod logical; -pub mod physical; diff --git a/rust/lance-encoding/src/testing.rs b/rust/lance-encoding/src/testing.rs index 9825cb5949d..8555a7951cc 100644 --- a/rust/lance-encoding/src/testing.rs +++ b/rust/lance-encoding/src/testing.rs @@ -21,27 +21,344 @@ use futures::{FutureExt, StreamExt, future::BoxFuture}; use log::{debug, info, trace}; use tokio::sync::mpsc::{self, UnboundedSender}; -use lance_core::{Result, utils::bit::pad_bytes}; +use lance_core::{Error, Result, datatypes::Field as LanceField, utils::bit::pad_bytes}; use lance_datagen::{ArrayGenerator, RowCount, Seed, array, gen_batch}; use crate::{ EncodingsIo, buffer::LanceBuffer, + compression::{ + BlockCompressor, CompressionStrategy, field_metadata_params, finalize_miniblock_compressor, + reject_packed_struct_per_value, try_bitpacking_block, try_bitpacking_miniblock, + try_byte_stream_split_miniblock, try_child_rle_miniblock, + try_fixed_packed_struct_miniblock, try_fixed_u8_rle_block, try_fixed_u8_rle_miniblock, + try_general_block, try_raw_block, try_raw_fixed_size_list_miniblock, + try_raw_fixed_width_miniblock, try_raw_per_value, try_uncompressed_fixed_width_miniblock, + try_variable_packed_struct_per_value, try_variable_rle_block, try_variable_width_miniblock, + try_variable_width_per_value, + }, + compression_config::{CompressionFieldParams, CompressionParams}, + data::DataBlock, decoder::{ ColumnInfo, DecodeBatchScheduler, DecoderMessage, DecoderPlugins, FilterExpression, PageInfo, create_decode_stream, }, encoder::{ ColumnIndexSequence, EncodedColumn, EncodedPage, EncodingOptions, FieldEncoder, - MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers, default_encoding_strategy, + FieldEncodingContext, FieldEncodingStrategy, MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers, + structural::{ + PrimitiveFieldEncoding, PrimitivePageEncoding, try_create_binary_blob, try_create_list, + try_create_map, try_create_struct, try_create_structural_blob, + try_create_structural_fixed_size_list, + }, }, + encodings::logical::primitive::{fullzip::PerValueCompressor, miniblock::MiniBlockCompressor}, repdef::RepDefBuilder, - version::LanceFileVersion, }; const MAX_PAGE_BYTES: u64 = 32 * 1024 * 1024; const TEST_ALIGNMENT: usize = MIN_PAGE_BUFFER_ALIGNMENT as usize; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TestEncoding { + Array, + StructuralU16, + StructuralU32, + StructuralSparse, +} + +impl TestEncoding { + fn all() -> impl Iterator { + [ + Self::Array, + Self::StructuralU16, + Self::StructuralU32, + Self::StructuralSparse, + ] + .into_iter() + } + + fn is_structural(self) -> bool { + self != Self::Array + } +} + +impl std::fmt::Display for TestEncoding { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Array => write!(f, "array"), + Self::StructuralU16 => write!(f, "structural-u16"), + Self::StructuralU32 => write!(f, "structural-u32"), + Self::StructuralSparse => write!(f, "structural-sparse"), + } + } +} + +#[derive(Debug, Clone)] +struct TestCompressionStrategy { + encoding: TestEncoding, + params: CompressionParams, +} + +impl TestCompressionStrategy { + fn field_params(&self, field: &LanceField) -> CompressionFieldParams { + let mut params = self + .params + .get_field_params(&field.name, &field.data_type()); + let mut metadata = field_metadata_params(field); + if self.encoding == TestEncoding::StructuralU16 + && metadata + .minichunk_size + .is_some_and(|size| size >= 32 * 1024) + { + metadata.minichunk_size = None; + } + params.merge(&metadata); + params + } +} + +impl CompressionStrategy for TestCompressionStrategy { + fn create_miniblock_compressor( + &self, + field: &LanceField, + data: &DataBlock, + ) -> Result> { + let params = self.field_params(field); + let compressor = + if let Some(compressor) = try_uncompressed_fixed_width_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = try_byte_stream_split_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = match self.encoding { + TestEncoding::StructuralSparse => try_child_rle_miniblock(data, ¶ms), + TestEncoding::Array | TestEncoding::StructuralU16 | TestEncoding::StructuralU32 => { + try_fixed_u8_rle_miniblock(data, ¶ms) + } + } { + compressor + } else if let Some(compressor) = try_bitpacking_miniblock(data) { + compressor + } else if let Some(compressor) = try_raw_fixed_width_miniblock(data) { + compressor + } else if let Some(compressor) = try_variable_width_miniblock(field, data, ¶ms)? { + compressor + } else if let Some(compressor) = try_fixed_packed_struct_miniblock(data)? { + compressor + } else if let Some(compressor) = try_raw_fixed_size_list_miniblock(data) { + compressor + } else { + return Err(lance_core::Error::not_supported_source( + format!( + "Mini-block compression not yet supported for block type {}", + data.name() + ) + .into(), + )); + }; + finalize_miniblock_compressor(data, compressor, ¶ms) + } + + fn create_per_value( + &self, + field: &LanceField, + data: &DataBlock, + ) -> Result> { + let params = self.field_params(field); + if let Some(compressor) = try_raw_per_value(data) { + return Ok(compressor); + } + let packed = match self.encoding { + TestEncoding::Array | TestEncoding::StructuralU16 => { + reject_packed_struct_per_value(field, data)? + } + TestEncoding::StructuralU32 | TestEncoding::StructuralSparse => { + try_variable_packed_struct_per_value(Arc::new(self.clone()), field, data)? + } + }; + if let Some(compressor) = packed { + return Ok(compressor); + } + if let Some(compressor) = try_variable_width_per_value(field, data, ¶ms)? { + return Ok(compressor); + } + Err(lance_core::Error::not_supported_source( + format!( + "Per-value compression not yet supported for block type {}", + data.name() + ) + .into(), + )) + } + + fn create_block_compressor( + &self, + field: &LanceField, + data: &DataBlock, + ) -> Result<(Box, CompressiveEncoding)> { + let params = self.field_params(field); + let rle = match self.encoding { + TestEncoding::Array | TestEncoding::StructuralU16 => None, + TestEncoding::StructuralU32 => try_fixed_u8_rle_block(data, ¶ms)?, + TestEncoding::StructuralSparse => try_variable_rle_block(data, ¶ms)?, + }; + if let Some(compressor) = rle { + return Ok(compressor); + } + if let Some(compressor) = try_bitpacking_block(data) { + return Ok(compressor); + } + if matches!( + self.encoding, + TestEncoding::StructuralU32 | TestEncoding::StructuralSparse + ) && let Some(compressor) = try_general_block(data, ¶ms)? + { + return Ok(compressor); + } + if let Some(compressor) = try_raw_block(data) { + return Ok(compressor); + } + Err(lance_core::Error::not_supported_source( + format!( + "Block compression not yet supported for block type {}", + data.name() + ) + .into(), + )) + } +} + +pub fn test_compression_strategy( + encoding: TestEncoding, + params: CompressionParams, +) -> Arc { + Arc::new(TestCompressionStrategy { encoding, params }) +} + +#[derive(Debug)] +struct TestFieldEncodingStrategy { + encoding: TestEncoding, + primitive: PrimitiveFieldEncoding, +} + +impl FieldEncodingStrategy for TestFieldEncodingStrategy { + fn create_field_encoder( + &self, + field: &LanceField, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, + ) -> Result> { + if let Some(encoder) = + try_create_binary_blob(&self.primitive, field, column_index, context)? + { + return Ok(encoder); + } + if self.encoding != TestEncoding::StructuralU16 + && let Some(encoder) = + try_create_structural_blob(&self.primitive, field, column_index, context)? + { + return Ok(encoder); + } + if field.is_blob() { + return Err(Error::invalid_input_source( + format!( + "Blob encoding is not available for field '{}' with data type {}", + field.name, + field.data_type() + ) + .into(), + )); + } + if self.encoding != TestEncoding::StructuralU16 { + if let Some(encoder) = try_create_map(field, column_index, context)? { + return Ok(encoder); + } + if let Some(encoder) = + try_create_structural_fixed_size_list(field, column_index, context)? + { + return Ok(encoder); + } + } + if let Some(encoder) = self.primitive.try_create(field, column_index, context)? { + return Ok(encoder); + } + if self.encoding == TestEncoding::StructuralU16 { + if matches!( + field.data_type(), + DataType::FixedSizeList(item, _) + if matches!(item.data_type(), DataType::Struct(_)) + ) { + return Err(Error::not_supported_source( + "FixedSizeList is not enabled by the selected file format".into(), + )); + } + if matches!(field.data_type(), DataType::Map(_, _)) { + return Err(Error::not_supported_source( + "Map data type is not enabled by the selected file format".into(), + )); + } + } + if let Some(encoder) = try_create_list(field, column_index, context)? { + return Ok(encoder); + } + if let Some(encoder) = try_create_struct(field, column_index, context)? { + return Ok(encoder); + } + Err(Error::not_supported_source( + format!( + "{} has no field encoding for '{}' with data type {}", + self.encoding, + field.name, + field.data_type() + ) + .into(), + )) + } +} + +pub fn test_encoding_strategy(encoding: TestEncoding) -> Box { + if encoding == TestEncoding::Array { + return Box::new(crate::array_encoding::ArrayFieldEncodingStrategy::new()); + } + + let compression = test_compression_strategy(encoding, CompressionParams::default()); + let page_encodings = match encoding { + TestEncoding::StructuralU16 => vec![ + PrimitivePageEncoding::reject_sparse(), + PrimitivePageEncoding::dense_u16(compression), + ], + TestEncoding::StructuralU32 => vec![ + PrimitivePageEncoding::reject_sparse(), + PrimitivePageEncoding::constant(), + PrimitivePageEncoding::dense_u32(compression), + ], + TestEncoding::StructuralSparse => vec![ + PrimitivePageEncoding::sparse(compression.clone()), + PrimitivePageEncoding::constant(), + PrimitivePageEncoding::dense_u32(compression), + ], + TestEncoding::Array => unreachable!(), + }; + Box::new(TestFieldEncodingStrategy { + encoding, + primitive: PrimitiveFieldEncoding::new(page_encodings), + }) +} + +pub fn create_test_field_encoder( + strategy: &dyn FieldEncodingStrategy, + field: &lance_core::datatypes::Field, + column_index: &mut ColumnIndexSequence, + options: &EncodingOptions, +) -> Result> { + let context = FieldEncodingContext { + strategy, + options, + root_field_metadata: &field.metadata, + }; + strategy.create_field_encoder(field, column_index, &context) +} + #[derive(Debug)] pub(crate) struct SimulatedScheduler { data: Bytes, @@ -342,24 +659,22 @@ pub async fn check_round_trip_encoding_generated( let lance_field = lance_core::datatypes::Field::try_from(&field).unwrap(); for page_size in test_cases.page_sizes.iter().copied() { debug!("Testing random data with a page size of {}", page_size); - let encoder_factory = |version: LanceFileVersion| { - let encoding_strategy = default_encoding_strategy(version); + let encoder_factory = |encoding: TestEncoding| { + let encoding_strategy = test_encoding_strategy(encoding); let mut column_index_seq = ColumnIndexSequence::default(); let encoding_options = EncodingOptions { max_page_bytes: MAX_PAGE_BYTES, cache_bytes_per_column: page_size, keep_original_array: true, buffer_alignment: MIN_PAGE_BUFFER_ALIGNMENT, - version, }; - encoding_strategy - .create_field_encoder( - encoding_strategy.as_ref(), - &lance_field, - &mut column_index_seq, - &encoding_options, - ) - .unwrap() + create_test_field_encoder( + encoding_strategy.as_ref(), + &lance_field, + &mut column_index_seq, + &encoding_options, + ) + .unwrap() }; check_round_trip_random( @@ -372,9 +687,9 @@ pub async fn check_round_trip_encoding_generated( } } -fn supports_nulls(data_type: &DataType, version: LanceFileVersion) -> bool { +fn supports_nulls(data_type: &DataType, encoding: TestEncoding) -> bool { if let DataType::Struct(fields) = data_type { - if version == LanceFileVersion::V2_0 { + if encoding == TestEncoding::Array { // 2.0 doesn't support nullability for structs false } else if fields.is_empty() { @@ -389,7 +704,7 @@ fn supports_nulls(data_type: &DataType, version: LanceFileVersion) -> bool { } } -type EncodingVerificationFn = dyn Fn(&[EncodedColumn], &LanceFileVersion); +type EncodingVerificationFn = dyn Fn(&[EncodedColumn], &TestEncoding); // The default will just test the full read #[derive(Clone)] @@ -400,8 +715,7 @@ pub struct TestCases { skip_validation: bool, max_page_size: Option, page_sizes: Vec, - min_file_version: Option, - max_file_version: Option, + encodings: Vec, verify_encoding: Option>, expected_encoding: Option>, } @@ -415,8 +729,7 @@ impl Default for TestCases { skip_validation: false, max_page_size: None, page_sizes: vec![4096, 1024 * 1024], - min_file_version: None, - max_file_version: None, + encodings: TestEncoding::all().collect(), verify_encoding: None, expected_encoding: None, } @@ -459,16 +772,40 @@ impl TestCases { self } - pub fn with_min_file_version(mut self, version: LanceFileVersion) -> Self { - self.min_file_version = Some(version); + pub fn with_encoding(mut self, encoding: TestEncoding) -> Self { + self.encodings = vec![encoding]; self } - pub fn with_max_file_version(mut self, version: LanceFileVersion) -> Self { - self.max_file_version = Some(version); + pub fn with_encodings(mut self, encodings: impl IntoIterator) -> Self { + self.encodings = encodings.into_iter().collect(); self } + pub fn with_structural_encodings(self) -> Self { + self.with_encodings([ + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + TestEncoding::StructuralSparse, + ]) + } + + pub fn with_u32_structural_encodings(self) -> Self { + self.with_encodings([TestEncoding::StructuralU32, TestEncoding::StructuralSparse]) + } + + pub fn with_dense_encodings(self) -> Self { + self.with_encodings([ + TestEncoding::Array, + TestEncoding::StructuralU16, + TestEncoding::StructuralU32, + ]) + } + + pub fn with_array_and_u16_encodings(self) -> Self { + self.with_encodings([TestEncoding::Array, TestEncoding::StructuralU16]) + } + pub fn with_page_sizes(mut self, page_sizes: Vec) -> Self { self.page_sizes = page_sizes; self @@ -483,22 +820,8 @@ impl TestCases { self.max_page_size.unwrap_or(MAX_PAGE_BYTES) } - fn get_versions(&self) -> Vec { - LanceFileVersion::iter_non_legacy() - .filter(|v| { - if let Some(min_file_version) = &self.min_file_version - && v < min_file_version - { - return false; - } - if let Some(max_file_version) = &self.max_file_version - && v > max_file_version - { - return false; - } - true - }) - .collect() + fn encodings(&self) -> impl Iterator + '_ { + self.encodings.iter().copied() } pub fn with_verify_encoding(mut self, verify_encoding: Arc) -> Self { @@ -506,9 +829,9 @@ impl TestCases { self } - fn verify_encoding(&self, encoding: &[EncodedColumn], version: &LanceFileVersion) { + fn verify_encoding(&self, columns: &[EncodedColumn], encoding: &TestEncoding) { if let Some(verify_encoding) = self.verify_encoding.as_ref() { - verify_encoding(encoding, version); + verify_encoding(columns, encoding); } } @@ -704,8 +1027,8 @@ fn verify_page_encoding( return Ok(()); } } - PageEncoding::Legacy(_) => { - // We don't need to care about legacy. + PageEncoding::Array(_) => { + // We don't need to care about the v2.0 array encoding. } } @@ -748,28 +1071,26 @@ pub async fn check_round_trip_encoding_of_data_with_expected( let mut field = Field::new("", example_data.data_type().clone(), true); field = field.with_metadata(metadata); let lance_field = lance_core::datatypes::Field::try_from(&field).unwrap(); - for file_version in test_cases.get_versions() { + for encoding in test_cases.encodings() { for page_size in test_cases.page_sizes.iter() { - let encoding_strategy = default_encoding_strategy(file_version); + let encoding_strategy = test_encoding_strategy(encoding); let mut column_index_seq = ColumnIndexSequence::default(); let encoding_options = EncodingOptions { cache_bytes_per_column: *page_size, max_page_bytes: test_cases.get_max_page_size(), keep_original_array: true, buffer_alignment: MIN_PAGE_BUFFER_ALIGNMENT, - version: file_version, }; - let encoder = encoding_strategy - .create_field_encoder( - encoding_strategy.as_ref(), - &lance_field, - &mut column_index_seq, - &encoding_options, - ) - .unwrap(); + let encoder = create_test_field_encoder( + encoding_strategy.as_ref(), + &lance_field, + &mut column_index_seq, + &encoding_options, + ) + .unwrap(); info!( - "Testing round trip encoding of data with file version {} and page size {}", - file_version, page_size + "Testing round trip encoding of data with test encoding {} and page size {}", + encoding, page_size ); check_round_trip_encoding_inner( encoder, @@ -777,7 +1098,7 @@ pub async fn check_round_trip_encoding_of_data_with_expected( data.clone(), expected_override.clone(), test_cases, - file_version, + encoding, ) .await } @@ -850,7 +1171,7 @@ async fn check_round_trip_encoding_inner( data: Vec>, expected_override: Option>, test_cases: &TestCases, - file_version: LanceFileVersion, + encoding: TestEncoding, ) { let mut writer = SimulatedWriter::new(encoder.num_columns()); @@ -891,7 +1212,7 @@ async fn check_round_trip_encoding_inner( log_page(&encoded_page); // For V2.1, verify encoding in the page if expected - if file_version >= LanceFileVersion::V2_1 + if encoding.is_structural() && let Some(ref expected) = test_cases.expected_encoding { verify_page_encoding(&encoded_page, expected, encoded_page.column_idx as usize) @@ -913,7 +1234,7 @@ async fn check_round_trip_encoding_inner( log_page(&encoded_page); // For V2.1, verify encoding in the page if expected - if file_version >= LanceFileVersion::V2_1 + if encoding.is_structural() && let Some(ref expected) = test_cases.expected_encoding { verify_page_encoding(&encoded_page, expected, encoded_page.column_idx as usize) @@ -925,7 +1246,7 @@ async fn check_round_trip_encoding_inner( let mut external_buffers = writer.new_external_buffers(); let encoded_columns = encoder.finish(&mut external_buffers).await.unwrap(); - test_cases.verify_encoding(&encoded_columns, &file_version); + test_cases.verify_encoding(&encoded_columns, &encoding); for buffer in external_buffers.take_buffers() { writer.write_lance_buffer(buffer); } @@ -978,7 +1299,7 @@ async fn check_round_trip_encoding_inner( let expected_data = expected_override.clone().or_else(|| concat_data.clone()); - let is_structural_encoding = file_version >= LanceFileVersion::V2_1; + let is_structural_encoding = encoding.is_structural(); let decode_field = if is_structural_encoding { let mut lance_field = lance_core::datatypes::Field::try_from(field).unwrap(); @@ -1118,20 +1439,20 @@ const NUM_RANDOM_ROWS: u32 = 10000; /// /// To test specific test cases use the async fn check_round_trip_random( - encoder_factory: impl Fn(LanceFileVersion) -> Box, + encoder_factory: impl Fn(TestEncoding) -> Box, field: Field, array_generator_provider: Box, test_cases: &TestCases, ) { for null_rate in [None, Some(0.5), Some(1.0)] { for use_slicing in [false, true] { - for file_version in test_cases.get_versions() { + for encoding in test_cases.encodings() { if null_rate != Some(1.0) && matches!(field.data_type(), DataType::Null) { continue; } let field = if null_rate.is_some() { - if !supports_nulls(field.data_type(), file_version) { + if !supports_nulls(field.data_type(), encoding) { continue; } field.clone().with_nullable(true) @@ -1189,8 +1510,8 @@ async fn check_round_trip_random( } info!( - "Testing version {} with {} rows divided across {} batches for {} rows per batch with null_rate={:?} and use_slicing={}", - file_version, + "Testing encoding {} with {} rows divided across {} batches for {} rows per batch with null_rate={:?} and use_slicing={}", + encoding, NUM_RANDOM_ROWS, num_ingest_batches, rows_per_batch, @@ -1198,12 +1519,12 @@ async fn check_round_trip_random( use_slicing ); check_round_trip_encoding_inner( - encoder_factory(file_version), + encoder_factory(encoding), &field, data, None, test_cases, - file_version, + encoding, ) .await } diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index 5549457a058..9ba9e346f93 100644 --- a/rust/lance-file/src/reader.rs +++ b/rust/lance-file/src/reader.rs @@ -1214,7 +1214,7 @@ impl FileReader { let num_rows = page.length; let encoding = match file_version { LanceFileVersion::V2_0 => { - PageEncoding::Legacy(Self::fetch_encoding::( + PageEncoding::Array(Self::fetch_encoding::( page.encoding.as_ref().ok_or_else(|| { Error::invalid_input_source( format!( @@ -2725,10 +2725,10 @@ mod tests { use lance_encoding::{ constants::{STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_SPARSE}, decoder::{ - DecodeBatchScheduler, DecoderPlugins, FilterExpression, PageEncoding, ReadBatchTask, - decode_batch, + DecodeBatchScheduler, DecoderPlugins, EncodedBatchLayout, FilterExpression, + PageEncoding, ReadBatchTask, decode_batch, }, - encoder::{EncodedBatch, EncodingOptions, default_encoding_strategy, encode_batch}, + encoder::{EncodedBatch, EncodingOptions, encode_batch}, format::pb21, version::LanceFileVersion, }; @@ -3121,10 +3121,9 @@ mod tests { max_page_bytes: 32 * 1024 * 1024, keep_original_array: true, buffer_alignment: 64, - version, }; - let encoding_strategy = default_encoding_strategy(version); + let encoding_strategy = crate::versions::v2_0::encoding_strategy(); let encoded_batch = encode_batch( &data, @@ -3146,7 +3145,7 @@ mod tests { &FilterExpression::no_filter(), Arc::::default(), false, - version, + EncodedBatchLayout::Array, None, ) .await @@ -3165,7 +3164,7 @@ mod tests { &FilterExpression::no_filter(), Arc::::default(), false, - version, + EncodedBatchLayout::Array, None, ) .await diff --git a/rust/lance-file/src/versions/mod.rs b/rust/lance-file/src/versions/mod.rs index 613eb15d52d..7eb1ca7f357 100644 --- a/rust/lance-file/src/versions/mod.rs +++ b/rust/lance-file/src/versions/mod.rs @@ -4,3 +4,7 @@ //! Exact Lance file-version implementations. pub mod v1; +pub mod v2_0; +pub mod v2_1; +pub mod v2_2; +pub mod v2_3; diff --git a/rust/lance-file/src/versions/v2_0/mod.rs b/rust/lance-file/src/versions/v2_0/mod.rs new file mode 100644 index 00000000000..25e20816bd8 --- /dev/null +++ b/rust/lance-file/src/versions/v2_0/mod.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Lance v2.0 encoding composition. + +use std::sync::Arc; + +use lance_encoding::{array_encoding::ArrayFieldEncodingStrategy, encoder::FieldEncodingStrategy}; + +/// Compose the v2.0 field encoding mechanisms. +pub fn encoding_strategy() -> Arc { + Arc::new(ArrayFieldEncodingStrategy::new()) +} diff --git a/rust/lance-file/src/versions/v2_1/compression.rs b/rust/lance-file/src/versions/v2_1/compression.rs new file mode 100644 index 00000000000..051e131e9a9 --- /dev/null +++ b/rust/lance-file/src/versions/v2_1/compression.rs @@ -0,0 +1,127 @@ +use lance_core::{Error, Result, datatypes::Field}; +use lance_encoding::{ + compression::{ + BlockCompressor, CompressionStrategy, field_metadata_params, finalize_miniblock_compressor, + reject_packed_struct_per_value, try_bitpacking_block, try_bitpacking_miniblock, + try_byte_stream_split_miniblock, try_fixed_packed_struct_miniblock, + try_fixed_u8_rle_miniblock, try_raw_block, try_raw_fixed_size_list_miniblock, + try_raw_fixed_width_miniblock, try_raw_per_value, try_uncompressed_fixed_width_miniblock, + try_variable_width_miniblock, try_variable_width_per_value, + }, + compression_config::{CompressionFieldParams, CompressionParams}, + data::DataBlock, + encodings::logical::primitive::{fullzip::PerValueCompressor, miniblock::MiniBlockCompressor}, + format::pb21::CompressiveEncoding, +}; + +#[derive(Debug, Clone)] +pub(super) struct Strategy { + params: CompressionParams, +} + +impl Strategy { + pub(super) fn new(params: CompressionParams) -> Self { + Self { params } + } + + fn field_params(&self, field: &Field) -> CompressionFieldParams { + let mut params = self + .params + .get_field_params(&field.name, &field.data_type()); + let mut metadata = field_metadata_params(field); + if metadata + .minichunk_size + .is_some_and(|size| size >= 32 * 1024) + { + log::warn!( + "minichunk_size '{}' is too large for the selected u16 miniblock layout, using default", + metadata.minichunk_size.unwrap() + ); + metadata.minichunk_size = None; + } + params.merge(&metadata); + params + } +} + +impl CompressionStrategy for Strategy { + fn create_miniblock_compressor( + &self, + field: &Field, + data: &DataBlock, + ) -> Result> { + let params = self.field_params(field); + let compressor = + if let Some(compressor) = try_uncompressed_fixed_width_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = try_byte_stream_split_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = try_fixed_u8_rle_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = try_bitpacking_miniblock(data) { + compressor + } else if let Some(compressor) = try_raw_fixed_width_miniblock(data) { + compressor + } else if let Some(compressor) = try_variable_width_miniblock(field, data, ¶ms)? { + compressor + } else if let Some(compressor) = try_fixed_packed_struct_miniblock(data)? { + compressor + } else if let Some(compressor) = try_raw_fixed_size_list_miniblock(data) { + compressor + } else { + return Err(Error::not_supported_source( + format!( + "Mini-block compression not yet supported for block type {}", + data.name() + ) + .into(), + )); + }; + finalize_miniblock_compressor(data, compressor, ¶ms) + } + + fn create_per_value( + &self, + field: &Field, + data: &DataBlock, + ) -> Result> { + let params = self.field_params(field); + if let Some(compressor) = try_raw_per_value(data) { + return Ok(compressor); + } + if let Some(compressor) = reject_packed_struct_per_value(field, data)? { + return Ok(compressor); + } + if let Some(compressor) = try_variable_width_per_value(field, data, ¶ms)? { + return Ok(compressor); + } + Err(Error::not_supported_source( + format!( + "Per-value compression not yet supported for block type {}", + data.name() + ) + .into(), + )) + } + + fn create_block_compressor( + &self, + field: &Field, + data: &DataBlock, + ) -> Result<(Box, CompressiveEncoding)> { + let _params = self.field_params(field); + if let Some(compressor) = try_bitpacking_block(data) { + return Ok(compressor); + } + if let Some(compressor) = try_raw_block(data) { + return Ok(compressor); + } + Err(Error::not_supported_source( + format!( + "Block compression not yet supported for block type {}", + data.name() + ) + .into(), + )) + } +} diff --git a/rust/lance-file/src/versions/v2_1/mod.rs b/rust/lance-file/src/versions/v2_1/mod.rs new file mode 100644 index 00000000000..4a8e5072c79 --- /dev/null +++ b/rust/lance-file/src/versions/v2_1/mod.rs @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Lance v2.1 encoding composition. + +use std::sync::Arc; + +use lance_core::{Error, Result, datatypes::Field}; +use lance_encoding::{ + compression_config::CompressionParams, + encoder::{ + ColumnIndexSequence, FieldEncoder, FieldEncodingContext, FieldEncodingStrategy, + structural::{ + PrimitiveFieldEncoding, PrimitivePageEncoding, try_create_binary_blob, try_create_list, + try_create_struct, + }, + }, +}; + +mod compression; + +#[derive(Debug)] +struct FieldStrategy { + primitive: PrimitiveFieldEncoding, +} + +impl FieldEncodingStrategy for FieldStrategy { + fn create_field_encoder( + &self, + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, + ) -> Result> { + if let Some(encoder) = + try_create_binary_blob(&self.primitive, field, column_index, context)? + { + return Ok(encoder); + } + if field.is_blob() { + return Err(Error::invalid_input_source( + format!( + "Blob encoding is not available for field '{}' with data type {}", + field.name, + field.data_type() + ) + .into(), + )); + } + if let Some(encoder) = self.primitive.try_create(field, column_index, context)? { + return Ok(encoder); + } + if matches!( + field.data_type(), + arrow_schema::DataType::FixedSizeList(item, _) + if matches!(item.data_type(), arrow_schema::DataType::Struct(_)) + ) { + return Err(Error::not_supported_source( + "FixedSizeList is not enabled by the selected file format".into(), + )); + } + if matches!(field.data_type(), arrow_schema::DataType::Map(_, _)) { + return Err(Error::not_supported_source( + "Map data type is not enabled by the selected file format".into(), + )); + } + if let Some(encoder) = try_create_list(field, column_index, context)? { + return Ok(encoder); + } + if let Some(encoder) = try_create_struct(field, column_index, context)? { + return Ok(encoder); + } + Err(Error::not_supported_source( + format!( + "Lance v2.1 has no field encoding for '{}' with data type {}", + field.name, + field.data_type() + ) + .into(), + )) + } +} + +/// Compose the v2.1 field encoding mechanisms. +pub fn encoding_strategy(params: CompressionParams) -> Arc { + let compression = Arc::new(compression::Strategy::new(params)); + Arc::new(FieldStrategy { + primitive: PrimitiveFieldEncoding::new([ + PrimitivePageEncoding::reject_sparse(), + PrimitivePageEncoding::dense_u16(compression), + ]), + }) +} diff --git a/rust/lance-file/src/versions/v2_2/compression.rs b/rust/lance-file/src/versions/v2_2/compression.rs new file mode 100644 index 00000000000..a374adc4bec --- /dev/null +++ b/rust/lance-file/src/versions/v2_2/compression.rs @@ -0,0 +1,127 @@ +use std::sync::Arc; + +use lance_core::{Error, Result, datatypes::Field}; +use lance_encoding::{ + compression::{ + BlockCompressor, CompressionStrategy, field_metadata_params, finalize_miniblock_compressor, + try_bitpacking_block, try_bitpacking_miniblock, try_byte_stream_split_miniblock, + try_fixed_packed_struct_miniblock, try_fixed_u8_rle_block, try_fixed_u8_rle_miniblock, + try_general_block, try_raw_block, try_raw_fixed_size_list_miniblock, + try_raw_fixed_width_miniblock, try_raw_per_value, try_uncompressed_fixed_width_miniblock, + try_variable_packed_struct_per_value, try_variable_width_miniblock, + try_variable_width_per_value, + }, + compression_config::{CompressionFieldParams, CompressionParams}, + data::DataBlock, + encodings::logical::primitive::{fullzip::PerValueCompressor, miniblock::MiniBlockCompressor}, + format::pb21::CompressiveEncoding, +}; + +#[derive(Debug, Clone)] +pub(super) struct Strategy { + params: CompressionParams, +} + +impl Strategy { + pub(super) fn new(params: CompressionParams) -> Self { + Self { params } + } + + fn field_params(&self, field: &Field) -> CompressionFieldParams { + let mut params = self + .params + .get_field_params(&field.name, &field.data_type()); + params.merge(&field_metadata_params(field)); + params + } +} + +impl CompressionStrategy for Strategy { + fn create_miniblock_compressor( + &self, + field: &Field, + data: &DataBlock, + ) -> Result> { + let params = self.field_params(field); + let compressor = + if let Some(compressor) = try_uncompressed_fixed_width_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = try_byte_stream_split_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = try_fixed_u8_rle_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = try_bitpacking_miniblock(data) { + compressor + } else if let Some(compressor) = try_raw_fixed_width_miniblock(data) { + compressor + } else if let Some(compressor) = try_variable_width_miniblock(field, data, ¶ms)? { + compressor + } else if let Some(compressor) = try_fixed_packed_struct_miniblock(data)? { + compressor + } else if let Some(compressor) = try_raw_fixed_size_list_miniblock(data) { + compressor + } else { + return Err(Error::not_supported_source( + format!( + "Mini-block compression not yet supported for block type {}", + data.name() + ) + .into(), + )); + }; + finalize_miniblock_compressor(data, compressor, ¶ms) + } + + fn create_per_value( + &self, + field: &Field, + data: &DataBlock, + ) -> Result> { + let params = self.field_params(field); + if let Some(compressor) = try_raw_per_value(data) { + return Ok(compressor); + } + if let Some(compressor) = + try_variable_packed_struct_per_value(Arc::new(self.clone()), field, data)? + { + return Ok(compressor); + } + if let Some(compressor) = try_variable_width_per_value(field, data, ¶ms)? { + return Ok(compressor); + } + Err(Error::not_supported_source( + format!( + "Per-value compression not yet supported for block type {}", + data.name() + ) + .into(), + )) + } + + fn create_block_compressor( + &self, + field: &Field, + data: &DataBlock, + ) -> Result<(Box, CompressiveEncoding)> { + let params = self.field_params(field); + if let Some(compressor) = try_fixed_u8_rle_block(data, ¶ms)? { + return Ok(compressor); + } + if let Some(compressor) = try_bitpacking_block(data) { + return Ok(compressor); + } + if let Some(compressor) = try_general_block(data, ¶ms)? { + return Ok(compressor); + } + if let Some(compressor) = try_raw_block(data) { + return Ok(compressor); + } + Err(Error::not_supported_source( + format!( + "Block compression not yet supported for block type {}", + data.name() + ) + .into(), + )) + } +} diff --git a/rust/lance-file/src/versions/v2_2/mod.rs b/rust/lance-file/src/versions/v2_2/mod.rs new file mode 100644 index 00000000000..8bf0b251add --- /dev/null +++ b/rust/lance-file/src/versions/v2_2/mod.rs @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Lance v2.2 encoding composition. + +use std::sync::Arc; + +use lance_core::{Error, Result, datatypes::Field}; +use lance_encoding::{ + compression_config::CompressionParams, + encoder::{ + ColumnIndexSequence, FieldEncoder, FieldEncodingContext, FieldEncodingStrategy, + structural::{ + PrimitiveFieldEncoding, PrimitivePageEncoding, try_create_binary_blob, try_create_list, + try_create_map, try_create_struct, try_create_structural_blob, + try_create_structural_fixed_size_list, + }, + }, +}; + +mod compression; + +#[derive(Debug)] +struct FieldStrategy { + primitive: PrimitiveFieldEncoding, +} + +impl FieldEncodingStrategy for FieldStrategy { + fn create_field_encoder( + &self, + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, + ) -> Result> { + if let Some(encoder) = + try_create_binary_blob(&self.primitive, field, column_index, context)? + { + return Ok(encoder); + } + if let Some(encoder) = + try_create_structural_blob(&self.primitive, field, column_index, context)? + { + return Ok(encoder); + } + if field.is_blob() { + return Err(Error::invalid_input_source( + format!( + "Blob encoding is not available for field '{}' with data type {}", + field.name, + field.data_type() + ) + .into(), + )); + } + if let Some(encoder) = try_create_map(field, column_index, context)? { + return Ok(encoder); + } + if let Some(encoder) = try_create_structural_fixed_size_list(field, column_index, context)? + { + return Ok(encoder); + } + if let Some(encoder) = self.primitive.try_create(field, column_index, context)? { + return Ok(encoder); + } + if let Some(encoder) = try_create_list(field, column_index, context)? { + return Ok(encoder); + } + if let Some(encoder) = try_create_struct(field, column_index, context)? { + return Ok(encoder); + } + Err(Error::not_supported_source( + format!( + "Lance v2.2 has no field encoding for '{}' with data type {}", + field.name, + field.data_type() + ) + .into(), + )) + } +} + +/// Compose the v2.2 field encoding mechanisms. +pub fn encoding_strategy(params: CompressionParams) -> Arc { + let compression = Arc::new(compression::Strategy::new(params)); + Arc::new(FieldStrategy { + primitive: PrimitiveFieldEncoding::new([ + PrimitivePageEncoding::reject_sparse(), + PrimitivePageEncoding::constant(), + PrimitivePageEncoding::dense_u32(compression), + ]), + }) +} diff --git a/rust/lance-file/src/versions/v2_3/compression.rs b/rust/lance-file/src/versions/v2_3/compression.rs new file mode 100644 index 00000000000..108f0fa63a4 --- /dev/null +++ b/rust/lance-file/src/versions/v2_3/compression.rs @@ -0,0 +1,127 @@ +use std::sync::Arc; + +use lance_core::{Error, Result, datatypes::Field}; +use lance_encoding::{ + compression::{ + BlockCompressor, CompressionStrategy, field_metadata_params, finalize_miniblock_compressor, + try_bitpacking_block, try_bitpacking_miniblock, try_byte_stream_split_miniblock, + try_child_rle_miniblock, try_fixed_packed_struct_miniblock, try_general_block, + try_raw_block, try_raw_fixed_size_list_miniblock, try_raw_fixed_width_miniblock, + try_raw_per_value, try_uncompressed_fixed_width_miniblock, + try_variable_packed_struct_per_value, try_variable_rle_block, try_variable_width_miniblock, + try_variable_width_per_value, + }, + compression_config::{CompressionFieldParams, CompressionParams}, + data::DataBlock, + encodings::logical::primitive::{fullzip::PerValueCompressor, miniblock::MiniBlockCompressor}, + format::pb21::CompressiveEncoding, +}; + +#[derive(Debug, Clone)] +pub(super) struct Strategy { + params: CompressionParams, +} + +impl Strategy { + pub(super) fn new(params: CompressionParams) -> Self { + Self { params } + } + + fn field_params(&self, field: &Field) -> CompressionFieldParams { + let mut params = self + .params + .get_field_params(&field.name, &field.data_type()); + params.merge(&field_metadata_params(field)); + params + } +} + +impl CompressionStrategy for Strategy { + fn create_miniblock_compressor( + &self, + field: &Field, + data: &DataBlock, + ) -> Result> { + let params = self.field_params(field); + let compressor = + if let Some(compressor) = try_uncompressed_fixed_width_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = try_byte_stream_split_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = try_child_rle_miniblock(data, ¶ms) { + compressor + } else if let Some(compressor) = try_bitpacking_miniblock(data) { + compressor + } else if let Some(compressor) = try_raw_fixed_width_miniblock(data) { + compressor + } else if let Some(compressor) = try_variable_width_miniblock(field, data, ¶ms)? { + compressor + } else if let Some(compressor) = try_fixed_packed_struct_miniblock(data)? { + compressor + } else if let Some(compressor) = try_raw_fixed_size_list_miniblock(data) { + compressor + } else { + return Err(Error::not_supported_source( + format!( + "Mini-block compression not yet supported for block type {}", + data.name() + ) + .into(), + )); + }; + finalize_miniblock_compressor(data, compressor, ¶ms) + } + + fn create_per_value( + &self, + field: &Field, + data: &DataBlock, + ) -> Result> { + let params = self.field_params(field); + if let Some(compressor) = try_raw_per_value(data) { + return Ok(compressor); + } + if let Some(compressor) = + try_variable_packed_struct_per_value(Arc::new(self.clone()), field, data)? + { + return Ok(compressor); + } + if let Some(compressor) = try_variable_width_per_value(field, data, ¶ms)? { + return Ok(compressor); + } + Err(Error::not_supported_source( + format!( + "Per-value compression not yet supported for block type {}", + data.name() + ) + .into(), + )) + } + + fn create_block_compressor( + &self, + field: &Field, + data: &DataBlock, + ) -> Result<(Box, CompressiveEncoding)> { + let params = self.field_params(field); + if let Some(compressor) = try_variable_rle_block(data, ¶ms)? { + return Ok(compressor); + } + if let Some(compressor) = try_bitpacking_block(data) { + return Ok(compressor); + } + if let Some(compressor) = try_general_block(data, ¶ms)? { + return Ok(compressor); + } + if let Some(compressor) = try_raw_block(data) { + return Ok(compressor); + } + Err(Error::not_supported_source( + format!( + "Block compression not yet supported for block type {}", + data.name() + ) + .into(), + )) + } +} diff --git a/rust/lance-file/src/versions/v2_3/mod.rs b/rust/lance-file/src/versions/v2_3/mod.rs new file mode 100644 index 00000000000..beceff92cad --- /dev/null +++ b/rust/lance-file/src/versions/v2_3/mod.rs @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Lance v2.3 encoding composition. + +use std::sync::Arc; + +use lance_core::{Error, Result, datatypes::Field}; +use lance_encoding::{ + compression_config::CompressionParams, + encoder::{ + ColumnIndexSequence, FieldEncoder, FieldEncodingContext, FieldEncodingStrategy, + structural::{ + PrimitiveFieldEncoding, PrimitivePageEncoding, try_create_binary_blob, try_create_list, + try_create_map, try_create_struct, try_create_structural_blob, + try_create_structural_fixed_size_list, + }, + }, +}; + +mod compression; + +#[derive(Debug)] +struct FieldStrategy { + primitive: PrimitiveFieldEncoding, +} + +impl FieldEncodingStrategy for FieldStrategy { + fn create_field_encoder( + &self, + field: &Field, + column_index: &mut ColumnIndexSequence, + context: &FieldEncodingContext<'_>, + ) -> Result> { + if let Some(encoder) = + try_create_binary_blob(&self.primitive, field, column_index, context)? + { + return Ok(encoder); + } + if let Some(encoder) = + try_create_structural_blob(&self.primitive, field, column_index, context)? + { + return Ok(encoder); + } + if field.is_blob() { + return Err(Error::invalid_input_source( + format!( + "Blob encoding is not available for field '{}' with data type {}", + field.name, + field.data_type() + ) + .into(), + )); + } + if let Some(encoder) = try_create_map(field, column_index, context)? { + return Ok(encoder); + } + if let Some(encoder) = try_create_structural_fixed_size_list(field, column_index, context)? + { + return Ok(encoder); + } + if let Some(encoder) = self.primitive.try_create(field, column_index, context)? { + return Ok(encoder); + } + if let Some(encoder) = try_create_list(field, column_index, context)? { + return Ok(encoder); + } + if let Some(encoder) = try_create_struct(field, column_index, context)? { + return Ok(encoder); + } + Err(Error::not_supported_source( + format!( + "Lance v2.3 has no field encoding for '{}' with data type {}", + field.name, + field.data_type() + ) + .into(), + )) + } +} + +/// Compose the v2.3 field encoding mechanisms. +pub fn encoding_strategy(params: CompressionParams) -> Arc { + let compression = Arc::new(compression::Strategy::new(params)); + Arc::new(FieldStrategy { + primitive: PrimitiveFieldEncoding::new([ + PrimitivePageEncoding::sparse(compression.clone()), + PrimitivePageEncoding::constant(), + PrimitivePageEncoding::dense_u32(compression), + ]), + }) +} diff --git a/rust/lance-file/src/writer.rs b/rust/lance-file/src/writer.rs index a9ee95a3e16..0719f69820e 100644 --- a/rust/lance-file/src/writer.rs +++ b/rust/lance-file/src/writer.rs @@ -15,10 +15,11 @@ use futures::stream::FuturesOrdered; use lance_core::datatypes::{Field, Schema as LanceSchema}; use lance_core::utils::bit::pad_bytes; use lance_core::{Error, Result}; +use lance_encoding::compression_config::CompressionParams; use lance_encoding::decoder::PageEncoding; use lance_encoding::encoder::{ BatchEncoder, EncodeTask, EncodedBatch, EncodedPage, EncodingOptions, FieldEncoder, - FieldEncodingStrategy, OutOfLineBuffers, default_encoding_strategy, + FieldEncodingStrategy, OutOfLineBuffers, }; use lance_encoding::repdef::RepDefBuilder; use lance_io::object_store::ObjectStore; @@ -37,6 +38,7 @@ use crate::format::pb; use crate::format::pbfile; use crate::format::pbfile::DirectEncoding; use crate::version::{ConcreteFileVersion, LanceFileVersion}; +use crate::versions; /// Pages buffers are aligned to 64 bytes pub(crate) const PAGE_BUFFER_ALIGNMENT: usize = 64; @@ -49,6 +51,39 @@ const PAD_BUFFER: [u8; PAGE_BUFFER_ALIGNMENT] = [72; PAGE_BUFFER_ALIGNMENT]; const MAX_PAGE_BYTES: usize = 32 * 1024 * 1024; const ENV_LANCE_FILE_WRITER_MAX_PAGE_BYTES: &str = "LANCE_FILE_WRITER_MAX_PAGE_BYTES"; +#[cfg(test)] +fn encoding_strategy_with_params( + version: LanceFileVersion, + params: CompressionParams, +) -> Result> { + match version.resolve() { + LanceFileVersion::Legacy | LanceFileVersion::V2_0 => Err(Error::invalid_input( + "Compression parameters are only supported in Lance file version 2.1 and later", + )), + LanceFileVersion::V2_1 => Ok(versions::v2_1::encoding_strategy(params)), + LanceFileVersion::V2_2 => Ok(versions::v2_2::encoding_strategy(params)), + LanceFileVersion::V2_3 => Ok(versions::v2_3::encoding_strategy(params)), + LanceFileVersion::Stable | LanceFileVersion::Next => { + unreachable!("resolved file-version selector must be exact") + } + } +} + +fn encoding_strategy(version: LanceFileVersion) -> Arc { + match version.resolve() { + LanceFileVersion::Legacy => { + panic!("legacy v1 files require versions::v1::writer::FileWriter") + } + LanceFileVersion::V2_0 => versions::v2_0::encoding_strategy(), + LanceFileVersion::V2_1 => versions::v2_1::encoding_strategy(CompressionParams::default()), + LanceFileVersion::V2_2 => versions::v2_2::encoding_strategy(CompressionParams::default()), + LanceFileVersion::V2_3 => versions::v2_3::encoding_strategy(CompressionParams::default()), + LanceFileVersion::Stable | LanceFileVersion::Next => { + unreachable!("resolved file-version selector must be exact") + } + } +} + /// Summary of a completed Lance file write. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct FileWriteSummary { @@ -343,7 +378,7 @@ impl FileWriter { Self::do_write_buffer(&mut self.writer, &buffer).await?; } let encoded_encoding = match encoded_page.description { - PageEncoding::Legacy(array_encoding) => Any::from_msg(&array_encoding)?.encode_to_vec(), + PageEncoding::Array(array_encoding) => Any::from_msg(&array_encoding)?.encode_to_vec(), PageEncoding::Structural(page_layout) => Any::from_msg(&page_layout)?.encode_to_vec(), }; let page = pbfile::column_metadata::Page { @@ -457,17 +492,17 @@ impl FileWriter { schema.validate()?; let keep_original_array = self.options.keep_original_array.unwrap_or(false); - let encoding_strategy = self.options.encoding_strategy.clone().unwrap_or_else(|| { - let version = self.version(); - default_encoding_strategy(version).into() - }); + let encoding_strategy = self + .options + .encoding_strategy + .clone() + .unwrap_or_else(|| encoding_strategy(self.version())); let encoding_options = EncodingOptions { cache_bytes_per_column, max_page_bytes, keep_original_array, buffer_alignment: PAGE_BUFFER_ALIGNMENT as u64, - version: self.version(), }; let encoder = BatchEncoder::try_new(&schema, encoding_strategy.as_ref(), &encoding_options)?; @@ -1008,7 +1043,7 @@ fn concat_lance_footer( .iter() .map(|page_info| { let encoded_encoding = match &page_info.encoding { - PageEncoding::Legacy(array_encoding) => { + PageEncoding::Array(array_encoding) => { Any::from_msg(array_encoding)?.encode_to_vec() } PageEncoding::Structural(page_layout) => { @@ -1984,15 +2019,12 @@ mod tests { // In a real implementation, you could add other compression types here // Build encoding strategy with compression parameters - let encoding_strategy = lance_encoding::encoder::default_encoding_strategy_with_params( - LanceFileVersion::V2_1, - params, - ) - .unwrap(); + let encoding_strategy = + super::encoding_strategy_with_params(LanceFileVersion::V2_1, params).unwrap(); // Configure file writer options let options = FileWriterOptions { - encoding_strategy: Some(Arc::from(encoding_strategy)), + encoding_strategy: Some(encoding_strategy), format_version: Some(LanceFileVersion::V2_1), max_page_bytes: Some(64 * 1024), // 64KB pages ..Default::default() @@ -2132,14 +2164,11 @@ mod tests { // Create encoding strategy that will read from field metadata let params = CompressionParams::new(); - let encoding_strategy = lance_encoding::encoder::default_encoding_strategy_with_params( - LanceFileVersion::V2_1, - params, - ) - .unwrap(); + let encoding_strategy = + super::encoding_strategy_with_params(LanceFileVersion::V2_1, params).unwrap(); let options = FileWriterOptions { - encoding_strategy: Some(Arc::from(encoding_strategy)), + encoding_strategy: Some(encoding_strategy), format_version: Some(LanceFileVersion::V2_1), ..Default::default() }; @@ -2231,14 +2260,11 @@ mod tests { // Create encoding strategy that will read from field metadata let params = CompressionParams::new(); - let encoding_strategy = lance_encoding::encoder::default_encoding_strategy_with_params( - LanceFileVersion::V2_1, - params, - ) - .unwrap(); + let encoding_strategy = + super::encoding_strategy_with_params(LanceFileVersion::V2_1, params).unwrap(); let options = FileWriterOptions { - encoding_strategy: Some(Arc::from(encoding_strategy)), + encoding_strategy: Some(encoding_strategy), format_version: Some(LanceFileVersion::V2_1), ..Default::default() }; diff --git a/rust/lance/benches/s3_file_reader_diagnostics.rs b/rust/lance/benches/s3_file_reader_diagnostics.rs index 1762048d495..9f5e7c29976 100644 --- a/rust/lance/benches/s3_file_reader_diagnostics.rs +++ b/rust/lance/benches/s3_file_reader_diagnostics.rs @@ -574,7 +574,7 @@ fn projection_name(columns: &Option>) -> String { fn page_layout_kind(encoding: &PageEncoding) -> &'static str { match encoding { - PageEncoding::Legacy(_) => "legacy", + PageEncoding::Array(_) => "legacy", PageEncoding::Structural(layout) => match layout.layout.as_ref() { Some(pb21::page_layout::Layout::MiniBlockLayout(_)) => "miniblock", Some(pb21::page_layout::Layout::ConstantLayout(_)) => "constant", diff --git a/rust/lance/src/dataset/optimize/binary_copy.rs b/rust/lance/src/dataset/optimize/binary_copy.rs index d22b69ab6fc..a5bdc8d29e0 100644 --- a/rust/lance/src/dataset/optimize/binary_copy.rs +++ b/rust/lance/src/dataset/optimize/binary_copy.rs @@ -394,7 +394,7 @@ pub async fn rewrite_files_binary_copy( let encoding = if page.encoding.is_structural() { PageEncoding::Structural(page.encoding.as_structural().clone()) } else { - PageEncoding::Legacy(page.encoding.as_legacy().clone()) + PageEncoding::Array(page.encoding.as_array().clone()) }; // `priority` acts as the global row offset for this page, ensuring // downstream iterators maintain the correct logical order across @@ -545,7 +545,7 @@ async fn flush_footer( .iter() .map(|page_info| { let encoded_encoding = match &page_info.encoding { - PageEncoding::Legacy(array_encoding) => { + PageEncoding::Array(array_encoding) => { Any::from_msg(array_encoding)?.encode_to_vec() } PageEncoding::Structural(page_layout) => {