feat(encoding): encode generic mini-block offsets

rust/lance-encoding/src/compression/block.rs

@@ -65,20 +65,16 @@ impl BlockValueType {
 mod factory;
 pub(crate) mod fixed;
 
-#[cfg(test)]
-pub(crate) use factory::encode_scalar;
-#[cfg(all(test, feature = "bitpacking"))]
-pub(crate) use factory::out_of_line_payload_bytes;
 pub(crate) use factory::{
-    create_block_decompressor, infer_block_value_type, validate_fixed_payload_len,
+    create_block_decompressor, encode_scalar, infer_block_value_type, validate_fixed_payload_len,
 };
 #[cfg(feature = "bitpacking")]
-pub(crate) use factory::{validate_inline_bitpacking_payload, validate_out_of_line_payload};
-#[cfg(test)]
-pub(crate) use fixed::fixed_from_u64_values;
-#[cfg(any(test, feature = "bitpacking"))]
-pub(crate) use fixed::visit_unsigned_values;
-pub(crate) use fixed::{fixed_block, read_unsigned_values};
+pub(crate) use factory::{
+    out_of_line_payload_bytes, validate_inline_bitpacking_payload, validate_out_of_line_payload,
+};
+pub(crate) use fixed::{
+    fixed_block, fixed_from_u64_values, read_unsigned_values, visit_unsigned_values,
+};
 
 #[cfg(test)]
 mod tests;

rust/lance-encoding/src/compression/block/factory.rs

@@ -3,7 +3,6 @@
 
 //! Fallible construction of concrete generic block decompressors.
 
-#[cfg(test)]
 use bytes::Bytes;
 
 use super::*;
@@ -525,7 +524,6 @@ fn decode_scalar(bytes: &[u8], value_type: BlockValueType, label: &str) -> Resul
     })
 }
 
-#[cfg(test)]
 pub fn encode_scalar(value: u64, value_type: BlockValueType) -> Result<Bytes> {
     if value > value_type.max_value() {
         return Err(Error::invalid_input(format!(

rust/lance-encoding/src/compression/block/fixed.rs

@@ -16,7 +16,6 @@ pub fn fixed_block(value_type: BlockValueType, num_values: u64, data: LanceBuffe
     })
 }
 
-#[cfg(test)]
 pub fn fixed_from_u64_values(
     values: &[u64],
     value_type: BlockValueType,
@@ -63,7 +62,6 @@ pub fn fixed_from_u64_values(
     })
 }
 
-#[cfg(any(test, feature = "bitpacking"))]
 pub fn visit_unsigned_values(
     block: &FixedWidthDataBlock,
     value_type: BlockValueType,

rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs

@@ -75,6 +75,19 @@ impl MiniBlockCompressionContext {
             allow_generic_offsets,
         }
     }
+
+    /// Returns the padded per-chunk header bytes for the supplied value buffers.
+    pub(crate) fn chunk_header_bytes(self, value_buffers: u64) -> u64 {
+        let value_buffer_size_bytes = if self.support_large_chunk { 4 } else { 2 };
+        2_u64
+            .saturating_add(self.common_chunk_buffers.saturating_mul(2))
+            .saturating_add(value_buffers.saturating_mul(value_buffer_size_bytes))
+            .next_multiple_of(8)
+    }
+
+    pub(crate) fn allows_generic_offsets(self) -> bool {
+        self.allow_generic_offsets
+    }
 }
 
 /// Describes the size of a mini-block chunk of data
@@ -156,6 +169,21 @@ mod tests {
         assert_eq!(parse_max_miniblock_values(), 4096);
     }
 
+    #[test]
+    fn compression_context_matches_padded_chunk_headers() {
+        let expected = [(0, 8, 16), (1, 8, 16), (2, 16, 16)];
+        for (common_buffers, one_value_buffer, two_value_buffers) in expected {
+            let context = MiniBlockCompressionContext::new(common_buffers, true, true);
+            assert_eq!(context.chunk_header_bytes(1), one_value_buffer);
+            assert_eq!(context.chunk_header_bytes(2), two_value_buffers);
+        }
+
+        let legacy_context = MiniBlockCompressionContext::new(2, false, false);
+        assert_eq!(legacy_context.chunk_header_bytes(1), 8);
+        assert_eq!(legacy_context.chunk_header_bytes(2), 16);
+        assert!(!legacy_context.allows_generic_offsets());
+    }
+
     #[test]
     #[serial]
     fn test_parse_custom_value() {