Add support for repetition index in full-zip encoding

Author: westonpace

rust/lance-encoding/src/buffer.rs

@@ -6,6 +6,7 @@
 use std::{ops::Deref, panic::RefUnwindSafe, ptr::NonNull, sync::Arc};
 
 use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer, ScalarBuffer};
+use itertools::Either;
 use snafu::{location, Location};
 
 use lance_core::{utils::bit::is_pwr_two, Error, Result};
@@ -104,6 +105,18 @@ impl LanceBuffer {
         hex::encode_upper(self)
     }
 
+    /// Combine multiple buffers into a single buffer
+    ///
+    /// This does involve a data copy (and allocation of a new buffer)
+    pub fn concat(buffers: &[Self]) -> Self {
+        let total_len = buffers.iter().map(|b| b.len()).sum();
+        let mut data = Vec::with_capacity(total_len);
+        for buffer in buffers {
+            data.extend_from_slice(buffer.as_ref());
+        }
+        Self::Owned(data)
+    }
+
     /// Converts the buffer into a hex string, inserting a space
     /// between words
     pub fn as_spaced_hex(&self, bytes_per_word: u32) -> String {
@@ -390,6 +403,40 @@ impl From<Buffer> for LanceBuffer {
     }
 }
 
+// An iterator that keeps a clone of a borrowed LanceBuffer so we
+// can have a 'static lifetime
+pub struct BorrowedBufferIter {
+    buffer: arrow_buffer::Buffer,
+    index: usize,
+}
+
+impl Iterator for BorrowedBufferIter {
+    type Item = u8;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index >= self.buffer.len() {
+            None
+        } else {
+            // SAFETY: we just checked that index is in bounds
+            let byte = unsafe { self.buffer.get_unchecked(self.index) };
+            self.index += 1;
+            Some(*byte)
+        }
+    }
+}
+
+impl IntoIterator for LanceBuffer {
+    type Item = u8;
+    type IntoIter = Either<std::vec::IntoIter<u8>, BorrowedBufferIter>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        match self {
+            Self::Borrowed(buffer) => Either::Right(BorrowedBufferIter { buffer, index: 0 }),
+            Self::Owned(buffer) => Either::Left(buffer.into_iter()),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use arrow_buffer::Buffer;

rust/lance-encoding/src/encoder.rs

@@ -1221,15 +1221,14 @@ impl StructuralEncodingStrategy {
                 | DataType::LargeUtf8,
         )
     }
-}
 
