feat: support large_string/large_binary in lance format v2.1 (#3967)

This PR addresses #3350. Now largeutf8 and largeBinary(64bits) are
support in lance format v2.1

---------

Co-authored-by: Weston Pace <weston.pace@gmail.com>

Author: HaochengLIU

Cargo.lock

@@ -105,6 +105,7 @@ allow = [
     "ISC",
     "BSD-2-Clause",
     "BSD-3-Clause",
+    "BSL-1.0",
     "0BSD",
     "OpenSSL",
     "Zlib",

deny.toml

@@ -37,6 +37,7 @@ rand.workspace = true
 snafu.workspace = true
 tokio.workspace = true
 tracing.workspace = true
+xxhash-rust = { version = "0.8.15", features = ["xxh3"] }
 zstd.workspace = true
 bytemuck = "1.14"
 arrayref = "0.3.7"

python/Cargo.lock

@@ -20,10 +20,11 @@ use std::{
 };
 
 use arrow::array::{ArrayData, ArrayDataBuilder, AsArray};
-use arrow_array::{new_empty_array, new_null_array, Array, ArrayRef, UInt64Array};
-use arrow_buffer::{ArrowNativeType, BooleanBuffer, BooleanBufferBuilder, NullBuffer};
+use arrow_array::{new_empty_array, new_null_array, Array, ArrayRef, OffsetSizeTrait, UInt64Array};
+use arrow_buffer::{
+    ArrowNativeType, BooleanBuffer, BooleanBufferBuilder, NullBuffer, ScalarBuffer,
+};
 use arrow_schema::DataType;
-use bytemuck::try_cast_slice;
 use lance_arrow::DataTypeExt;
 use snafu::location;
 
@@ -252,43 +253,42 @@ impl FixedWidthDataBlock {
 }
 
 #[derive(Debug)]
-pub struct VariableWidthDataBlockBuilder {
-    offsets: Vec<u32>,
+pub struct VariableWidthDataBlockBuilder<T: OffsetSizeTrait> {
+    offsets: Vec<T>,
     bytes: Vec<u8>,
 }
 
-impl VariableWidthDataBlockBuilder {
+impl<T: OffsetSizeTrait> VariableWidthDataBlockBuilder<T> {
     fn new(estimated_size_bytes: u64) -> Self {
         Self {
-            offsets: vec![0u32],
+            offsets: vec![T::from_usize(0).unwrap()],
             bytes: Vec::with_capacity(estimated_size_bytes as usize),
         }
     }
 }
 
-impl DataBlockBuilderImpl for VariableWidthDataBlockBuilder {
+impl<T: OffsetSizeTrait> DataBlockBuilderImpl for VariableWidthDataBlockBuilder<T> {
     fn append(&mut self, data_block: &DataBlock, selection: Range<u64>) {
         let block = data_block.as_variable_width_ref().unwrap();
-        assert!(block.bits_per_offset == 32);
+        assert!(block.bits_per_offset == T::get_byte_width() as u8 * 8);
 
-        let offsets: &[u32] = try_cast_slice(&block.offsets)
-            .expect("cast from a bits_per_offset=32 `VariableWidthDataBlock's offsets field field to &[32] should be fine.");
+        let offsets: ScalarBuffer<T> = block.offsets.try_clone().unwrap().borrow_to_typed_slice();
 
         let start_offset = offsets[selection.start as usize];
         let end_offset = offsets[selection.end as usize];
         let mut previous_len = self.bytes.len();
 
         self.bytes
-            .extend_from_slice(&block.data[start_offset as usize..end_offset as usize]);
+            .extend_from_slice(&block.data[start_offset.as_usize()..end_offset.as_usize()]);
 
         self.offsets.extend(
             offsets[selection.start as usize..selection.end as usize]
                 .iter()
                 .zip(&offsets[selection.start as usize + 1..=selection.end as usize])
                 .map(|(&current, &next)| {
                     let this_value_len = next - current;
-                    previous_len += this_value_len as usize;
-                    previous_len as u32
+                    previous_len += this_value_len.as_usize();
+                    T::from_usize(previous_len).unwrap()
                 }),
         );
     }
@@ -298,7 +298,7 @@ impl DataBlockBuilderImpl for VariableWidthDataBlockBuilder {
         DataBlock::VariableWidth(VariableWidthBlock {
             data: LanceBuffer::Owned(self.bytes),
             offsets: LanceBuffer::reinterpret_vec(self.offsets),
-            bits_per_offset: 32,
+            bits_per_offset: T::get_byte_width() as u8 * 8,
             num_values,
             block_info: BlockInfo::new(),
         })
@@ -1085,7 +1085,13 @@ impl DataBlock {
             }
             Self::VariableWidth(inner) => {
                 if inner.bits_per_offset == 32 {
-                    Box::new(VariableWidthDataBlockBuilder::new(estimated_size_bytes))
+                    Box::new(VariableWidthDataBlockBuilder::<i32>::new(
+                        estimated_size_bytes,
+                    ))
+                } else if inner.bits_per_offset == 64 {
+                    Box::new(VariableWidthDataBlockBuilder::<i64>::new(
+                        estimated_size_bytes,
+                    ))
                 } else {
                     todo!()
                 }

rust/lance-encoding/Cargo.toml

@@ -517,9 +517,9 @@ impl DecompressorStrategy for CoreDecompressorStrategy {
             pb::array_encoding::ArrayEncoding::InlineBitpacking(description) => {
                 Ok(Box::new(InlineBitpacking::from_description(description)))
             }
-            pb::array_encoding::ArrayEncoding::Variable(_) => {
-                Ok(Box::new(BinaryMiniBlockDecompressor::default()))
-            }
+            pb::array_encoding::ArrayEncoding::Variable(description) => Ok(Box::new(
+                BinaryMiniBlockDecompressor::from_variable(description),
+            )),
             pb::array_encoding::ArrayEncoding::Fsst(description) => {
                 Ok(Box::new(FsstMiniBlockDecompressor::new(description)))
             }
@@ -786,6 +786,7 @@ impl CoreFieldDecoderStrategy {
             let scheduler = Box::new(StructuralPrimitiveFieldScheduler::try_new(
                 column_info.as_ref(),
                 self.decompressor_strategy.as_ref(),
+                32, // Use 32-bit offsets for standard primitive fields
             )?);
 
             // advance to the next top level column
@@ -800,6 +801,7 @@ impl CoreFieldDecoderStrategy {
                     let scheduler = Box::new(StructuralPrimitiveFieldScheduler::try_new(
                         column_info.as_ref(),
                         self.decompressor_strategy.as_ref(),
+                        32, // Use 32-bit offsets for packed struct fields
                     )?);
 
                     // advance to the next top level column
@@ -825,6 +827,17 @@ impl CoreFieldDecoderStrategy {
                 let scheduler = Box::new(StructuralPrimitiveFieldScheduler::try_new(
                     column_info.as_ref(),
                     self.decompressor_strategy.as_ref(),
+                    32, // Use 32-bit offsets for Binary/Utf8
+                )?);
+                column_infos.next_top_level();
+                Ok(scheduler)
+            }
+            DataType::LargeBinary | DataType::LargeUtf8 => {
+                let column_info = column_infos.expect_next()?;
+                let scheduler = Box::new(StructuralPrimitiveFieldScheduler::try_new(
+                    column_info.as_ref(),
+                    self.decompressor_strategy.as_ref(),
+                    64, // Use 64-bit offsets for LargeBinary/LargeUtf8
                 )?);
                 column_infos.next_top_level();
                 Ok(scheduler)

rust/lance-encoding/src/data.rs

@@ -851,8 +851,11 @@ impl CompressionStrategy for CoreArrayEncodingStrategy {
                     } else {
                         Ok(Box::new(BinaryMiniBlockEncoder::default()))
                     }
+                } else if variable_width_data.bits_per_offset == 64 {
+                    // TODO: Support FSSTMiniBlockEncoder
+                    Ok(Box::new(BinaryMiniBlockEncoder::default()))
                 } else {
-                    todo!("Implement MiniBlockCompression for VariableWidth DataBlock with 64 bits offsets.")
+                    todo!("Implement MiniBlockCompression for VariableWidth DataBlock with {} bits offsets.", variable_width_data.bits_per_offset)
                 }
             }
             DataBlock::Struct(struct_data_block) => {

rust/lance-encoding/src/decoder.rs

@@ -1328,6 +1328,7 @@ impl MiniBlockScheduler {
             .collect::<Vec<_>>();
         let value_decompressor = decompressors
             .create_miniblock_decompressor(layout.value_compression.as_ref().unwrap())?;
+
         let dictionary = if let Some(dictionary_encoding) = layout.dictionary.as_ref() {
             let num_dictionary_items = layout.num_dictionary_items;
             match dictionary_encoding.array_encoding.as_ref().unwrap() {
@@ -2780,6 +2781,7 @@ impl StructuralPrimitiveFieldScheduler {
     pub fn try_new(
         column_info: &ColumnInfo,
         decompressors: &dyn DecompressorStrategy,
+        bits_per_offset: u8,
     ) -> Result<Self> {
         let page_schedulers = column_info
             .page_infos
@@ -2791,6 +2793,7 @@ impl StructuralPrimitiveFieldScheduler {
                     page_index,
                     column_info.index as usize,
                     decompressors,
+                    bits_per_offset,
                 )
             })
             .collect::<Result<Vec<_>>>()?;
@@ -2805,6 +2808,7 @@ impl StructuralPrimitiveFieldScheduler {
         page_index: usize,
         _column_index: usize,
         decompressors: &dyn DecompressorStrategy,
+        bits_per_offset: u8,
     ) -> Result<PageInfoAndScheduler> {
         let scheduler: Box<dyn StructuralPageScheduler> =
             match page_info.encoding.as_structural().layout.as_ref() {
@@ -2824,7 +2828,7 @@ impl StructuralPrimitiveFieldScheduler {
                         page_info.num_rows,
                         full_zip,
                         decompressors,
-                        /*bits_per_offset=*/ 32,
+                        bits_per_offset,
                     )?)
                 }
                 Some(pb::page_layout::Layout::AllNullLayout(all_null)) => {