1use arrow_buffer::{BooleanBufferBuilder, bit_util};
5
6use crate::buffer::LanceBuffer;
7use crate::compression::{
8 BlockCompressor, BlockDecompressor, FixedPerValueDecompressor, MiniBlockDecompressor,
9 require_block_payload,
10};
11use crate::data::{
12 AllNullDataBlock, BlockInfo, DataBlock, FixedSizeListBlock, FixedWidthDataBlock,
13 NullableDataBlock,
14};
15use crate::encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock};
16use crate::encodings::logical::primitive::miniblock::{
17 MAX_MINIBLOCK_BYTES, MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed,
18 MiniBlockCompressionContext, MiniBlockCompressor,
19};
20use crate::format::EncodingUtils;
21use crate::format::metadata::compressive_encoding::Compression;
22use crate::format::metadata::{self, CompressiveEncoding};
23
24use lance_core::{Error, Result};
25
26#[derive(Debug, Default)]
28pub struct ValueEncoder {}
29
30impl ValueEncoder {
31 fn find_log_vals_per_chunk(bytes_per_word: u64, values_per_word: u64) -> Result<(u64, u64)> {
33 let mut size_bytes = 2 * bytes_per_word;
34 let (mut log_num_vals, mut num_vals) = match values_per_word {
35 1 => (1, 2),
36 8 => (3, 8),
37 _ => unreachable!(),
38 };
39
40 if size_bytes >= MAX_MINIBLOCK_BYTES {
41 let num_values = 2 * values_per_word;
42 return Err(Error::invalid_input(format!(
43 "Value is too wide for miniblock encoding: {} values require {} bytes but a \
44 miniblock chunk is limited to {} bytes.",
45 num_values, size_bytes, MAX_MINIBLOCK_BYTES
46 )));
47 }
48
49 while 2 * size_bytes < MAX_MINIBLOCK_BYTES && 2 * num_vals <= *MAX_MINIBLOCK_VALUES {
50 log_num_vals += 1;
51 size_bytes *= 2;
52 num_vals *= 2;
53 }
54
55 Ok((log_num_vals, num_vals))
56 }
57
58 fn chunk_data(data: FixedWidthDataBlock) -> Result<MiniBlockCompressed> {
59 let (bytes_per_word, values_per_word) = if data.bits_per_value.is_multiple_of(8) {
64 (data.bits_per_value / 8, 1)
65 } else {
66 (data.bits_per_value, 8)
67 };
68
69 let (log_vals_per_chunk, vals_per_chunk) =
71 Self::find_log_vals_per_chunk(bytes_per_word, values_per_word)?;
72 let num_chunks = bit_util::ceil(data.num_values as usize, vals_per_chunk as usize);
73 debug_assert_eq!(vals_per_chunk % values_per_word, 0);
74 let bytes_per_chunk = bytes_per_word * (vals_per_chunk / values_per_word);
75 let bytes_per_chunk = u32::try_from(bytes_per_chunk).unwrap();
76 debug_assert!(bytes_per_chunk > 0);
77
78 let data_buffer = data.data;
79
80 let mut row_offset = 0;
81 let mut chunks = Vec::with_capacity(num_chunks);
82
83 let mut bytes_counter = 0;
84 loop {
85 if row_offset + vals_per_chunk <= data.num_values {
86 chunks.push(MiniBlockChunk {
88 log_num_values: log_vals_per_chunk as u8,
89 buffer_sizes: vec![bytes_per_chunk],
90 });
91 row_offset += vals_per_chunk;
92 bytes_counter += bytes_per_chunk as u64;
93 } else if row_offset < data.num_values {
94 let num_bytes = data_buffer.len() as u64 - bytes_counter;
96 let num_bytes = u32::try_from(num_bytes).unwrap();
97 chunks.push(MiniBlockChunk {
98 log_num_values: 0,
99 buffer_sizes: vec![num_bytes],
100 });
101 break;
102 } else {
103 break;
105 }
106 }
107
108 debug_assert_eq!(chunks.len(), num_chunks);
109
110 Ok(MiniBlockCompressed {
111 chunks,
112 data: vec![data_buffer],
113 num_values: data.num_values,
114 })
115 }
116}
117
118#[derive(Debug)]
119struct MiniblockFslLayer {
120 validity: Option<LanceBuffer>,
121 dimension: u64,
122}
123
124impl ValueEncoder {
141 fn make_fsl_encoding(layers: &[MiniblockFslLayer], bits_per_value: u64) -> CompressiveEncoding {
142 let mut encoding = EncodingUtils::flat(bits_per_value, None);
143 for layer in layers.iter().rev() {
144 let has_validity = layer.validity.is_some();
145 let dimension = layer.dimension;
146 encoding = EncodingUtils::fsl(dimension, has_validity, encoding);
147 }
148 encoding
149 }
150
151 fn extract_fsl_chunk(
152 data: &FixedWidthDataBlock,
153 layers: &[MiniblockFslLayer],
154 row_offset: usize,
155 num_rows: usize,
156 validity_buffers: &mut [Vec<u8>],
157 ) -> Vec<u32> {
158 let mut row_offset = row_offset;
159 let mut num_values = num_rows;
160 let mut buffer_counter = 0;
161 let mut buffer_sizes = Vec::with_capacity(validity_buffers.len() + 1);
162 for layer in layers {
163 row_offset *= layer.dimension as usize;
164 num_values *= layer.dimension as usize;
165 if let Some(validity) = &layer.validity {
166 let validity_slice = validity
167 .clone()
168 .bit_slice_le_with_length(row_offset, num_values);
169 validity_buffers[buffer_counter].extend_from_slice(&validity_slice);
170 buffer_sizes.push(validity_slice.len() as u32);
171 buffer_counter += 1;
172 }
173 }
174
175 let bits_in_chunk = data.bits_per_value * num_values as u64;
176 let bytes_in_chunk = bits_in_chunk.div_ceil(8);
177 let bytes_in_chunk = u32::try_from(bytes_in_chunk).unwrap();
178 debug_assert!(bytes_in_chunk > 0);
179 buffer_sizes.push(bytes_in_chunk);
180
181 buffer_sizes
182 }
183
184 fn chunk_fsl(
185 data: FixedWidthDataBlock,
186 layers: Vec<MiniblockFslLayer>,
187 num_rows: u64,
188 ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
189 let mut ceil_bytes_validity = 0;
191 let mut cum_dim = 1;
192 let mut num_validity_buffers = 0;
193 for layer in &layers {
194 cum_dim *= layer.dimension;
195 if layer.validity.is_some() {
196 ceil_bytes_validity += cum_dim.div_ceil(8);
197 num_validity_buffers += 1;
198 }
199 }
200 let cum_bits_per_value = data.bits_per_value * cum_dim;
202 let (cum_bytes_per_word, vals_per_word) = if cum_bits_per_value.is_multiple_of(8) {
203 (cum_bits_per_value / 8, 1)
204 } else {
205 (cum_bits_per_value, 8)
206 };
207 let est_bytes_per_word = (ceil_bytes_validity * vals_per_word) + cum_bytes_per_word;
208 let (log_rows_per_chunk, rows_per_chunk) =
209 Self::find_log_vals_per_chunk(est_bytes_per_word, vals_per_word)?;
210
211 let num_chunks = num_rows.div_ceil(rows_per_chunk) as usize;
212
213 let mut chunks = Vec::with_capacity(num_chunks);
215 let mut validity_buffers: Vec<Vec<u8>> = Vec::with_capacity(num_validity_buffers);
216 cum_dim = 1;
217 for layer in &layers {
218 cum_dim *= layer.dimension;
219 if let Some(validity) = &layer.validity {
220 let layer_bytes_validity = cum_dim.div_ceil(8);
221 let validity_with_padding =
222 layer_bytes_validity as usize * num_chunks * rows_per_chunk as usize;
223 debug_assert!(validity_with_padding >= validity.len());
224 validity_buffers.push(Vec::with_capacity(
225 layer_bytes_validity as usize * num_chunks,
226 ));
227 }
228 }
229
230 let mut row_offset = 0;
232 while row_offset + rows_per_chunk <= num_rows {
233 let buffer_sizes = Self::extract_fsl_chunk(
234 &data,
235 &layers,
236 row_offset as usize,
237 rows_per_chunk as usize,
238 &mut validity_buffers,
239 );
240 row_offset += rows_per_chunk;
241 chunks.push(MiniBlockChunk {
242 log_num_values: log_rows_per_chunk as u8,
243 buffer_sizes,
244 })
245 }
246 let rows_in_chunk = num_rows - row_offset;
247 if rows_in_chunk > 0 {
248 let buffer_sizes = Self::extract_fsl_chunk(
249 &data,
250 &layers,
251 row_offset as usize,
252 rows_in_chunk as usize,
253 &mut validity_buffers,
254 );
255 chunks.push(MiniBlockChunk {
256 log_num_values: 0,
257 buffer_sizes,
258 });
259 }
260
261 let encoding = Self::make_fsl_encoding(&layers, data.bits_per_value);
262 let buffers = validity_buffers
264 .into_iter()
265 .map(LanceBuffer::from)
266 .chain(std::iter::once(data.data))
267 .collect::<Vec<_>>();
268
269 Ok((
270 MiniBlockCompressed {
271 chunks,
272 data: buffers,
273 num_values: num_rows,
274 },
275 encoding,
276 ))
277 }
278
279 fn miniblock_fsl(data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
280 let num_rows = data.num_values();
281 let fsl = data.as_fixed_size_list().unwrap();
282 let mut layers = Vec::new();
283 let mut child = *fsl.child;
284 let mut cur_layer = MiniblockFslLayer {
285 validity: None,
286 dimension: fsl.dimension,
287 };
288 loop {
289 if let DataBlock::Nullable(nullable) = child {
290 cur_layer.validity = Some(nullable.nulls);
291 child = *nullable.data;
292 }
293 match child {
294 DataBlock::FixedSizeList(inner) => {
295 layers.push(cur_layer);
296 cur_layer = MiniblockFslLayer {
297 validity: None,
298 dimension: inner.dimension,
299 };
300 child = *inner.child;
301 }
302 DataBlock::FixedWidth(inner) => {
303 layers.push(cur_layer);
304 return Self::chunk_fsl(inner, layers, num_rows);
305 }
306 _ => unreachable!("Unexpected data block type in value encoder's miniblock_fsl"),
307 }
308 }
309 }
310}
311
312struct PerValueFslValidityIter {
313 buffer: LanceBuffer,
314 bits_per_row: usize,
315 offset: usize,
316}
317
318impl ValueEncoder {
323 fn fsl_to_encoding(fsl: &FixedSizeListBlock) -> CompressiveEncoding {
324 let mut inner = fsl.child.as_ref();
325 let mut has_validity = false;
326 if let DataBlock::Nullable(nullable) = inner {
327 has_validity = true;
328 inner = nullable.data.as_ref();
329 }
330 let inner_encoding = match inner {
331 DataBlock::AllNull(_) => {
335 has_validity = true;
336 EncodingUtils::constant(None)
337 }
338 DataBlock::FixedWidth(fixed_width) => {
339 EncodingUtils::flat(fixed_width.bits_per_value, None)
340 }
341 DataBlock::FixedSizeList(inner) => Self::fsl_to_encoding(inner),
342 _ => unreachable!(
343 "Unexpected data block type in value encoder's fsl_to_encoding: {}",
344 inner.name()
345 ),
346 };
347 EncodingUtils::fsl(fsl.dimension, has_validity, inner_encoding)
348 }
349
350 fn simple_per_value_fsl(fsl: FixedSizeListBlock) -> (PerValueDataBlock, CompressiveEncoding) {
351 let encoding = Self::fsl_to_encoding(&fsl);
353 let Some(flat) = fsl.try_into_flat() else {
354 unreachable!("per_value_fsl only sends FSL blocks without nullable children here")
355 };
356 let data = FixedWidthDataBlock {
357 block_info: BlockInfo::new(),
358 ..flat
359 };
360 (PerValueDataBlock::Fixed(data), encoding)
361 }
362
363 fn nullable_per_value_fsl(fsl: FixedSizeListBlock) -> (PerValueDataBlock, CompressiveEncoding) {
364 let encoding = Self::fsl_to_encoding(&fsl);
366 let num_values = fsl.num_values();
367 let mut bytes_per_row = 0;
368 let mut cum_dim = 1;
369 let mut current = fsl;
370 let mut validity_iters: Vec<PerValueFslValidityIter> = Vec::new();
371 let data_bytes_per_row: usize;
372 let data_buffer: LanceBuffer;
373 loop {
374 cum_dim *= current.dimension;
375 let mut child = *current.child;
376 if let DataBlock::Nullable(nullable) = child {
377 bytes_per_row += cum_dim.div_ceil(8) as usize;
379 validity_iters.push(PerValueFslValidityIter {
380 buffer: nullable.nulls,
381 bits_per_row: cum_dim as usize,
382 offset: 0,
383 });
384 child = *nullable.data;
385 };
386 match child {
387 DataBlock::FixedSizeList(inner) => {
388 current = inner;
389 }
390 DataBlock::FixedWidth(fixed_width) => {
391 data_bytes_per_row =
392 (fixed_width.bits_per_value.div_ceil(8) * cum_dim) as usize;
393 bytes_per_row += data_bytes_per_row;
394 data_buffer = fixed_width.data;
395 break;
396 }
397 DataBlock::AllNull(_) => {
398 bytes_per_row += cum_dim.div_ceil(8) as usize;
402 validity_iters.push(PerValueFslValidityIter {
403 buffer: LanceBuffer::from(vec![
404 0u8;
405 cum_dim.div_ceil(8) as usize
406 * num_values as usize
407 ]),
408 bits_per_row: cum_dim as usize,
409 offset: 0,
410 });
411 data_bytes_per_row = 0;
412 data_buffer = LanceBuffer::empty();
413 break;
414 }
415 _ => unreachable!(
416 "Unexpected data block type in value encoder's nullable_per_value_fsl: {:?}",
417 child
418 ),
419 }
420 }
421
422 let bytes_needed = bytes_per_row * num_values as usize;
423 let mut zipped = Vec::with_capacity(bytes_needed);
424 let data_slice = &data_buffer;
425 for i in 0..num_values as usize {
427 for validity in validity_iters.iter_mut() {
428 let validity_slice = validity
429 .buffer
430 .bit_slice_le_with_length(validity.offset, validity.bits_per_row);
431 zipped.extend_from_slice(&validity_slice);
432 validity.offset += validity.bits_per_row;
433 }
434 let start = i * data_bytes_per_row;
435 let end = start + data_bytes_per_row;
436 zipped.extend_from_slice(&data_slice[start..end]);
437 }
438
439 let zipped = LanceBuffer::from(zipped);
440 let data = PerValueDataBlock::Fixed(FixedWidthDataBlock {
441 bits_per_value: bytes_per_row as u64 * 8,
442 num_values,
443 data: zipped,
444 block_info: BlockInfo::new(),
445 });
446 (data, encoding)
447 }
448
449 fn per_value_fsl(fsl: FixedSizeListBlock) -> (PerValueDataBlock, CompressiveEncoding) {
450 if !fsl.child.is_nullable() {
451 Self::simple_per_value_fsl(fsl)
452 } else {
453 Self::nullable_per_value_fsl(fsl)
454 }
455 }
456}
457
458impl BlockCompressor for ValueEncoder {
459 fn compress(&self, data: DataBlock) -> Result<(Option<LanceBuffer>, CompressiveEncoding)> {
460 if matches!(data, DataBlock::FixedSizeList(_)) {
461 let (PerValueDataBlock::Fixed(values), encoding) =
462 PerValueCompressor::compress(self, data)?
463 else {
464 return Err(Error::internal(
465 "Fixed-size-list block has variable row width",
466 ));
467 };
468 return Ok((Some(values.data), encoding));
469 }
470 let DataBlock::FixedWidth(fixed_width) = data else {
471 return Err(Error::invalid_input(format!(
472 "ValueEncoder cannot compress a {} block",
473 data.name()
474 )));
475 };
476 Ok((
477 Some(fixed_width.data),
478 EncodingUtils::flat(fixed_width.bits_per_value, None),
479 ))
480 }
481}
482
483impl MiniBlockCompressor for ValueEncoder {
484 fn compress(
485 &self,
486 _context: MiniBlockCompressionContext,
487 chunk: DataBlock,
488 ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
489 match chunk {
490 DataBlock::FixedWidth(fixed_width) => {
491 let encoding = EncodingUtils::flat(fixed_width.bits_per_value, None);
492 Ok((Self::chunk_data(fixed_width)?, encoding))
493 }
494 DataBlock::FixedSizeList(_) => Self::miniblock_fsl(chunk),
495 _ => Err(Error::invalid_input_source(
496 format!(
497 "Cannot compress a data block of type {} with ValueEncoder",
498 chunk.name()
499 )
500 .into(),
501 )),
502 }
503 }
504}
505
506#[derive(Debug)]
507struct ValueFslDesc {
508 dimension: u64,
509 has_validity: bool,
510}
511
512#[derive(Debug)]
515pub struct ValueDecompressor {
516 bits_per_item: u64,
518 bits_per_value: u64,
523 layers: Vec<ValueFslDesc>,
524}
525
526impl ValueDecompressor {
527 pub fn from_flat(description: &metadata::Flat) -> Self {
528 Self {
529 bits_per_item: description.bits_per_value,
530 bits_per_value: description.bits_per_value,
531 layers: Vec::default(),
532 }
533 }
534
535 pub fn from_fsl(mut description: &metadata::FixedSizeList) -> Result<Self> {
536 let mut layers = Vec::new();
537 let mut cum_dim = 1_u64;
538 let mut bytes_per_value = 0_u64;
539 loop {
540 if description.items_per_value == 0 {
541 return Err(Error::invalid_input("FSL dimension must be positive"));
542 }
543 layers.push(ValueFslDesc {
544 has_validity: description.has_validity,
545 dimension: description.items_per_value,
546 });
547 cum_dim = cum_dim
548 .checked_mul(description.items_per_value)
549 .ok_or_else(|| Error::invalid_input("FSL dimension product overflows"))?;
550 if description.has_validity {
551 bytes_per_value = bytes_per_value
552 .checked_add(cum_dim.div_ceil(8))
553 .ok_or_else(|| Error::invalid_input("FSL validity size overflows"))?;
554 }
555 let encoding = description
556 .values
557 .as_ref()
558 .ok_or_else(|| Error::invalid_input("FSL encoding missing inner values field"))?
559 .compression
560 .as_ref()
561 .ok_or_else(|| {
562 Error::invalid_input("FSL encoding missing inner compression field")
563 })?;
564 match encoding {
565 Compression::FixedSizeList(inner) => {
566 description = inner;
567 }
568 Compression::Flat(flat) => {
569 if flat.bits_per_value == 0 || flat.data.is_some() {
570 return Err(Error::invalid_input(
571 "FSL requires positive plain Flat leaf width",
572 ));
573 }
574 let bits_per_value = flat
575 .bits_per_value
576 .checked_mul(cum_dim)
577 .and_then(|bits| bytes_per_value.checked_mul(8)?.checked_add(bits))
578 .ok_or_else(|| Error::invalid_input("FSL value bit width overflows"))?;
579 return Ok(Self {
580 bits_per_item: flat.bits_per_value,
581 bits_per_value,
582 layers,
583 });
584 }
585 Compression::Constant(_) => {
588 return Ok(Self {
589 bits_per_item: 0,
590 bits_per_value: bytes_per_value.checked_mul(8).ok_or_else(|| {
591 Error::invalid_input("FSL validity bit width overflows")
592 })?,
593 layers,
594 });
595 }
596 _ => {
597 return Err(Error::invalid_input(format!(
598 "Unexpected inner encoding type in FSL descriptor: {:?}",
599 encoding
600 )));
601 }
602 }
603 }
604 }
605
606 fn buffer_to_block(&self, data: LanceBuffer, num_values: u64) -> DataBlock {
607 DataBlock::FixedWidth(FixedWidthDataBlock {
608 bits_per_value: self.bits_per_item,
609 num_values,
610 data,
611 block_info: BlockInfo::new(),
612 })
613 }
614}
615
616impl BlockDecompressor for ValueDecompressor {
617 fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
618 let data = require_block_payload(data, "Flat block")?;
619 if !self.layers.is_empty() {
620 let expected = num_values
621 .checked_mul(self.bits_per_value)
622 .and_then(|bits| usize::try_from(bits.div_ceil(8)).ok())
623 .ok_or_else(|| Error::invalid_input("Fixed-size-list block length overflows"))?;
624 if data.len() != expected {
625 return Err(Error::invalid_input(format!(
626 "Fixed-size-list block has {} bytes, expected {expected}",
627 data.len()
628 )));
629 }
630 return FixedPerValueDecompressor::decompress(
631 self,
632 FixedWidthDataBlock {
633 data,
634 bits_per_value: self.bits_per_value,
635 num_values,
636 block_info: BlockInfo::new(),
637 },
638 num_values,
639 );
640 }
641 let block = self.buffer_to_block(data, num_values);
642 assert_eq!(block.num_values(), num_values);
643 Ok(block)
644 }
645}
646
647impl MiniBlockDecompressor for ValueDecompressor {
648 fn num_buffers(&self) -> usize {
649 1 + self
650 .layers
651 .iter()
652 .filter(|layer| layer.has_validity)
653 .count()
654 }
655
656 fn decompress(&self, data: Vec<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
657 if data.len() != self.num_buffers() {
658 return Err(Error::invalid_input(format!(
659 "Flat/FSL mini-block has {} buffers, expected {}",
660 data.len(),
661 self.num_buffers()
662 )));
663 }
664 let mut child_values = num_values;
665 let mut buffers = data.iter();
666 for layer in &self.layers {
667 child_values = child_values
668 .checked_mul(layer.dimension)
669 .ok_or_else(|| Error::invalid_input("FSL child value count overflows"))?;
670 if layer.has_validity {
671 validate_fixed_buffer(
672 buffers
673 .next()
674 .ok_or_else(|| Error::invalid_input("FSL validity buffer is missing"))?,
675 child_values,
676 1,
677 )?;
678 }
679 }
680 validate_fixed_buffer(
681 buffers
682 .next()
683 .ok_or_else(|| Error::invalid_input("Flat value buffer is missing"))?,
684 child_values,
685 self.bits_per_item,
686 )?;
687 let num_items = child_values;
688 let mut buffer_iter = data.into_iter().rev();
689
690 let data_buf = buffer_iter.next().unwrap();
692 let items = self.buffer_to_block(data_buf, num_items);
693 let mut lists = items;
694
695 for layer in self.layers.iter().rev() {
696 if layer.has_validity {
697 let validity_buf = buffer_iter.next().unwrap();
698 lists = DataBlock::Nullable(NullableDataBlock {
699 data: Box::new(lists),
700 nulls: validity_buf,
701 block_info: BlockInfo::default(),
702 });
703 }
704 lists = DataBlock::FixedSizeList(FixedSizeListBlock {
705 child: Box::new(lists),
706 dimension: layer.dimension,
707 })
708 }
709
710 assert_eq!(lists.num_values(), num_values);
711 Ok(lists)
712 }
713
714 fn decoded_size_bytes(&self, num_values: u64) -> Option<u64> {
715 if self.has_validity() {
716 return None;
717 }
718 num_values
719 .checked_mul(self.bits_per_value)
720 .map(|bits| bits.div_ceil(8))
721 }
722}
723
724fn validate_fixed_buffer(data: &LanceBuffer, num_values: u64, bits: u64) -> Result<()> {
725 let expected = num_values
726 .checked_mul(bits)
727 .map(|bits| bits.div_ceil(8))
728 .ok_or_else(|| Error::invalid_input("Flat decoded bit length overflows"))?;
729 if data.len() as u64 != expected {
730 return Err(Error::invalid_input(format!(
731 "Flat buffer has {} bytes, expected {expected}",
732 data.len()
733 )));
734 }
735 Ok(())
736}
737
738struct FslDecompressorValidityBuilder {
739 buffer: BooleanBufferBuilder,
740 bits_per_row: usize,
741 bytes_per_row: usize,
742}
743
744impl ValueDecompressor {
746 fn has_validity(&self) -> bool {
747 self.layers.iter().any(|layer| layer.has_validity)
748 }
749
750 fn simple_decompress(&self, data: FixedWidthDataBlock, num_rows: u64) -> DataBlock {
752 let mut cum_dim = 1;
753 for layer in &self.layers {
754 cum_dim *= layer.dimension;
755 }
756 debug_assert_eq!(self.bits_per_item, data.bits_per_value / cum_dim);
757 let mut block = DataBlock::FixedWidth(FixedWidthDataBlock {
758 bits_per_value: self.bits_per_item,
759 num_values: num_rows * cum_dim,
760 data: data.data,
761 block_info: BlockInfo::new(),
762 });
763 for layer in self.layers.iter().rev() {
764 block = DataBlock::FixedSizeList(FixedSizeListBlock {
765 child: Box::new(block),
766 dimension: layer.dimension,
767 });
768 }
769 debug_assert_eq!(num_rows, block.num_values());
770 block
771 }
772
773 fn unzip_decompress(&self, data: FixedWidthDataBlock, num_rows: usize) -> DataBlock {
775 assert_eq!(self.bits_per_item % 8, 0);
777 let bytes_per_item = self.bits_per_item / 8;
778 let mut buffer_builders = Vec::with_capacity(self.layers.len());
779 let mut cum_dim = 1;
780 let mut total_size_bytes = 0;
781 for layer in &self.layers {
783 cum_dim *= layer.dimension as usize;
784 if layer.has_validity {
785 let validity_size_bits = cum_dim;
786 let validity_size_bytes = validity_size_bits.div_ceil(8);
787 total_size_bytes += num_rows * validity_size_bytes;
788 buffer_builders.push(FslDecompressorValidityBuilder {
789 buffer: BooleanBufferBuilder::new(validity_size_bits * num_rows),
790 bits_per_row: cum_dim,
791 bytes_per_row: validity_size_bytes,
792 })
793 }
794 }
795 let num_items = num_rows * cum_dim;
796 let data_size = num_items * bytes_per_item as usize;
797 total_size_bytes += data_size;
798 let mut data_buffer = Vec::with_capacity(data_size);
799
800 assert_eq!(data.data.len(), total_size_bytes);
801
802 let bytes_per_value = bytes_per_item as usize;
803 let data_bytes_per_row = bytes_per_value * cum_dim;
804
805 let mut data_offset = 0;
807 while data_offset < total_size_bytes {
808 for builder in buffer_builders.iter_mut() {
809 let start = data_offset * 8;
810 let end = start + builder.bits_per_row;
811 builder.buffer.append_packed_range(start..end, &data.data);
812 data_offset += builder.bytes_per_row;
813 }
814 let end = data_offset + data_bytes_per_row;
815 data_buffer.extend_from_slice(&data.data[data_offset..end]);
816 data_offset += data_bytes_per_row;
817 }
818
819 let mut block = if self.bits_per_item == 0 {
825 debug_assert!(data_buffer.is_empty());
826 DataBlock::AllNull(AllNullDataBlock {
827 num_values: num_items as u64,
828 })
829 } else {
830 DataBlock::FixedWidth(FixedWidthDataBlock {
831 bits_per_value: self.bits_per_item,
832 num_values: num_items as u64,
833 data: LanceBuffer::from(data_buffer),
834 block_info: BlockInfo::new(),
835 })
836 };
837
838 let mut validity_bufs = buffer_builders
839 .into_iter()
840 .rev()
841 .map(|mut b| LanceBuffer::from(b.buffer.finish().into_inner()));
842 for layer in self.layers.iter().rev() {
843 if layer.has_validity {
844 let nullable = NullableDataBlock {
845 data: Box::new(block),
846 nulls: validity_bufs.next().unwrap(),
847 block_info: BlockInfo::new(),
848 };
849 block = DataBlock::Nullable(nullable);
850 }
851 block = DataBlock::FixedSizeList(FixedSizeListBlock {
852 child: Box::new(block),
853 dimension: layer.dimension,
854 });
855 }
856
857 assert_eq!(num_rows, block.num_values() as usize);
858
859 block
860 }
861}
862
863impl FixedPerValueDecompressor for ValueDecompressor {
864 fn decompress(&self, data: FixedWidthDataBlock, num_rows: u64) -> Result<DataBlock> {
865 if self.has_validity() {
866 Ok(self.unzip_decompress(data, num_rows as usize))
867 } else {
868 Ok(self.simple_decompress(data, num_rows))
869 }
870 }
871
872 fn bits_per_value(&self) -> u64 {
873 self.bits_per_value
874 }
875
876 fn decoded_size_bytes(&self, num_values: u64) -> Option<u64> {
877 if self.has_validity() {
878 return None;
879 }
880 num_values
881 .checked_mul(self.bits_per_value)
882 .map(|bits| bits.div_ceil(8))
883 }
884}
885
886impl PerValueCompressor for ValueEncoder {
887 fn compress(&self, data: DataBlock) -> Result<(PerValueDataBlock, CompressiveEncoding)> {
888 let (data, encoding) = match data {
889 DataBlock::FixedWidth(fixed_width) => {
890 let encoding = EncodingUtils::flat(fixed_width.bits_per_value, None);
891 (PerValueDataBlock::Fixed(fixed_width), encoding)
892 }
893 DataBlock::FixedSizeList(fixed_size_list) => Self::per_value_fsl(fixed_size_list),
894 _ => unimplemented!(
895 "Cannot compress block of type {} with ValueEncoder",
896 data.name()
897 ),
898 };
899 Ok((data, encoding))
900 }
901}
902
903#[cfg(test)]
905mod tests {
906 use std::{collections::HashMap, sync::Arc};
907
908 use arrow_array::{
909 Array, ArrayRef, Decimal128Array, FixedSizeListArray, Int32Array, ListArray, NullArray,
910 UInt8Array, make_array, new_null_array, types::UInt32Type,
911 };
912 use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer};
913 use arrow_schema::{DataType, Field, TimeUnit};
914 use lance_datagen::{ArrayGeneratorExt, Dimension, RowCount, Seed, array, gen_batch};
915
916 use crate::{
917 compression::{
918 DecompressionStrategy, DefaultDecompressionStrategy, FixedPerValueDecompressor,
919 MiniBlockDecompressor, try_fixed_size_list_block,
920 },
921 compression_config::CompressionFieldParams,
922 data::DataBlock,
923 encodings::{
924 logical::primitive::{
925 fullzip::{PerValueCompressor, PerValueDataBlock},
926 miniblock::{MiniBlockCompressionContext, MiniBlockCompressor},
927 },
928 physical::value::ValueDecompressor,
929 },
930 format::metadata::compressive_encoding::Compression,
931 testing::{
932 TestCases, TestEncoding, check_basic_random_case, check_round_trip_encoding_of_data,
933 },
934 };
935
936 use super::ValueEncoder;
937
938 #[rstest::rstest]
939 #[case::raw("none", false)]
940 #[case::lz4("lz4", false)]
941 #[case::zstd("zstd", false)]
942 #[case::all_null_children("lz4", true)]
943 #[test]
944 fn test_vector_dictionary_block_lengths(#[case] compression: &str, #[case] all_null: bool) {
945 let array = FixedSizeListArray::new(
946 Arc::new(Field::new("item", DataType::Int32, true)),
947 3,
948 Arc::new(if all_null {
949 Int32Array::new_null(6)
950 } else {
951 Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), None])
952 }),
953 None,
954 );
955 let data_type = array.data_type().clone();
956 let data = DataBlock::from_array(array.clone());
957 let params = CompressionFieldParams {
958 compression: Some(compression.to_string()),
959 ..Default::default()
960 };
961 let encoder = try_fixed_size_list_block(&data, ¶ms).unwrap().unwrap();
962 let (payload, encoding) = encoder.compress(data).unwrap();
963 let decoder = DefaultDecompressionStrategy::default()
964 .create_block_decompressor(&encoding)
965 .unwrap();
966 let decoded = decoder.decompress(payload.clone(), 2).unwrap();
967 assert_eq!(
968 decoded.into_arrow(data_type, true).unwrap(),
969 array.to_data()
970 );
971 assert!(decoder.decompress(payload.clone(), 1).is_err());
972 assert!(decoder.decompress(payload.clone(), 3).is_err());
973 let payload = payload.unwrap();
974 assert!(
975 decoder
976 .decompress(Some(payload.slice_with_length(0, payload.len() - 1)), 2)
977 .is_err()
978 );
979 }
980
981 fn miniblock_context() -> MiniBlockCompressionContext {
982 MiniBlockCompressionContext::new(0, true, true)
983 }
984
985 const PRIMITIVE_TYPES: &[DataType] = &[
986 DataType::Null,
987 DataType::FixedSizeBinary(2),
988 DataType::Date32,
989 DataType::Date64,
990 DataType::Int8,
991 DataType::Int16,
992 DataType::Int32,
993 DataType::Int64,
994 DataType::UInt8,
995 DataType::UInt16,
996 DataType::UInt32,
997 DataType::UInt64,
998 DataType::Float16,
999 DataType::Float32,
1000 DataType::Float64,
1001 DataType::Decimal128(10, 10),
1002 DataType::Decimal256(10, 10),
1003 DataType::Timestamp(TimeUnit::Nanosecond, None),
1004 DataType::Time32(TimeUnit::Second),
1005 DataType::Time64(TimeUnit::Nanosecond),
1006 DataType::Duration(TimeUnit::Second),
1007 ];
1011
1012 #[test_log::test(tokio::test)]
1013 async fn test_simple_value() {
1014 let items = Arc::new(Int32Array::from(vec![
1015 Some(0),
1016 None,
1017 Some(2),
1018 Some(3),
1019 Some(4),
1020 Some(5),
1021 ]));
1022
1023 let test_cases = TestCases::default()
1024 .with_range(0..3)
1025 .with_range(0..2)
1026 .with_range(1..3)
1027 .with_indices(vec![0, 1, 2])
1028 .with_indices(vec![1])
1029 .with_indices(vec![2])
1030 .with_structural_encodings();
1031
1032 check_round_trip_encoding_of_data(vec![items], &test_cases, HashMap::default()).await;
1033 }
1034
1035 #[test_log::test(tokio::test)]
1036 async fn test_simple_range() {
1037 let items = Arc::new(Int32Array::from_iter(
1038 (0..5000).map(|i| if i % 2 == 0 { Some(i) } else { None }),
1039 ));
1040
1041 let test_cases = TestCases::default().with_structural_encodings();
1042
1043 check_round_trip_encoding_of_data(vec![items], &test_cases, HashMap::default()).await;
1044 }
1045
1046 #[test_log::test(tokio::test)]
1047 async fn test_value_primitive() {
1048 const NUM_ROWS: u32 = 1025;
1049
1050 let test_cases = TestCases::default()
1051 .with_batch_size(NUM_ROWS)
1052 .with_page_sizes(vec![4096])
1053 .with_expected_encoding("flat");
1054 let value_metadata =
1055 HashMap::from([("lance-encoding:compression".to_string(), "none".to_string())]);
1056
1057 for (seed, data_type) in PRIMITIVE_TYPES.iter().enumerate() {
1058 log::info!("Testing encoding for {:?}", data_type);
1059 let data = gen_batch()
1060 .with_seed(Seed::from(seed as u64))
1061 .anon_col(array::rand_type(data_type))
1062 .into_batch_rows(RowCount::from(NUM_ROWS as u64))
1063 .unwrap()
1064 .column(0)
1065 .clone();
1066
1067 check_round_trip_encoding_of_data(vec![data], &test_cases, value_metadata.clone())
1068 .await;
1069 }
1070 }
1071
1072 #[rstest::rstest]
1073 #[test_log::test(tokio::test)]
1074 async fn test_large_primitive(
1075 #[values(
1076 TestEncoding::Array,
1077 TestEncoding::StructuralU16,
1078 TestEncoding::StructuralU32,
1079 TestEncoding::StructuralSparse
1080 )]
1081 encoding: TestEncoding,
1082 #[values(4096, 1024 * 1024)] page_size: u64,
1083 #[values(false, true)] use_slicing: bool,
1084 ) {
1085 let data_type =
1086 DataType::FixedSizeList(Arc::new(Field::new("", DataType::Int32, false)), 128);
1087 let field = Field::new("", data_type, false);
1088 check_basic_random_case(field, encoding, page_size, use_slicing).await;
1089 }
1090
1091 #[rstest::rstest]
1092 #[test_log::test(tokio::test)]
1093 async fn test_decimal128_dictionary_encoding(
1094 #[values(
1095 TestEncoding::StructuralU16,
1096 TestEncoding::StructuralU32,
1097 TestEncoding::StructuralSparse
1098 )]
1099 encoding: TestEncoding,
1100 ) {
1101 let test_cases = TestCases::default()
1102 .with_encoding(encoding)
1103 .with_expected_encoding("dictionary");
1104 let decimals: Vec<i32> = (0..100).collect();
1105 let repeated_strings: Vec<_> = decimals
1106 .iter()
1107 .cycle()
1108 .take(decimals.len() * 1000)
1109 .map(|&v| Some(v as i128))
1110 .collect();
1111 let decimal_array = Arc::new(Decimal128Array::from(repeated_strings)) as ArrayRef;
1112 check_round_trip_encoding_of_data(vec![decimal_array], &test_cases, HashMap::new()).await;
1113 }
1114
1115 #[rstest::rstest]
1116 #[test_log::test(tokio::test)]
1117 async fn test_miniblock_stress(
1118 #[values(false, true)] mixed_validity: bool,
1119 #[values(10, 100, 1500, 15000)] batch_size: u32,
1120 #[values(1000, 2000, 3000, 60000)] page_size: u64,
1121 #[values(
1122 TestEncoding::StructuralU16,
1123 TestEncoding::StructuralU32,
1124 TestEncoding::StructuralSparse
1125 )]
1126 encoding: TestEncoding,
1127 ) {
1128 let data = (0..100)
1132 .map(|_| {
1133 if mixed_validity {
1134 Arc::new(Int32Array::from_iter(
1135 (0..100).map(|i| if i % 2 == 0 { Some(i) } else { None }),
1136 )) as Arc<dyn Array>
1137 } else {
1138 Arc::new(Int32Array::from_iter_values(0..100)) as Arc<dyn Array>
1139 }
1140 })
1141 .collect::<Vec<_>>();
1142
1143 let test_cases = TestCases::default()
1144 .with_page_sizes(vec![page_size])
1145 .with_batch_size(batch_size)
1146 .with_encoding(encoding);
1147
1148 check_round_trip_encoding_of_data(data, &test_cases, HashMap::new()).await;
1149 }
1150
1151 fn create_simple_fsl() -> FixedSizeListArray {
1152 let items = Arc::new(Int32Array::from(vec![
1154 Some(0),
1155 Some(1),
1156 Some(2),
1157 Some(3),
1158 None,
1159 None,
1160 None,
1161 None,
1162 Some(8),
1163 Some(9),
1164 None,
1165 Some(11),
1166 ]));
1167 let items_field = Arc::new(Field::new("item", DataType::Int32, true));
1168 let inner_list_nulls = BooleanBuffer::from(vec![true, false, false, false, true, true]);
1169 let inner_list = Arc::new(FixedSizeListArray::new(
1170 items_field.clone(),
1171 2,
1172 items,
1173 Some(NullBuffer::new(inner_list_nulls)),
1174 ));
1175 let inner_list_field = Arc::new(Field::new(
1176 "item",
1177 DataType::FixedSizeList(items_field, 2),
1178 true,
1179 ));
1180 FixedSizeListArray::new(inner_list_field, 2, inner_list, None)
1181 }
1182
1183 #[test]
1184 fn test_fsl_value_compression_miniblock() {
1185 let sample_list = create_simple_fsl();
1186
1187 let starting_data = DataBlock::from_array(sample_list.clone());
1188
1189 let encoder = ValueEncoder::default();
1190 let (data, compression) =
1191 MiniBlockCompressor::compress(&encoder, miniblock_context(), starting_data).unwrap();
1192
1193 assert_eq!(data.num_values, 3);
1194 assert_eq!(data.data.len(), 3);
1195 assert_eq!(data.chunks.len(), 1);
1196 assert_eq!(data.chunks[0].buffer_sizes, vec![1, 2, 48]);
1197 assert_eq!(data.chunks[0].log_num_values, 0);
1198
1199 let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else {
1200 panic!()
1201 };
1202
1203 let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap();
1204
1205 let decompressed =
1206 MiniBlockDecompressor::decompress(&decompressor, data.data, data.num_values).unwrap();
1207
1208 let decompressed = make_array(
1209 decompressed
1210 .into_arrow(sample_list.data_type().clone(), true)
1211 .unwrap(),
1212 );
1213
1214 assert_eq!(decompressed.as_ref(), &sample_list);
1215 }
1216
1217 fn wide_fixed_size_binary() -> ArrayRef {
1218 let wide_value = vec![0xABu8; 5000];
1219 Arc::new(
1220 arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
1221 std::iter::repeat_n(Some(wide_value.as_slice()), 4),
1222 5000,
1223 )
1224 .unwrap(),
1225 )
1226 }
1227
1228 fn wide_fixed_size_list_bool() -> ArrayRef {
1229 let dimension = 4095;
1232 let values = arrow_array::BooleanArray::from(vec![false; dimension * 2]);
1233 let field = Arc::new(Field::new("item", DataType::Boolean, true));
1234 Arc::new(FixedSizeListArray::new(
1235 field,
1236 dimension as i32,
1237 Arc::new(values),
1238 None,
1239 ))
1240 }
1241
1242 #[rstest::rstest]
1243 #[case::fixed_size_binary(wide_fixed_size_binary(), 2)]
1244 #[case::fixed_size_list_bool(wide_fixed_size_list_bool(), 16)]
1245 fn test_wide_value_miniblock_returns_error(
1246 #[case] array: ArrayRef,
1247 #[case] expected_min_values: u64,
1248 ) {
1249 let starting_data = DataBlock::from_array(array);
1250
1251 let encoder = ValueEncoder::default();
1252 let result = MiniBlockCompressor::compress(&encoder, miniblock_context(), starting_data);
1253
1254 let err = result.expect_err("wide values should not be encodable as miniblock");
1255 assert!(
1256 matches!(err, lance_core::Error::InvalidInput { .. }),
1257 "expected InvalidInput, got {err:?}"
1258 );
1259 let msg = err.to_string();
1260 assert!(
1261 msg.contains("too wide for miniblock encoding"),
1262 "unexpected error message: {msg}"
1263 );
1264 assert!(
1265 msg.contains(&format!("{expected_min_values} values require")),
1266 "unexpected error message: {msg}"
1267 );
1268 }
1269
1270 #[test]
1271 fn test_fsl_value_compression_per_value() {
1272 let sample_list = create_simple_fsl();
1273
1274 let starting_data = DataBlock::from_array(sample_list.clone());
1275
1276 let encoder = ValueEncoder::default();
1277 let (data, compression) = PerValueCompressor::compress(&encoder, starting_data).unwrap();
1278
1279 let PerValueDataBlock::Fixed(data) = data else {
1280 panic!()
1281 };
1282
1283 assert_eq!(data.bits_per_value, 144);
1284 assert_eq!(data.num_values, 3);
1285 assert_eq!(data.data.len(), 18 * 3);
1286
1287 let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else {
1288 panic!()
1289 };
1290
1291 let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap();
1292
1293 let num_values = data.num_values;
1294 assert_eq!(
1295 FixedPerValueDecompressor::decoded_size_bytes(&decompressor, num_values),
1296 None,
1297 "nullable FSL output uses multiple buffers and requires the fallback estimate"
1298 );
1299 let decompressed =
1300 FixedPerValueDecompressor::decompress(&decompressor, data, num_values).unwrap();
1301
1302 let decompressed = make_array(
1303 decompressed
1304 .into_arrow(sample_list.data_type().clone(), true)
1305 .unwrap(),
1306 );
1307
1308 assert_eq!(decompressed.as_ref(), &sample_list);
1309 }
1310
1311 #[rstest::rstest]
1314 #[case::outer_wider(22, 3)]
1315 #[case::inner_wider(3, 22)]
1316 fn test_nested_fsl_simple_per_value(#[case] outer_dim: i32, #[case] inner_dim: i32) {
1317 let num_rows = 3;
1318 let inner_field = Arc::new(Field::new("item", DataType::Int32, false));
1319 let leaf = Arc::new(Int32Array::from_iter_values(
1320 0..num_rows * outer_dim * inner_dim,
1321 )) as ArrayRef;
1322 let inner = Arc::new(FixedSizeListArray::new(
1323 inner_field.clone(),
1324 inner_dim,
1325 leaf,
1326 None,
1327 )) as ArrayRef;
1328 let outer = Arc::new(FixedSizeListArray::new(
1329 Arc::new(Field::new(
1330 "item",
1331 DataType::FixedSizeList(inner_field, inner_dim),
1332 false,
1333 )),
1334 outer_dim,
1335 inner,
1336 None,
1337 )) as ArrayRef;
1338
1339 let encoder = ValueEncoder::default();
1340 let (data, compression) =
1341 PerValueCompressor::compress(&encoder, DataBlock::from_array(outer.clone())).unwrap();
1342 let PerValueDataBlock::Fixed(data) = data else {
1343 panic!()
1344 };
1345 assert_eq!(data.num_values, num_rows as u64);
1346 assert_eq!(data.bits_per_value, 32 * (outer_dim * inner_dim) as u64);
1347
1348 let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else {
1349 panic!()
1350 };
1351 let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap();
1352 let num_values = data.num_values;
1353 let decompressed =
1354 FixedPerValueDecompressor::decompress(&decompressor, data, num_values).unwrap();
1355 let decompressed = make_array(
1356 decompressed
1357 .into_arrow(outer.data_type().clone(), true)
1358 .unwrap(),
1359 );
1360 assert_eq!(decompressed.as_ref(), outer.as_ref());
1361 }
1362
1363 #[test_log::test(tokio::test)]
1364 async fn test_fsl_all_null() {
1365 let items = new_null_array(&DataType::Int32, 12);
1366 let items_field = Arc::new(Field::new("item", DataType::Int32, true));
1367 let list_nulls = BooleanBuffer::from(vec![true, false, false, false, true, true]);
1368 let list_array =
1369 FixedSizeListArray::new(items_field, 2, items, Some(NullBuffer::new(list_nulls)));
1370
1371 let test_cases = TestCases::default().with_structural_encodings();
1372
1373 check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
1374 .await;
1375 }
1376
1377 #[rstest::rstest]
1380 #[test_log::test(tokio::test)]
1381 async fn test_fsl_of_null_type(#[values(false, true)] has_outer_nulls: bool) {
1382 let items_field = Arc::new(Field::new("item", DataType::Null, true));
1383 let outer_nulls = has_outer_nulls.then(|| NullBuffer::from(vec![true, false, true]));
1384 let list_array =
1385 FixedSizeListArray::new(items_field, 2, Arc::new(NullArray::new(6)), outer_nulls);
1386
1387 let test_cases = TestCases::default().with_structural_encodings();
1388
1389 check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
1390 .await;
1391 }
1392
1393 #[test_log::test(tokio::test)]
1397 async fn test_fsl_nullable_child_all_null_no_outer_nulls() {
1398 let items = new_null_array(&DataType::Float32, 8);
1403 let items_field = Arc::new(Field::new("item", DataType::Float32, true));
1404 let fsl = FixedSizeListArray::new(items_field, 4, items, None);
1405
1406 let test_cases = TestCases::default().with_structural_encodings();
1407 check_round_trip_encoding_of_data(vec![Arc::new(fsl)], &test_cases, HashMap::new()).await;
1408 }
1409
1410 #[test_log::test(tokio::test)]
1411 async fn regress_list_fsl() {
1412 let offsets = ScalarBuffer::<i32>::from(vec![0, 393, 755, 1156, 1536]);
1415 let data = UInt8Array::from(vec![0; 1536 * 16]);
1416 let fsl_field = Arc::new(Field::new("item", DataType::UInt8, true));
1417 let fsl = FixedSizeListArray::new(fsl_field, 16, Arc::new(data), None);
1418 let list_field = Arc::new(Field::new("item", fsl.data_type().clone(), false));
1419 let list_arr = ListArray::new(list_field, OffsetBuffer::new(offsets), Arc::new(fsl), None);
1420
1421 let test_cases = TestCases::default()
1422 .with_structural_encodings()
1423 .with_batch_size(1);
1424
1425 check_round_trip_encoding_of_data(vec![Arc::new(list_arr)], &test_cases, HashMap::new())
1426 .await;
1427 }
1428
1429 fn create_random_fsl() -> Arc<dyn Array> {
1430 let inner = array::rand_type(&DataType::Int32).with_random_nulls(0.1);
1432 let list_one = array::cycle_vec(inner, Dimension::from(4)).with_random_nulls(0.1);
1433 let list_two = array::cycle_vec(list_one, Dimension::from(4)).with_random_nulls(0.1);
1434 let list_three = array::cycle_vec(list_two, Dimension::from(2));
1435
1436 let batch = gen_batch()
1438 .anon_col(list_three)
1439 .into_batch_rows(RowCount::from(8 * 1024))
1440 .unwrap();
1441 batch.column(0).clone()
1442 }
1443
1444 #[test]
1445 fn fsl_value_miniblock_stress() {
1446 let sample_array = create_random_fsl();
1447
1448 let starting_data = DataBlock::from_arrays(
1449 std::slice::from_ref(&sample_array),
1450 sample_array.len() as u64,
1451 );
1452
1453 let encoder = ValueEncoder::default();
1454 let (data, compression) =
1455 MiniBlockCompressor::compress(&encoder, miniblock_context(), starting_data).unwrap();
1456
1457 let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else {
1458 panic!()
1459 };
1460
1461 let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap();
1462
1463 let decompressed =
1464 MiniBlockDecompressor::decompress(&decompressor, data.data, data.num_values).unwrap();
1465
1466 let decompressed = make_array(
1467 decompressed
1468 .into_arrow(sample_array.data_type().clone(), true)
1469 .unwrap(),
1470 );
1471
1472 assert_eq!(decompressed.as_ref(), sample_array.as_ref());
1473 }
1474
1475 #[test]
1476 fn fsl_value_per_value_stress() {
1477 let sample_array = create_random_fsl();
1478
1479 let starting_data = DataBlock::from_arrays(
1480 std::slice::from_ref(&sample_array),
1481 sample_array.len() as u64,
1482 );
1483
1484 let encoder = ValueEncoder::default();
1485 let (data, compression) = PerValueCompressor::compress(&encoder, starting_data).unwrap();
1486
1487 let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else {
1488 panic!()
1489 };
1490
1491 let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap();
1492
1493 let PerValueDataBlock::Fixed(data) = data else {
1494 panic!()
1495 };
1496
1497 let num_values = data.num_values;
1498 let decompressed =
1499 FixedPerValueDecompressor::decompress(&decompressor, data, num_values).unwrap();
1500
1501 let decompressed = make_array(
1502 decompressed
1503 .into_arrow(sample_array.data_type().clone(), true)
1504 .unwrap(),
1505 );
1506
1507 assert_eq!(decompressed.as_ref(), sample_array.as_ref());
1508 }
1509
1510 #[rstest::rstest]
1511 #[test_log::test(tokio::test)]
1512 async fn test_fsl_nullable_items(
1513 #[values(
1514 TestEncoding::Array,
1515 TestEncoding::StructuralU16,
1516 TestEncoding::StructuralU32,
1517 TestEncoding::StructuralSparse
1518 )]
1519 encoding: TestEncoding,
1520 ) {
1521 let mut generator =
1522 gen_batch()
1523 .with_seed(Seed::from(0))
1524 .anon_col(array::rand_vec_nullable::<UInt32Type>(
1525 Dimension::from(128),
1526 0.5,
1527 ));
1528 generator.with_random_nulls(0.2);
1529 let source = generator
1530 .into_batch_rows(RowCount::from(1026))
1531 .unwrap()
1532 .column(0)
1533 .clone();
1534 let test_cases = TestCases::default()
1535 .with_page_sizes(vec![4096])
1536 .with_encoding(encoding)
1537 .with_batch_size(257)
1538 .with_range(510..515)
1539 .with_indices(vec![0, 511, 512, 1024]);
1540
1541 check_round_trip_encoding_of_data(
1542 vec![source.slice(1, 512), source.slice(513, 513)],
1543 &test_cases,
1544 HashMap::new(),
1545 )
1546 .await;
1547 }
1548
1549 #[test_log::test(tokio::test)]
1550 async fn test_value_encoding_verification() {
1551 use std::collections::HashMap;
1552
1553 let test_cases = TestCases::default()
1554 .with_expected_encoding("flat")
1555 .with_structural_encodings();
1556
1557 let mut metadata_explicit = HashMap::new();
1561 metadata_explicit.insert("lance-encoding:compression".to_string(), "none".to_string());
1562 metadata_explicit.insert("lance-encoding:bss".to_string(), "off".to_string());
1563
1564 let arr_explicit =
1565 Arc::new(Int32Array::from((0..1000).collect::<Vec<i32>>())) as Arc<dyn Array>;
1566 check_round_trip_encoding_of_data(vec![arr_explicit], &test_cases, metadata_explicit).await;
1567
1568 let mut metadata = HashMap::new();
1572 metadata.insert("lance-encoding:bss".to_string(), "off".to_string());
1573
1574 let arr_fallback = Arc::new(Int32Array::from(
1575 (0..100).map(|i| i * 73 + 19).collect::<Vec<i32>>(),
1576 )) as Arc<dyn Array>;
1577 check_round_trip_encoding_of_data(vec![arr_fallback], &test_cases, metadata).await;
1578 }
1579
1580 #[test_log::test(tokio::test)]
1581 async fn test_mixed_page_validity() {
1582 let no_nulls = Arc::new(Int32Array::from_iter_values([1, 2]));
1583 let has_nulls = Arc::new(Int32Array::from_iter([Some(3), None, Some(5)]));
1584
1585 let test_cases = TestCases::default().with_page_sizes(vec![1]);
1586 check_round_trip_encoding_of_data(vec![no_nulls, has_nulls], &test_cases, HashMap::new())
1587 .await;
1588 }
1589}