-impl FieldEncodingStrategy for StructuralEncodingStrategy {
-    fn create_field_encoder(
+    fn do_create_field_encoder(
         &self,
         _encoding_strategy_root: &dyn FieldEncodingStrategy,
         field: &Field,
         column_index: &mut ColumnIndexSequence,
         options: &EncodingOptions,
+        root_field_metadata: &HashMap<String, String>,
     ) -> Result<Box<dyn FieldEncoder>> {
         let data_type = field.data_type();
         if Self::is_primitive_type(&data_type) {
@@ -1238,16 +1237,18 @@ impl FieldEncodingStrategy for StructuralEncodingStrategy {
                 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("List should have a child");
-                    let child_encoder = self.create_field_encoder(
+                    let child_encoder = self.do_create_field_encoder(
                         _encoding_strategy_root,
                         child,
                         column_index,
                         options,
+                        root_field_metadata,
                     )?;
                     Ok(Box::new(ListStructuralEncoder::new(child_encoder)))
                 }
@@ -1258,17 +1259,19 @@ impl FieldEncodingStrategy for StructuralEncodingStrategy {
                             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.create_field_encoder(
+                                self.do_create_field_encoder(
                                     _encoding_strategy_root,
                                     field,
                                     column_index,
                                     options,
+                                    root_field_metadata,
                                 )
                             })
                             .collect::<Result<Vec<_>>>()?;
@@ -1283,6 +1286,7 @@ impl FieldEncodingStrategy for StructuralEncodingStrategy {
                             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
@@ -1299,6 +1303,24 @@ impl FieldEncodingStrategy for StructuralEncodingStrategy {
     }
 }
 
+impl FieldEncodingStrategy for StructuralEncodingStrategy {
+    fn create_field_encoder(
+        &self,
+        encoding_strategy_root: &dyn FieldEncodingStrategy,
+        field: &Field,
+        column_index: &mut ColumnIndexSequence,
+        options: &EncodingOptions,
+    ) -> Result<Box<dyn FieldEncoder>> {
+        self.do_create_field_encoder(
+            encoding_strategy_root,
+            field,
+            column_index,
+            options,
+            &field.metadata,
+        )
+    }
+}
+
 /// A batch encoder that encodes RecordBatch objects by delegating
 /// to field encoders for each top-level field in the batch.
 pub struct BatchEncoder {

rust/lance-encoding/src/encodings/logical/list.rs

@@ -1465,6 +1465,9 @@ mod tests {
     };
     use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer};
     use arrow_schema::{DataType, Field, Fields};
+    use lance_core::datatypes::{
+        STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK,
+    };
     use rstest::rstest;
 
     use crate::{
@@ -1484,8 +1487,16 @@ mod tests {
     #[test_log::test(tokio::test)]
     async fn test_list(
         #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion,
+        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
+        structural_encoding: &str,
     ) {
-        let field = Field::new("", make_list_type(DataType::Int32), true);
+        let mut field_metadata = HashMap::new();
+        field_metadata.insert(
+            STRUCTURAL_ENCODING_META_KEY.to_string(),
+            structural_encoding.into(),
+        );
+        let field =
+            Field::new("", make_list_type(DataType::Int32), true).with_metadata(field_metadata);
         check_round_trip_encoding_random(field, version).await;
     }
 
@@ -1544,6 +1555,8 @@ mod tests {
     #[test_log::test(tokio::test)]
     async fn test_simple_list(
         #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion,
+        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
+        structural_encoding: &str,
     ) {
         let items_builder = Int32Builder::new();
         let mut list_builder = ListBuilder::new(items_builder);
@@ -1553,18 +1566,23 @@ mod tests {
         list_builder.append_value([Some(6), Some(7), Some(8)]);
         let list_array = list_builder.finish();
 
+        let mut field_metadata = HashMap::new();
+        field_metadata.insert(
+            STRUCTURAL_ENCODING_META_KEY.to_string(),
+            structural_encoding.into(),
+        );
+
         let test_cases = TestCases::default()
             .with_range(0..2)
             .with_range(0..3)
             .with_range(1..3)
             .with_indices(vec![1, 3])
             .with_indices(vec![2])
             .with_file_version(version);
-        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
+        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, field_metadata)
             .await;
     }
 
-    #[rstest]
     #[test_log::test(tokio::test)]
     async fn test_simple_sliced_list() {
         let items_builder = Int32Builder::new();
@@ -1587,7 +1605,6 @@ mod tests {
             .await;
     }
 
-    #[rstest]
     #[test_log::test(tokio::test)]
     async fn test_list_with_garbage_nulls() {
         // In Arrow, list nulls are allowed to be non-empty, with masked garbage values
@@ -1613,8 +1630,12 @@ mod tests {
             .await;
     }
 
+    #[rstest]
     #[test_log::test(tokio::test)]
-    async fn test_simple_two_page_list() {
+    async fn test_simple_two_page_list(
+        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
+        structural_encoding: &str,
+    ) {
         // This is a simple pre-defined list that spans two pages.  This test is useful for
         // debugging the repetition index
         let items_builder = Int64Builder::new();
@@ -1632,14 +1653,20 @@ mod tests {
         }
         let list_array_2 = list_builder.finish();
 
+        let mut metadata = HashMap::new();
+        metadata.insert(
+            STRUCTURAL_ENCODING_META_KEY.to_string(),
+            structural_encoding.into(),
+        );
+
         let test_cases = TestCases::default()
             .with_file_version(LanceFileVersion::V2_1)
             .with_page_sizes(vec![100])
             .with_range(800..900);
         check_round_trip_encoding_of_data(
             vec![Arc::new(list_array_1), Arc::new(list_array_2)],
             &test_cases,
-            HashMap::new(),
+            metadata,
         )
         .await;
     }