Skip to main content

lance_encoding/encodings/logical/
primitive.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::{
5    any::Any,
6    collections::{HashMap, VecDeque},
7    env,
8    fmt::Debug,
9    iter,
10    ops::Range,
11    sync::Arc,
12    vec,
13};
14
15use crate::{
16    constants::{
17        STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK,
18    },
19    data::DictionaryDataBlock,
20    encodings::logical::primitive::blob::{BlobDescriptionPageScheduler, BlobPageScheduler},
21    format::{
22        ProtobufUtils21,
23        pb21::{self, CompressiveEncoding, PageLayout, compressive_encoding::Compression},
24    },
25};
26use arrow_array::{Array, ArrayRef, PrimitiveArray, cast::AsArray, make_array, types::UInt64Type};
27use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBuffer, ScalarBuffer};
28use arrow_schema::{DataType, Field as ArrowField};
29use bytes::Bytes;
30use futures::{FutureExt, TryStreamExt, future::BoxFuture, stream::FuturesOrdered};
31use itertools::Itertools;
32use lance_arrow::DataTypeExt;
33use lance_arrow::deepcopy::deep_copy_nulls;
34use lance_core::{
35    cache::{CacheKey, Context, DeepSizeOf},
36    error::{Error, LanceOptionExt},
37    utils::bit::pad_bytes,
38};
39use log::trace;
40
41use crate::encodings::logical::primitive::miniblock::MiniBlockChunk;
42use crate::utils::bytepack::ByteUnpacker;
43use crate::{
44    compression::{
45        BlockDecompressor, CompressionStrategy, DecompressionStrategy, MiniBlockDecompressor,
46    },
47    data::{AllNullDataBlock, DataBlock, VariableWidthBlock},
48    utils::bytepack::BytepackedIntegerEncoder,
49};
50use crate::{
51    compression::{FixedPerValueDecompressor, VariablePerValueDecompressor},
52    encodings::logical::primitive::fullzip::PerValueDataBlock,
53};
54use crate::{
55    encodings::logical::primitive::miniblock::MiniBlockCompressed,
56    statistics::{ComputeStat, GetStat, Stat},
57};
58use crate::{
59    repdef::{
60        CompositeRepDefUnraveler, ControlWordIterator, ControlWordParser, DefinitionInterpretation,
61        MiniBlockRepDefBudget, RepDefSlicer, SerializedRepDefs, build_control_word_iterator,
62    },
63    utils::accumulation::AccumulationQueue,
64};
65use lance_core::{Result, datatypes::Field, utils::tokio::spawn_cpu};
66
67use crate::constants::{
68    COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_DIVISOR_META_KEY,
69    DICT_SIZE_RATIO_META_KEY, DICT_VALUES_COMPRESSION_ENV_VAR,
70    DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR, DICT_VALUES_COMPRESSION_LEVEL_META_KEY,
71    DICT_VALUES_COMPRESSION_META_KEY,
72};
73use crate::version::LanceFileVersion;
74use crate::{
75    EncodingsIo,
76    buffer::LanceBuffer,
77    data::{BlockInfo, DataBlockBuilder, FixedWidthDataBlock},
78    decoder::{
79        ColumnInfo, DecodePageTask, DecodedArray, DecodedPage, FilterExpression, LoadedPageShard,
80        MessageType, PageEncoding, PageInfo, ScheduledScanLine, SchedulerContext,
81        StructuralDecodeArrayTask, StructuralFieldDecoder, StructuralFieldScheduler,
82        StructuralPageDecoder, StructuralSchedulingJob, UnloadedPageShard,
83    },
84    encoder::{
85        EncodeTask, EncodedColumn, EncodedPage, EncodingOptions, FieldEncoder, OutOfLineBuffers,
86    },
87    repdef::{LevelBuffer, RepDefBuilder, RepDefUnraveler},
88};
89
90pub mod blob;
91pub mod constant;
92pub mod dict;
93pub mod fullzip;
94pub mod miniblock;
95
96const FILL_BYTE: u8 = 0xFE;
97const DEFAULT_DICT_DIVISOR: u64 = 2;
98const DEFAULT_DICT_MAX_CARDINALITY: u64 = 100_000;
99const DEFAULT_DICT_SIZE_RATIO: f64 = 0.8;
100const DEFAULT_DICT_VALUES_COMPRESSION: &str = "lz4";
101
102struct PageLoadTask {
103    decoder_fut: BoxFuture<'static, Result<Box<dyn StructuralPageDecoder>>>,
104    num_rows: u64,
105}
106
107/// A trait for figuring out how to schedule the data within
108/// a single page.
109trait StructuralPageScheduler: std::fmt::Debug + Send {
110    /// Fetches any metadata required for the page
111    fn initialize<'a>(
112        &'a mut self,
113        io: &Arc<dyn EncodingsIo>,
114    ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>>;
115    /// Loads metadata from a previous initialize call
116    fn load(&mut self, data: &Arc<dyn CachedPageData>);
117    /// Schedules the read of the given ranges in the page
118    ///
119    /// The read may be split into multiple "shards" if the page is extremely large.
120    /// Each shard maps to one or more rows and can be decoded independently.
121    ///
122    /// Note: this sharding is for splitting up very large pages into smaller reads to
123    /// avoid buffering too much data in memory.  It is not related to the batch size or
124    /// compute units in any way.
125    fn schedule_ranges(
126        &self,
127        ranges: &[Range<u64>],
128        io: &Arc<dyn EncodingsIo>,
129    ) -> Result<Vec<PageLoadTask>>;
130}
131
132/// Metadata describing the decoded size of a mini-block
133#[derive(Debug)]
134struct ChunkMeta {
135    num_values: u64,
136    chunk_size_bytes: u64,
137    offset_bytes: u64,
138}
139
140/// A mini-block chunk that has been decoded and decompressed
141#[derive(Debug, Clone)]
142struct DecodedMiniBlockChunk {
143    rep: Option<ScalarBuffer<u16>>,
144    def: Option<ScalarBuffer<u16>>,
145    values: DataBlock,
146}
147
148/// A task to decode a one or more mini-blocks of data into an output batch
149///
150/// Note: Two batches might share the same mini-block of data.  When this happens
151/// then each batch gets a copy of the block and each batch decodes the block independently.
152///
153/// This means we have duplicated work but it is necessary to avoid having to synchronize
154/// the decoding of the block. (TODO: test this theory)
155#[derive(Debug)]
156struct DecodeMiniBlockTask {
157    rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
158    def_decompressor: Option<Arc<dyn BlockDecompressor>>,
159    value_decompressor: Arc<dyn MiniBlockDecompressor>,
160    dictionary_data: Option<Arc<DataBlock>>,
161    def_meaning: Arc<[DefinitionInterpretation]>,
162    num_buffers: u64,
163    max_visible_level: u16,
164    instructions: Vec<(ChunkDrainInstructions, LoadedChunk)>,
165    has_large_chunk: bool,
166}
167
168impl DecodeMiniBlockTask {
169    fn decode_levels(
170        rep_decompressor: &dyn BlockDecompressor,
171        levels: LanceBuffer,
172        num_levels: u16,
173    ) -> Result<ScalarBuffer<u16>> {
174        let rep = rep_decompressor.decompress(levels, num_levels as u64)?;
175        let rep = rep.as_fixed_width().unwrap();
176        debug_assert_eq!(rep.num_values, num_levels as u64);
177        debug_assert_eq!(rep.bits_per_value, 16);
178        Ok(rep.data.borrow_to_typed_slice::<u16>())
179    }
180
181    // We are building a LevelBuffer (levels) and want to copy into it `total_len`
182    // values from `level_buf` starting at `offset`.
183    //
184    // We need to handle both the case where `levels` is None (no nulls encountered
185    // yet) and the case where `level_buf` is None (the input we are copying from has
186    // no nulls)
187    fn extend_levels(
188        range: Range<u64>,
189        levels: &mut Option<LevelBuffer>,
190        level_buf: &Option<impl AsRef<[u16]>>,
191        dest_offset: usize,
192    ) {
193        if let Some(level_buf) = level_buf {
194            if levels.is_none() {
195                // This is the first non-empty def buf we've hit, fill in the past
196                // with 0 (valid)
197                let mut new_levels_vec =
198                    LevelBuffer::with_capacity(dest_offset + (range.end - range.start) as usize);
199                new_levels_vec.extend(iter::repeat_n(0, dest_offset));
200                *levels = Some(new_levels_vec);
201            }
202            levels.as_mut().unwrap().extend(
203                level_buf.as_ref()[range.start as usize..range.end as usize]
204                    .iter()
205                    .copied(),
206            );
207        } else if let Some(levels) = levels {
208            let num_values = (range.end - range.start) as usize;
209            // This is an all-valid level_buf but we had nulls earlier and so we
210            // need to materialize it
211            levels.extend(iter::repeat_n(0, num_values));
212        }
213    }
214
215    /// Maps a range of rows to a range of items and a range of levels
216    ///
217    /// If there is no repetition information this just returns the range as-is.
218    ///
219    /// If there is repetition information then we need to do some work to figure out what
220    /// range of items corresponds to the requested range of rows.
221    ///
222    /// For example, if the data is [[1, 2, 3], [4, 5], [6, 7]] and the range is 1..2 (i.e. just row
223    /// 1) then the user actually wants items 3..5.  In the above case the rep levels would be:
224    ///
225    /// Idx: 0 1 2 3 4 5 6
226    /// Rep: 1 0 0 1 0 1 0
227    ///
228    /// So the start (1) maps to the second 1 (idx=3) and the end (2) maps to the third 1 (idx=5)
229    ///
230    /// If there are invisible items then we don't count them when calculating the range of items we
231    /// are interested in but we do count them when calculating the range of levels we are interested
232    /// in.  As a result we have to return both the item range (first return value) and the level range
233    /// (second return value).
234    ///
235    /// For example, if the data is [[1, 2, 3], [4, 5], NULL, [6, 7, 8]] and the range is 2..4 then the
236    /// user wants items 5..8 but they want levels 5..9.  In the above case the rep/def levels would be:
237    ///
238    /// Idx: 0 1 2 3 4 5 6 7 8
239    /// Rep: 1 0 0 1 0 1 1 0 0
240    /// Def: 0 0 0 0 0 1 0 0 0
241    /// Itm: 1 2 3 4 5 6 7 8
242    ///
243    /// Finally, we have to contend with the fact that chunks may or may not start with a "preamble" of
244    /// trailing values that finish up a list from the previous chunk.  In this case the first item does
245    /// not start at max_rep because it is a continuation of the previous chunk.  For our purposes we do
246    /// not consider this a "row" and so the range 0..1 will refer to the first row AFTER the preamble.
247    ///
248    /// We have a separate parameter (`preamble_action`) to control whether we want the preamble or not.
249    ///
250    /// Note that the "trailer" is considered a "row" and if we want it we should include it in the range.
251    fn map_range(
252        range: Range<u64>,
253        rep: Option<&impl AsRef<[u16]>>,
254        def: Option<&impl AsRef<[u16]>>,
255        max_rep: u16,
256        max_visible_def: u16,
257        // The total number of items (not rows) in the chunk.  This is not quite the same as
258        // rep.len() / def.len() because it doesn't count invisible items
259        total_items: u64,
260        preamble_action: PreambleAction,
261    ) -> (Range<u64>, Range<u64>) {
262        if let Some(rep) = rep {
263            let mut rep = rep.as_ref();
264            // If there is a preamble and we need to skip it then do that first.  The work is the same
265            // whether there is def information or not
266            let mut items_in_preamble = 0_u64;
267            let first_row_start = match preamble_action {
268                PreambleAction::Skip | PreambleAction::Take => {
269                    let first_row_start = if let Some(def) = def.as_ref() {
270                        let mut first_row_start = None;
271                        for (idx, (rep, def)) in rep.iter().zip(def.as_ref()).enumerate() {
272                            if *rep == max_rep {
273                                first_row_start = Some(idx as u64);
274                                break;
275                            }
276                            if *def <= max_visible_def {
277                                items_in_preamble += 1;
278                            }
279                        }
280                        first_row_start
281                    } else {
282                        let first_row_start =
283                            rep.iter().position(|&r| r == max_rep).map(|r| r as u64);
284                        items_in_preamble = first_row_start.unwrap_or(rep.len() as u64);
285                        first_row_start
286                    };
287                    // It is possible for a chunk to be entirely partial values but if it is then it
288                    // should never show up as a preamble to skip
289                    if first_row_start.is_none() {
290                        assert!(preamble_action == PreambleAction::Take);
291                        return (0..total_items, 0..rep.len() as u64);
292                    }
293                    let first_row_start = first_row_start.unwrap();
294                    rep = &rep[first_row_start as usize..];
295                    first_row_start
296                }
297                PreambleAction::Absent => {
298                    debug_assert!(rep[0] == max_rep);
299                    0
300                }
301            };
302
303            // We hit this case when all we needed was the preamble
304            if range.start == range.end {
305                debug_assert!(preamble_action == PreambleAction::Take);
306                debug_assert!(items_in_preamble <= total_items);
307                return (0..items_in_preamble, 0..first_row_start);
308            }
309            assert!(range.start < range.end);
310
311            let mut rows_seen = 0;
312            let mut new_start = 0;
313            let mut new_levels_start = 0;
314
315            if let Some(def) = def {
316                let def = &def.as_ref()[first_row_start as usize..];
317
318                // range.start == 0 always maps to 0 (even with invis items), otherwise we need to walk
319                let mut lead_invis_seen = 0;
320
321                if range.start > 0 {
322                    if def[0] > max_visible_def {
323                        lead_invis_seen += 1;
324                    }
325                    for (idx, (rep, def)) in rep.iter().zip(def).skip(1).enumerate() {
326                        if *rep == max_rep {
327                            rows_seen += 1;
328                            if rows_seen == range.start {
329                                new_start = idx as u64 + 1 - lead_invis_seen;
330                                new_levels_start = idx as u64 + 1;
331                                break;
332                            }
333                        }
334                        if *def > max_visible_def {
335                            lead_invis_seen += 1;
336                        }
337                    }
338                }
339
340                rows_seen += 1;
341
342                let mut new_end = u64::MAX;
343                let mut new_levels_end = rep.len() as u64;
344                let new_start_is_visible = def[new_levels_start as usize] <= max_visible_def;
345                let mut tail_invis_seen = if new_start_is_visible { 0 } else { 1 };
346                for (idx, (rep, def)) in rep[(new_levels_start + 1) as usize..]
347                    .iter()
348                    .zip(&def[(new_levels_start + 1) as usize..])
349                    .enumerate()
350                {
351                    if *rep == max_rep {
352                        rows_seen += 1;
353                        if rows_seen == range.end + 1 {
354                            new_end = idx as u64 + new_start + 1 - tail_invis_seen;
355                            new_levels_end = idx as u64 + new_levels_start + 1;
356                            break;
357                        }
358                    }
359                    if *def > max_visible_def {
360                        tail_invis_seen += 1;
361                    }
362                }
363
364                if new_end == u64::MAX {
365                    new_levels_end = rep.len() as u64;
366                    let total_invis_seen = lead_invis_seen + tail_invis_seen;
367                    new_end = rep.len() as u64 - total_invis_seen;
368                }
369
370                assert_ne!(new_end, u64::MAX);
371
372                // Adjust for any skipped preamble
373                if preamble_action == PreambleAction::Skip {
374                    new_start += items_in_preamble;
375                    new_end += items_in_preamble;
376                    new_levels_start += first_row_start;
377                    new_levels_end += first_row_start;
378                } else if preamble_action == PreambleAction::Take {
379                    debug_assert_eq!(new_start, 0);
380                    debug_assert_eq!(new_levels_start, 0);
381                    new_end += items_in_preamble;
382                    new_levels_end += first_row_start;
383                }
384
385                debug_assert!(new_end <= total_items);
386                (new_start..new_end, new_levels_start..new_levels_end)
387            } else {
388                // Easy case, there are no invisible items, so we don't need to check for them
389                // The items range and levels range will be the same.  We do still need to walk
390                // the rep levels to find the row boundaries
391
392                // range.start == 0 always maps to 0, otherwise we need to walk
393                if range.start > 0 {
394                    for (idx, rep) in rep.iter().skip(1).enumerate() {
395                        if *rep == max_rep {
396                            rows_seen += 1;
397                            if rows_seen == range.start {
398                                new_start = idx as u64 + 1;
399                                break;
400                            }
401                        }
402                    }
403                }
404                let mut new_end = rep.len() as u64;
405                // range.end == max_items always maps to rep.len(), otherwise we need to walk
406                if range.end < total_items {
407                    for (idx, rep) in rep[(new_start + 1) as usize..].iter().enumerate() {
408                        if *rep == max_rep {
409                            rows_seen += 1;
410                            if rows_seen == range.end {
411                                new_end = idx as u64 + new_start + 1;
412                                break;
413                            }
414                        }
415                    }
416                }
417
418                // Adjust for any skipped preamble
419                if preamble_action == PreambleAction::Skip {
420                    new_start += first_row_start;
421                    new_end += first_row_start;
422                } else if preamble_action == PreambleAction::Take {
423                    debug_assert_eq!(new_start, 0);
424                    new_end += first_row_start;
425                }
426
427                debug_assert!(new_end <= total_items);
428                (new_start..new_end, new_start..new_end)
429            }
430        } else {
431            // No repetition info, easy case, just use the range as-is and the item
432            // and level ranges are the same
433            (range.clone(), range)
434        }
435    }
436
437    // read `num_buffers` buffer sizes from `buf` starting at `offset`
438    fn read_buffer_sizes<const LARGE: bool>(
439        buf: &[u8],
440        offset: &mut usize,
441        num_buffers: u64,
442    ) -> Vec<u32> {
443        let read_size = if LARGE { 4 } else { 2 };
444        (0..num_buffers)
445            .map(|_| {
446                let bytes = &buf[*offset..*offset + read_size];
447                let size = if LARGE {
448                    u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
449                } else {
450                    // the buffer size is read from u16 but is stored as u32 after decoding for consistency
451                    u16::from_le_bytes([bytes[0], bytes[1]]) as u32
452                };
453                *offset += read_size;
454                size
455            })
456            .collect()
457    }
458
459    // Unserialize a miniblock into a collection of vectors
460    fn decode_miniblock_chunk(
461        &self,
462        buf: &LanceBuffer,
463        items_in_chunk: u64,
464    ) -> Result<DecodedMiniBlockChunk> {
465        let mut offset = 0;
466        let num_levels = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
467        offset += 2;
468
469        let rep_size = if self.rep_decompressor.is_some() {
470            let rep_size = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
471            offset += 2;
472            Some(rep_size)
473        } else {
474            None
475        };
476        let def_size = if self.def_decompressor.is_some() {
477            let def_size = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
478            offset += 2;
479            Some(def_size)
480        } else {
481            None
482        };
483
484        let buffer_sizes = if self.has_large_chunk {
485            Self::read_buffer_sizes::<true>(buf, &mut offset, self.num_buffers)
486        } else {
487            Self::read_buffer_sizes::<false>(buf, &mut offset, self.num_buffers)
488        };
489
490        offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
491
492        let rep = rep_size.map(|rep_size| {
493            let rep = buf.slice_with_length(offset, rep_size as usize);
494            offset += rep_size as usize;
495            offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
496            rep
497        });
498
499        let def = def_size.map(|def_size| {
500            let def = buf.slice_with_length(offset, def_size as usize);
501            offset += def_size as usize;
502            offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
503            def
504        });
505
506        let buffers = buffer_sizes
507            .into_iter()
508            .map(|buf_size| {
509                let buf = buf.slice_with_length(offset, buf_size as usize);
510                offset += buf_size as usize;
511                offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
512                buf
513            })
514            .collect::<Vec<_>>();
515
516        let values = self
517            .value_decompressor
518            .decompress(buffers, items_in_chunk)?;
519
520        let rep = rep
521            .map(|rep| {
522                Self::decode_levels(
523                    self.rep_decompressor.as_ref().unwrap().as_ref(),
524                    rep,
525                    num_levels,
526                )
527            })
528            .transpose()?;
529        let def = def
530            .map(|def| {
531                Self::decode_levels(
532                    self.def_decompressor.as_ref().unwrap().as_ref(),
533                    def,
534                    num_levels,
535                )
536            })
537            .transpose()?;
538
539        Ok(DecodedMiniBlockChunk { rep, def, values })
540    }
541}
542
543impl DecodePageTask for DecodeMiniBlockTask {
544    fn decode(self: Box<Self>) -> Result<DecodedPage> {
545        // First, we create output buffers for the rep and def and data
546        let mut repbuf: Option<LevelBuffer> = None;
547        let mut defbuf: Option<LevelBuffer> = None;
548
549        let max_rep = self.def_meaning.iter().filter(|l| l.is_list()).count() as u16;
550
551        // This is probably an over-estimate but it's quick and easy to calculate
552        let estimated_size_bytes = self
553            .instructions
554            .iter()
555            .map(|(_, chunk)| chunk.data.len())
556            .sum::<usize>()
557            * 2;
558        let mut data_builder =
559            DataBlockBuilder::with_capacity_estimate(estimated_size_bytes as u64);
560
561        // We need to keep track of the offset into repbuf/defbuf that we are building up
562        let mut level_offset = 0;
563
564        // Pre-compute caching needs for each chunk by checking if the next chunk is the same
565        let needs_caching: Vec<bool> = self
566            .instructions
567            .windows(2)
568            .map(|w| w[0].1.chunk_idx == w[1].1.chunk_idx)
569            .chain(std::iter::once(false)) // the last one never needs caching
570            .collect();
571
572        // Cache for storing decoded chunks when beneficial
573        let mut chunk_cache: Option<(usize, DecodedMiniBlockChunk)> = None;
574
575        // Now we iterate through each instruction and process it
576        for (idx, (instructions, chunk)) in self.instructions.iter().enumerate() {
577            let should_cache_this_chunk = needs_caching[idx];
578
579            let decoded_chunk = match &chunk_cache {
580                Some((cached_chunk_idx, cached_chunk)) if *cached_chunk_idx == chunk.chunk_idx => {
581                    // Clone only when we have a cache hit (much cheaper than decoding)
582                    cached_chunk.clone()
583                }
584                _ => {
585                    // Cache miss, need to decode
586                    let decoded = self.decode_miniblock_chunk(&chunk.data, chunk.items_in_chunk)?;
587
588                    // Only update cache if this chunk will benefit the next access
589                    if should_cache_this_chunk {
590                        chunk_cache = Some((chunk.chunk_idx, decoded.clone()));
591                    }
592                    decoded
593                }
594            };
595
596            let DecodedMiniBlockChunk { rep, def, values } = decoded_chunk;
597
598            // Our instructions tell us which rows we want to take from this chunk
599            let row_range_start =
600                instructions.rows_to_skip + instructions.chunk_instructions.rows_to_skip;
601            let row_range_end = row_range_start + instructions.rows_to_take;
602
603            // We use the rep info to map the row range to an item range / levels range
604            let (item_range, level_range) = Self::map_range(
605                row_range_start..row_range_end,
606                rep.as_ref(),
607                def.as_ref(),
608                max_rep,
609                self.max_visible_level,
610                chunk.items_in_chunk,
611                instructions.preamble_action,
612            );
613            if item_range.end - item_range.start > chunk.items_in_chunk {
614                return Err(lance_core::Error::internal(format!(
615                    "Item range {:?} is greater than chunk items in chunk {:?}",
616                    item_range, chunk.items_in_chunk
617                )));
618            }
619
620            // Now we append the data to the output buffers
621            Self::extend_levels(level_range.clone(), &mut repbuf, &rep, level_offset);
622            Self::extend_levels(level_range.clone(), &mut defbuf, &def, level_offset);
623            level_offset += (level_range.end - level_range.start) as usize;
624            data_builder.append(&values, item_range);
625        }
626
627        let mut data = data_builder.finish();
628
629        let unraveler =
630            RepDefUnraveler::new(repbuf, defbuf, self.def_meaning.clone(), data.num_values());
631
632        if let Some(dictionary) = &self.dictionary_data {
633            // Don't decode here, that happens later (if needed)
634            let DataBlock::FixedWidth(indices) = data else {
635                return Err(lance_core::Error::internal(format!(
636                    "Expected FixedWidth DataBlock for dictionary indices, got {:?}",
637                    data
638                )));
639            };
640            data = DataBlock::Dictionary(DictionaryDataBlock::from_parts(
641                indices,
642                dictionary.as_ref().clone(),
643            ));
644        }
645
646        Ok(DecodedPage {
647            data,
648            repdef: unraveler,
649        })
650    }
651}
652
653/// A chunk that has been loaded by the miniblock scheduler (but not
654/// yet decoded)
655#[derive(Debug)]
656struct LoadedChunk {
657    data: LanceBuffer,
658    items_in_chunk: u64,
659    byte_range: Range<u64>,
660    chunk_idx: usize,
661}
662
663impl Clone for LoadedChunk {
664    fn clone(&self) -> Self {
665        Self {
666            // Safe as we always create borrowed buffers here
667            data: self.data.clone(),
668            items_in_chunk: self.items_in_chunk,
669            byte_range: self.byte_range.clone(),
670            chunk_idx: self.chunk_idx,
671        }
672    }
673}
674
675/// Decodes mini-block formatted data.  See [`PrimitiveStructuralEncoder`] for more
676/// details on the different layouts.
677#[derive(Debug)]
678struct MiniBlockDecoder {
679    rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
680    def_decompressor: Option<Arc<dyn BlockDecompressor>>,
681    value_decompressor: Arc<dyn MiniBlockDecompressor>,
682    def_meaning: Arc<[DefinitionInterpretation]>,
683    loaded_chunks: VecDeque<LoadedChunk>,
684    instructions: VecDeque<ChunkInstructions>,
685    offset_in_current_chunk: u64,
686    num_rows: u64,
687    num_buffers: u64,
688    dictionary: Option<Arc<DataBlock>>,
689    has_large_chunk: bool,
690}
691
692/// See [`MiniBlockScheduler`] for more details on the scheduling and decoding
693/// process for miniblock encoded data.
694impl StructuralPageDecoder for MiniBlockDecoder {
695    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
696        let mut items_desired = num_rows;
697        let mut need_preamble = false;
698        let mut skip_in_chunk = self.offset_in_current_chunk;
699        let mut drain_instructions = Vec::new();
700        while items_desired > 0 || need_preamble {
701            let (instructions, consumed) = self
702                .instructions
703                .front()
704                .unwrap()
705                .drain_from_instruction(&mut items_desired, &mut need_preamble, &mut skip_in_chunk);
706
707            while self.loaded_chunks.front().unwrap().chunk_idx
708                != instructions.chunk_instructions.chunk_idx
709            {
710                self.loaded_chunks.pop_front();
711            }
712            drain_instructions.push((instructions, self.loaded_chunks.front().unwrap().clone()));
713            if consumed {
714                self.instructions.pop_front();
715            }
716        }
717        // We can throw away need_preamble here because it must be false.  If it were true it would mean
718        // we were still in the middle of loading rows.  We do need to latch skip_in_chunk though.
719        self.offset_in_current_chunk = skip_in_chunk;
720
721        let max_visible_level = self
722            .def_meaning
723            .iter()
724            .take_while(|l| !l.is_list())
725            .map(|l| l.num_def_levels())
726            .sum::<u16>();
727
728        Ok(Box::new(DecodeMiniBlockTask {
729            instructions: drain_instructions,
730            def_decompressor: self.def_decompressor.clone(),
731            rep_decompressor: self.rep_decompressor.clone(),
732            value_decompressor: self.value_decompressor.clone(),
733            dictionary_data: self.dictionary.clone(),
734            def_meaning: self.def_meaning.clone(),
735            num_buffers: self.num_buffers,
736            max_visible_level,
737            has_large_chunk: self.has_large_chunk,
738        }))
739    }
740
741    fn num_rows(&self) -> u64 {
742        self.num_rows
743    }
744}
745
746#[derive(Debug)]
747struct CachedComplexAllNullState {
748    rep: Option<ScalarBuffer<u16>>,
749    def: Option<ScalarBuffer<u16>>,
750}
751
752impl DeepSizeOf for CachedComplexAllNullState {
753    fn deep_size_of_children(&self, _ctx: &mut Context) -> usize {
754        self.rep.as_ref().map(|buf| buf.len() * 2).unwrap_or(0)
755            + self.def.as_ref().map(|buf| buf.len() * 2).unwrap_or(0)
756    }
757}
758
759impl CachedPageData for CachedComplexAllNullState {
760    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
761        self
762    }
763}
764
765/// A scheduler for all-null data that has repetition and definition levels
766///
767/// We still need to do some I/O in this case because we need to figure out what kind of null we
768/// are dealing with (null list, null struct, what level null struct, etc.)
769///
770/// TODO: Right now we just load the entire rep/def at initialization time and cache it.  This is a touch
771/// RAM aggressive and maybe we want something more lazy in the future.  On the other hand, it's simple
772/// and fast so...maybe not :)
773#[derive(Debug)]
774pub struct ComplexAllNullScheduler {
775    // Set from protobuf
776    buffer_offsets_and_sizes: Arc<[(u64, u64)]>,
777    def_meaning: Arc<[DefinitionInterpretation]>,
778    repdef: Option<Arc<CachedComplexAllNullState>>,
779    max_rep: u16,
780    max_visible_level: u16,
781    rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
782    def_decompressor: Option<Arc<dyn BlockDecompressor>>,
783    num_rep_values: u64,
784    num_def_values: u64,
785}
786
787impl ComplexAllNullScheduler {
788    pub fn new(
789        buffer_offsets_and_sizes: Arc<[(u64, u64)]>,
790        def_meaning: Arc<[DefinitionInterpretation]>,
791        rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
792        def_decompressor: Option<Arc<dyn BlockDecompressor>>,
793        num_rep_values: u64,
794        num_def_values: u64,
795    ) -> Self {
796        let max_rep = def_meaning.iter().filter(|l| l.is_list()).count() as u16;
797        let max_visible_level = def_meaning
798            .iter()
799            .take_while(|l| !l.is_list())
800            .map(|l| l.num_def_levels())
801            .sum::<u16>();
802        Self {
803            buffer_offsets_and_sizes,
804            def_meaning,
805            repdef: None,
806            max_rep,
807            max_visible_level,
808            rep_decompressor,
809            def_decompressor,
810            num_rep_values,
811            num_def_values,
812        }
813    }
814}
815
816impl StructuralPageScheduler for ComplexAllNullScheduler {
817    fn initialize<'a>(
818        &'a mut self,
819        io: &Arc<dyn EncodingsIo>,
820    ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
821        // Fully load the rep & def buffers, as needed
822        let (rep_pos, rep_size) = self.buffer_offsets_and_sizes[0];
823        let (def_pos, def_size) = self.buffer_offsets_and_sizes[1];
824        let has_rep = rep_size > 0;
825        let has_def = def_size > 0;
826
827        let mut reads = Vec::with_capacity(2);
828        if has_rep {
829            reads.push(rep_pos..rep_pos + rep_size);
830        }
831        if has_def {
832            reads.push(def_pos..def_pos + def_size);
833        }
834
835        let data = io.submit_request(reads, 0);
836        let rep_decompressor = self.rep_decompressor.clone();
837        let def_decompressor = self.def_decompressor.clone();
838        let num_rep_values = self.num_rep_values;
839        let num_def_values = self.num_def_values;
840
841        async move {
842            let data = data.await?;
843            let mut data_iter = data.into_iter();
844
845            let decompress_levels = |compressed_bytes: Bytes,
846                                     decompressor: &Arc<dyn BlockDecompressor>,
847                                     num_values: u64,
848                                     level_type: &str|
849             -> Result<ScalarBuffer<u16>> {
850                let compressed_buffer = LanceBuffer::from_bytes(compressed_bytes, 1);
851                let decompressed = decompressor.decompress(compressed_buffer, num_values)?;
852                match decompressed {
853                    DataBlock::FixedWidth(block) => {
854                        if block.num_values != num_values {
855                            return Err(Error::invalid_input_source(format!(
856                                "Unexpected {} level count after decompression: expected {}, got {}",
857                                level_type, num_values, block.num_values
858                            )
859                            .into()));
860                        }
861                        if block.bits_per_value != 16 {
862                            return Err(Error::invalid_input_source(format!(
863                                "Unexpected {} level bit width after decompression: expected 16, got {}",
864                                level_type, block.bits_per_value
865                            )
866                            .into()));
867                        }
868                        Ok(block.data.borrow_to_typed_slice::<u16>())
869                    }
870                    _ => Err(Error::invalid_input_source(format!(
871                        "Expected fixed-width data block for {} levels",
872                        level_type
873                    )
874                    .into())),
875                }
876            };
877
878            let rep = if has_rep {
879                let rep = data_iter.next().unwrap();
880                if let Some(rep_decompressor) = rep_decompressor.as_ref() {
881                    Some(decompress_levels(
882                        rep,
883                        rep_decompressor,
884                        num_rep_values,
885                        "repetition",
886                    )?)
887                } else {
888                    let rep = LanceBuffer::from_bytes(rep, 2);
889                    let rep = rep.borrow_to_typed_slice::<u16>();
890                    Some(rep)
891                }
892            } else {
893                None
894            };
895
896            let def = if has_def {
897                let def = data_iter.next().unwrap();
898                if let Some(def_decompressor) = def_decompressor.as_ref() {
899                    Some(decompress_levels(
900                        def,
901                        def_decompressor,
902                        num_def_values,
903                        "definition",
904                    )?)
905                } else {
906                    let def = LanceBuffer::from_bytes(def, 2);
907                    let def = def.borrow_to_typed_slice::<u16>();
908                    Some(def)
909                }
910            } else {
911                None
912            };
913
914            let repdef = Arc::new(CachedComplexAllNullState { rep, def });
915
916            self.repdef = Some(repdef.clone());
917
918            Ok(repdef as Arc<dyn CachedPageData>)
919        }
920        .boxed()
921    }
922
923    fn load(&mut self, data: &Arc<dyn CachedPageData>) {
924        self.repdef = Some(
925            data.clone()
926                .as_arc_any()
927                .downcast::<CachedComplexAllNullState>()
928                .unwrap(),
929        );
930    }
931
932    fn schedule_ranges(
933        &self,
934        ranges: &[Range<u64>],
935        _io: &Arc<dyn EncodingsIo>,
936    ) -> Result<Vec<PageLoadTask>> {
937        let ranges = VecDeque::from_iter(ranges.iter().cloned());
938        let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
939        let decoder = Box::new(ComplexAllNullPageDecoder {
940            ranges,
941            rep: self.repdef.as_ref().unwrap().rep.clone(),
942            def: self.repdef.as_ref().unwrap().def.clone(),
943            num_rows,
944            def_meaning: self.def_meaning.clone(),
945            max_rep: self.max_rep,
946            max_visible_level: self.max_visible_level,
947            cursor_row: 0,
948            cursor_level: 0,
949        }) as Box<dyn StructuralPageDecoder>;
950        let page_load_task = PageLoadTask {
951            decoder_fut: std::future::ready(Ok(decoder)).boxed(),
952            num_rows,
953        };
954        Ok(vec![page_load_task])
955    }
956}
957
958#[derive(Debug)]
959pub struct ComplexAllNullPageDecoder {
960    ranges: VecDeque<Range<u64>>,
961    rep: Option<ScalarBuffer<u16>>,
962    def: Option<ScalarBuffer<u16>>,
963    num_rows: u64,
964    def_meaning: Arc<[DefinitionInterpretation]>,
965    max_rep: u16,
966    max_visible_level: u16,
967    cursor_row: u64,
968    cursor_level: usize,
969}
970
971impl ComplexAllNullPageDecoder {
972    fn drain_ranges(&mut self, num_rows: u64) -> Vec<Range<u64>> {
973        let mut rows_desired = num_rows;
974        let mut ranges = Vec::with_capacity(self.ranges.len());
975        while rows_desired > 0 {
976            let front = self.ranges.front_mut().unwrap();
977            let avail = front.end - front.start;
978            if avail > rows_desired {
979                ranges.push(front.start..front.start + rows_desired);
980                front.start += rows_desired;
981                rows_desired = 0;
982            } else {
983                ranges.push(self.ranges.pop_front().unwrap());
984                rows_desired -= avail;
985            }
986        }
987        ranges
988    }
989
990    fn take_row(&mut self) -> Result<(Range<usize>, u64)> {
991        let start = self.cursor_level;
992        let end = if let Some(rep) = &self.rep {
993            if start >= rep.len() {
994                return Err(Error::internal(
995                    "Invalid complex all-null layout: repetition buffer too short",
996                ));
997            }
998            if rep[start] != self.max_rep {
999                return Err(Error::internal(
1000                    "Invalid complex all-null layout: row did not start at max repetition level",
1001                ));
1002            }
1003            let mut end = start + 1;
1004            while end < rep.len() && rep[end] != self.max_rep {
1005                end += 1;
1006            }
1007            end
1008        } else {
1009            start + 1
1010        };
1011
1012        let visible = if let Some(def) = &self.def {
1013            if end > def.len() {
1014                return Err(Error::internal(
1015                    "Invalid complex all-null layout: definition buffer too short",
1016                ));
1017            }
1018            def[start..end]
1019                .iter()
1020                .filter(|d| **d <= self.max_visible_level)
1021                .count() as u64
1022        } else {
1023            (end - start) as u64
1024        };
1025
1026        self.cursor_level = end;
1027        self.cursor_row += 1;
1028        Ok((start..end, visible))
1029    }
1030
1031    fn skip_to_row(&mut self, target_row: u64) -> Result<()> {
1032        while self.cursor_row < target_row {
1033            self.take_row()?;
1034        }
1035        Ok(())
1036    }
1037}
1038
1039impl StructuralPageDecoder for ComplexAllNullPageDecoder {
1040    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
1041        let drained_ranges = self.drain_ranges(num_rows);
1042        let mut level_slices: Vec<Range<usize>> = Vec::new();
1043        let mut visible_items_total = 0;
1044
1045        for range in drained_ranges {
1046            self.skip_to_row(range.start)?;
1047            for _ in range.start..range.end {
1048                let (level_range, visible) = self.take_row()?;
1049                visible_items_total += visible;
1050                if let Some(last) = level_slices.last_mut()
1051                    && last.end == level_range.start
1052                {
1053                    last.end = level_range.end;
1054                    continue;
1055                }
1056                level_slices.push(level_range);
1057            }
1058        }
1059
1060        Ok(Box::new(DecodeComplexAllNullTask {
1061            level_slices,
1062            visible_items_total,
1063            rep: self.rep.clone(),
1064            def: self.def.clone(),
1065            def_meaning: self.def_meaning.clone(),
1066            max_visible_level: self.max_visible_level,
1067        }))
1068    }
1069
1070    fn num_rows(&self) -> u64 {
1071        self.num_rows
1072    }
1073}
1074
1075/// We use `level_slices` to slice into `rep` and `def` and create rep/def buffers
1076/// for the null data.
1077#[derive(Debug)]
1078pub struct DecodeComplexAllNullTask {
1079    level_slices: Vec<Range<usize>>,
1080    visible_items_total: u64,
1081    rep: Option<ScalarBuffer<u16>>,
1082    def: Option<ScalarBuffer<u16>>,
1083    def_meaning: Arc<[DefinitionInterpretation]>,
1084    max_visible_level: u16,
1085}
1086
1087impl DecodeComplexAllNullTask {
1088    fn decode_level(&self, levels: &Option<ScalarBuffer<u16>>) -> Option<Vec<u16>> {
1089        levels.as_ref().map(|levels| {
1090            let num_levels = self
1091                .level_slices
1092                .iter()
1093                .map(|range| range.end - range.start)
1094                .sum();
1095            let mut referenced_levels = Vec::with_capacity(num_levels);
1096            for range in &self.level_slices {
1097                referenced_levels.extend(levels[range.start..range.end].iter().copied());
1098            }
1099            referenced_levels
1100        })
1101    }
1102}
1103
1104impl DecodePageTask for DecodeComplexAllNullTask {
1105    fn decode(self: Box<Self>) -> Result<DecodedPage> {
1106        let rep = self.decode_level(&self.rep);
1107        let def = self.decode_level(&self.def);
1108
1109        // If there are definition levels there may be empty / null lists which are not visible
1110        // in the items array.  We need to account for that here to figure out how many values
1111        // should be in the items array.
1112        let num_values = if let Some(def) = &def {
1113            def.iter().filter(|&d| *d <= self.max_visible_level).count() as u64
1114        } else {
1115            self.visible_items_total
1116        };
1117
1118        let data = DataBlock::AllNull(AllNullDataBlock { num_values });
1119        let unraveler = RepDefUnraveler::new(rep, def, self.def_meaning, num_values);
1120        Ok(DecodedPage {
1121            data,
1122            repdef: unraveler,
1123        })
1124    }
1125}
1126
1127/// A scheduler for simple all-null data
1128///
1129/// "simple" all-null data is data that is all null and only has a single level of definition and
1130/// no repetition.  We don't need to read any data at all in this case.
1131#[derive(Debug, Default)]
1132pub struct SimpleAllNullScheduler {}
1133
1134impl StructuralPageScheduler for SimpleAllNullScheduler {
1135    fn initialize<'a>(
1136        &'a mut self,
1137        _io: &Arc<dyn EncodingsIo>,
1138    ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
1139        std::future::ready(Ok(Arc::new(NoCachedPageData) as Arc<dyn CachedPageData>)).boxed()
1140    }
1141
1142    fn load(&mut self, _cache: &Arc<dyn CachedPageData>) {}
1143
1144    fn schedule_ranges(
1145        &self,
1146        ranges: &[Range<u64>],
1147        _io: &Arc<dyn EncodingsIo>,
1148    ) -> Result<Vec<PageLoadTask>> {
1149        let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
1150        let decoder =
1151            Box::new(SimpleAllNullPageDecoder { num_rows }) as Box<dyn StructuralPageDecoder>;
1152        let page_load_task = PageLoadTask {
1153            decoder_fut: std::future::ready(Ok(decoder)).boxed(),
1154            num_rows,
1155        };
1156        Ok(vec![page_load_task])
1157    }
1158}
1159
1160/// A page decode task for all-null data without any
1161/// repetition and only a single level of definition
1162#[derive(Debug)]
1163struct SimpleAllNullDecodePageTask {
1164    num_values: u64,
1165}
1166impl DecodePageTask for SimpleAllNullDecodePageTask {
1167    fn decode(self: Box<Self>) -> Result<DecodedPage> {
1168        let unraveler = RepDefUnraveler::new(
1169            None,
1170            Some(vec![1; self.num_values as usize]),
1171            Arc::new([DefinitionInterpretation::NullableItem]),
1172            self.num_values,
1173        );
1174        Ok(DecodedPage {
1175            data: DataBlock::AllNull(AllNullDataBlock {
1176                num_values: self.num_values,
1177            }),
1178            repdef: unraveler,
1179        })
1180    }
1181}
1182
1183#[derive(Debug)]
1184pub struct SimpleAllNullPageDecoder {
1185    num_rows: u64,
1186}
1187
1188impl StructuralPageDecoder for SimpleAllNullPageDecoder {
1189    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
1190        Ok(Box::new(SimpleAllNullDecodePageTask {
1191            num_values: num_rows,
1192        }))
1193    }
1194
1195    fn num_rows(&self) -> u64 {
1196        self.num_rows
1197    }
1198}
1199
1200#[derive(Debug, Clone)]
1201struct MiniBlockSchedulerDictionary {
1202    // These come from the protobuf
1203    dictionary_decompressor: Arc<dyn BlockDecompressor>,
1204    dictionary_buf_position_and_size: (u64, u64),
1205    dictionary_data_alignment: u64,
1206    num_dictionary_items: u64,
1207}
1208
1209/// Individual block metadata within a MiniBlock repetition index.
1210#[derive(Debug)]
1211struct MiniBlockRepIndexBlock {
1212    // The index of the first row that starts after the beginning of this block.  If the block
1213    // has a preamble this will be the row after the preamble.  If the block is entirely preamble
1214    // then this will be a row that starts in some future block.
1215    first_row: u64,
1216    // The number of rows in the block, including the trailer but not the preamble.
1217    // Can be 0 if the block is entirely preamble
1218    starts_including_trailer: u64,
1219    // Whether the block has a preamble
1220    has_preamble: bool,
1221    // Whether the block has a trailer
1222    has_trailer: bool,
1223}
1224
1225impl DeepSizeOf for MiniBlockRepIndexBlock {
1226    fn deep_size_of_children(&self, _context: &mut Context) -> usize {
1227        0
1228    }
1229}
1230
1231/// Repetition index for MiniBlock encoding.
1232///
1233/// Stores block-level offset information to enable efficient random
1234/// access to nested data structures within mini-blocks.
1235#[derive(Debug)]
1236struct MiniBlockRepIndex {
1237    blocks: Vec<MiniBlockRepIndexBlock>,
1238}
1239
1240impl DeepSizeOf for MiniBlockRepIndex {
1241    fn deep_size_of_children(&self, context: &mut Context) -> usize {
1242        self.blocks.deep_size_of_children(context)
1243    }
1244}
1245
1246impl MiniBlockRepIndex {
1247    /// Decode repetition index from chunk metadata using default values.
1248    ///
1249    /// This creates a repetition index where each chunk has no partial values
1250    /// and no trailers, suitable for simple sequential data layouts.
1251    pub fn default_from_chunks(chunks: &[ChunkMeta]) -> Self {
1252        let mut blocks = Vec::with_capacity(chunks.len());
1253        let mut offset: u64 = 0;
1254
1255        for c in chunks {
1256            blocks.push(MiniBlockRepIndexBlock {
1257                first_row: offset,
1258                starts_including_trailer: c.num_values,
1259                has_preamble: false,
1260                has_trailer: false,
1261            });
1262
1263            offset += c.num_values;
1264        }
1265
1266        Self { blocks }
1267    }
1268
1269    /// Decode repetition index from raw bytes in little-endian format.
1270    ///
1271    /// The bytes should contain u64 values arranged in groups of `stride` elements,
1272    /// where the first two values of each group represent ends_count and partial_count.
1273    /// Returns an empty index if no bytes are provided.
1274    pub fn decode_from_bytes(rep_bytes: &[u8], stride: usize) -> Self {
1275        // Convert bytes to u64 slice, handling alignment automatically
1276        let buffer = crate::buffer::LanceBuffer::from(rep_bytes.to_vec());
1277        let u64_slice = buffer.borrow_to_typed_slice::<u64>();
1278        let n = u64_slice.len() / stride;
1279
1280        let mut blocks = Vec::with_capacity(n);
1281        let mut chunk_has_preamble = false;
1282        let mut offset: u64 = 0;
1283
1284        // Extract first two values from each block: ends_count and partial_count
1285        for i in 0..n {
1286            let base_idx = i * stride;
1287            let ends = u64_slice[base_idx];
1288            let partial = u64_slice[base_idx + 1];
1289
1290            let has_trailer = partial > 0;
1291            // Convert branches to arithmetic for better compiler optimization
1292            let starts_including_trailer =
1293                ends + (has_trailer as u64) - (chunk_has_preamble as u64);
1294
1295            blocks.push(MiniBlockRepIndexBlock {
1296                first_row: offset,
1297                starts_including_trailer,
1298                has_preamble: chunk_has_preamble,
1299                has_trailer,
1300            });
1301
1302            chunk_has_preamble = has_trailer;
1303            offset += starts_including_trailer;
1304        }
1305
1306        Self { blocks }
1307    }
1308}
1309
1310/// State that is loaded once and cached for future lookups
1311#[derive(Debug)]
1312struct MiniBlockCacheableState {
1313    /// Metadata that describes each chunk in the page
1314    chunk_meta: Vec<ChunkMeta>,
1315    /// The decoded repetition index
1316    rep_index: MiniBlockRepIndex,
1317    /// The dictionary for the page, if any
1318    dictionary: Option<Arc<DataBlock>>,
1319}
1320
1321impl DeepSizeOf for MiniBlockCacheableState {
1322    fn deep_size_of_children(&self, context: &mut Context) -> usize {
1323        self.rep_index.deep_size_of_children(context)
1324            + self
1325                .dictionary
1326                .as_ref()
1327                .map(|dict| dict.data_size() as usize)
1328                .unwrap_or(0)
1329    }
1330}
1331
1332impl CachedPageData for MiniBlockCacheableState {
1333    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
1334        self
1335    }
1336}
1337
1338/// A scheduler for a page that has been encoded with the mini-block layout
1339///
1340/// Scheduling mini-block encoded data is simple in concept and somewhat complex
1341/// in practice.
1342///
1343/// First, during initialization, we load the chunk metadata, the repetition index,
1344/// and the dictionary (these last two may not be present)
1345///
1346/// Then, during scheduling, we use the user's requested row ranges and the repetition
1347/// index to determine which chunks we need and which rows we need from those chunks.
1348///
1349/// For example, if the repetition index is: [50, 3], [50, 0], [10, 0] and the range
1350/// from the user is 40..60 then we need to:
1351///
1352///  - Read the first chunk and skip the first 40 rows, then read 10 full rows, and
1353///    then read 3 items for the 11th row of our range.
1354///  - Read the second chunk and read the remaining items in our 11th row and then read
1355///    the remaining 9 full rows.
1356///
1357/// Then, if we are going to decode that in batches of 5, we need to make decode tasks.
1358/// The first two decode tasks will just need the first chunk.  The third decode task will
1359/// need the first chunk (for the trailer which has the 11th row in our range) and the second
1360/// chunk.  The final decode task will just need the second chunk.
1361///
1362/// The above prose descriptions are what are represented by `ChunkInstructions` and
1363/// `ChunkDrainInstructions`.
1364#[derive(Debug)]
1365pub struct MiniBlockScheduler {
1366    // These come from the protobuf
1367    buffer_offsets_and_sizes: Vec<(u64, u64)>,
1368    priority: u64,
1369    items_in_page: u64,
1370    repetition_index_depth: u16,
1371    num_buffers: u64,
1372    rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
1373    def_decompressor: Option<Arc<dyn BlockDecompressor>>,
1374    value_decompressor: Arc<dyn MiniBlockDecompressor>,
1375    def_meaning: Arc<[DefinitionInterpretation]>,
1376    dictionary: Option<MiniBlockSchedulerDictionary>,
1377    // This is set after initialization
1378    page_meta: Option<Arc<MiniBlockCacheableState>>,
1379    has_large_chunk: bool,
1380}
1381
1382impl MiniBlockScheduler {
1383    fn try_new(
1384        buffer_offsets_and_sizes: &[(u64, u64)],
1385        priority: u64,
1386        items_in_page: u64,
1387        layout: &pb21::MiniBlockLayout,
1388        decompressors: &dyn DecompressionStrategy,
1389    ) -> Result<Self> {
1390        let rep_decompressor = layout
1391            .rep_compression
1392            .as_ref()
1393            .map(|rep_compression| {
1394                decompressors
1395                    .create_block_decompressor(rep_compression)
1396                    .map(Arc::from)
1397            })
1398            .transpose()?;
1399        let def_decompressor = layout
1400            .def_compression
1401            .as_ref()
1402            .map(|def_compression| {
1403                decompressors
1404                    .create_block_decompressor(def_compression)
1405                    .map(Arc::from)
1406            })
1407            .transpose()?;
1408        let def_meaning = layout
1409            .layers
1410            .iter()
1411            .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
1412            .collect::<Vec<_>>();
1413        let value_decompressor = decompressors.create_miniblock_decompressor(
1414            layout.value_compression.as_ref().unwrap(),
1415            decompressors,
1416        )?;
1417
1418        let dictionary = if let Some(dictionary_encoding) = layout.dictionary.as_ref() {
1419            let num_dictionary_items = layout.num_dictionary_items;
1420            let dictionary_decompressor = decompressors
1421                .create_block_decompressor(dictionary_encoding)?
1422                .into();
1423            let dictionary_data_alignment = match dictionary_encoding.compression.as_ref().unwrap()
1424            {
1425                Compression::Variable(_) => 4,
1426                Compression::Flat(_) => 16,
1427                Compression::General(_) => 1,
1428                Compression::InlineBitpacking(_) | Compression::OutOfLineBitpacking(_) => {
1429                    crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT
1430                }
1431                _ => {
1432                    return Err(Error::invalid_input_source(
1433                        format!(
1434                            "Unsupported mini-block dictionary encoding: {:?}",
1435                            dictionary_encoding.compression.as_ref().unwrap()
1436                        )
1437                        .into(),
1438                    ));
1439                }
1440            };
1441            Some(MiniBlockSchedulerDictionary {
1442                dictionary_decompressor,
1443                dictionary_buf_position_and_size: buffer_offsets_and_sizes[2],
1444                dictionary_data_alignment,
1445                num_dictionary_items,
1446            })
1447        } else {
1448            None
1449        };
1450
1451        Ok(Self {
1452            buffer_offsets_and_sizes: buffer_offsets_and_sizes.to_vec(),
1453            rep_decompressor,
1454            def_decompressor,
1455            value_decompressor: value_decompressor.into(),
1456            repetition_index_depth: layout.repetition_index_depth as u16,
1457            num_buffers: layout.num_buffers,
1458            priority,
1459            items_in_page,
1460            dictionary,
1461            def_meaning: def_meaning.into(),
1462            page_meta: None,
1463            has_large_chunk: layout.has_large_chunk,
1464        })
1465    }
1466
1467    fn lookup_chunks(&self, chunk_indices: &[usize]) -> Vec<LoadedChunk> {
1468        let page_meta = self.page_meta.as_ref().unwrap();
1469        chunk_indices
1470            .iter()
1471            .map(|&chunk_idx| {
1472                let chunk_meta = &page_meta.chunk_meta[chunk_idx];
1473                let bytes_start = chunk_meta.offset_bytes;
1474                let bytes_end = bytes_start + chunk_meta.chunk_size_bytes;
1475                LoadedChunk {
1476                    byte_range: bytes_start..bytes_end,
1477                    items_in_chunk: chunk_meta.num_values,
1478                    chunk_idx,
1479                    data: LanceBuffer::empty(),
1480                }
1481            })
1482            .collect()
1483    }
1484}
1485
1486#[derive(Debug, PartialEq, Eq, Clone, Copy)]
1487enum PreambleAction {
1488    Take,
1489    Skip,
1490    Absent,
1491}
1492
1493// When we schedule a chunk we use the repetition index (or, if none exists, just the # of items
1494// in each chunk) to map a user requested range into a set of ChunkInstruction objects which tell
1495// us how exactly to read from the chunk.
1496//
1497// Examples:
1498//
1499// | Chunk 0     | Chunk 1   | Chunk 2   | Chunk 3 |
1500// | xxxxyyyyzzz | zzzzzzzzz | zzzzzzzzz | aaabbcc |
1501//
1502// Full read (0..6)
1503//
1504// Chunk 0: (several rows, ends with trailer)
1505//   preamble: absent
1506//   rows_to_skip: 0
1507//   rows_to_take: 3 (x, y, z)
1508//   take_trailer: true
1509//
1510// Chunk 1: (all preamble, ends with trailer)
1511//   preamble: take
1512//   rows_to_skip: 0
1513//   rows_to_take: 0
1514//   take_trailer: true
1515//
1516// Chunk 2: (all preamble, no trailer)
1517//   preamble: take
1518//   rows_to_skip: 0
1519//   rows_to_take: 0
1520//   take_trailer: false
1521//
1522// Chunk 3: (several rows, no trailer or preamble)
1523//   preamble: absent
1524//   rows_to_skip: 0
1525//   rows_to_take: 3 (a, b, c)
1526//   take_trailer: false
1527#[derive(Clone, Debug, PartialEq, Eq)]
1528struct ChunkInstructions {
1529    // The index of the chunk to read
1530    chunk_idx: usize,
1531    // A "preamble" is when a chunk begins with a continuation of the previous chunk's list.  If there
1532    // is no repetition index there is never a preamble.
1533    //
1534    // It's possible for a chunk to be entirely premable.  For example, if there is a really large list
1535    // that spans several chunks.
1536    preamble: PreambleAction,
1537    // How many complete rows (not including the preamble or trailer) to skip
1538    //
1539    // If this is non-zero then premable must not be Take
1540    rows_to_skip: u64,
1541    // How many rows to take.  If a row splits across chunks then we will count the row in the first
1542    // chunk that contains the row.
1543    rows_to_take: u64,
1544    // A "trailer" is when a chunk ends with a partial list.  If there is no repetition index there is
1545    // never a trailer.
1546    //
1547    // A chunk that is all preamble may or may not have a trailer.
1548    //
1549    // If this is true then we want to include the trailer
1550    take_trailer: bool,
1551}
1552
1553// First, we schedule a bunch of [`ChunkInstructions`] based on the users ranges.  Then we
1554// start decoding them, based on a batch size, which might not align with what we scheduled.
1555//
1556// This results in `ChunkDrainInstructions` which targets a contiguous slice of a `ChunkInstructions`
1557//
1558// So if `ChunkInstructions` is "skip preamble, skip 10, take 50, take trailer" and we are decoding in
1559// batches of size 10 we might have a `ChunkDrainInstructions` that targets that chunk and has its own
1560// skip of 17 and take of 10.  This would mean we decode the chunk, skip the preamble and 27 rows, and
1561// then take 10 rows.
1562//
1563// One very confusing bit is that `rows_to_take` includes the trailer.  So if we have two chunks:
1564//  -no preamble, skip 5, take 10, take trailer
1565//  -take preamble, skip 0, take 50, no trailer
1566//
1567// and we are draining 20 rows then the drain instructions for the first batch will be:
1568//  - no preamble, skip 0 (from chunk 0), take 11 (from chunk 0)
1569//  - take preamble (from chunk 1), skip 0 (from chunk 1), take 9 (from chunk 1)
1570#[derive(Debug, PartialEq, Eq)]
1571struct ChunkDrainInstructions {
1572    chunk_instructions: ChunkInstructions,
1573    rows_to_skip: u64,
1574    rows_to_take: u64,
1575    preamble_action: PreambleAction,
1576}
1577
1578impl ChunkInstructions {
1579    // Given a repetition index and a set of user ranges we need to figure out how to read from the chunks
1580    //
1581    // We assume that `user_ranges` are in sorted order and non-overlapping
1582    //
1583    // The output will be a set of `ChunkInstructions` which tell us how to read from the chunks
1584    fn schedule_instructions(
1585        rep_index: &MiniBlockRepIndex,
1586        user_ranges: &[Range<u64>],
1587    ) -> Vec<Self> {
1588        // This is an in-exact capacity guess but pretty good.  The actual capacity can be
1589        // smaller if instructions are merged.  It can be larger if there are multiple instructions
1590        // per row which can happen with lists.
1591        let mut chunk_instructions = Vec::with_capacity(user_ranges.len());
1592
1593        for user_range in user_ranges {
1594            let mut rows_needed = user_range.end - user_range.start;
1595            let mut need_preamble = false;
1596
1597            // Need to find the first chunk with a first row >= user_range.start.  If there are
1598            // multiple chunks with the same first row we need to take the first one.
1599            let mut block_index = match rep_index
1600                .blocks
1601                .binary_search_by_key(&user_range.start, |block| block.first_row)
1602            {
1603                Ok(idx) => {
1604                    // Slightly tricky case, we may need to walk backwards a bit to make sure we
1605                    // are grabbing first eligible chunk
1606                    let mut idx = idx;
1607                    while idx > 0 && rep_index.blocks[idx - 1].first_row == user_range.start {
1608                        idx -= 1;
1609                    }
1610                    idx
1611                }
1612                // Easy case.  idx is greater, and idx - 1 is smaller, so idx - 1 contains the start
1613                Err(idx) => idx - 1,
1614            };
1615
1616            let mut to_skip = user_range.start - rep_index.blocks[block_index].first_row;
1617
1618            while rows_needed > 0 || need_preamble {
1619                // Check if we've gone past the last block (should not happen)
1620                if block_index >= rep_index.blocks.len() {
1621                    log::warn!(
1622                        "schedule_instructions inconsistency: block_index >= rep_index.blocks.len(), exiting early"
1623                    );
1624                    break;
1625                }
1626
1627                let chunk = &rep_index.blocks[block_index];
1628                let rows_avail = chunk.starts_including_trailer.saturating_sub(to_skip);
1629
1630                // Handle blocks that are entirely preamble (rows_avail = 0)
1631                // These blocks have no rows to take but may have a preamble we need
1632                // We only look for preamble if to_skip == 0 (we're not skipping rows)
1633                if rows_avail == 0 && to_skip == 0 {
1634                    // Only process if this chunk has a preamble we need
1635                    if chunk.has_preamble && need_preamble {
1636                        chunk_instructions.push(Self {
1637                            chunk_idx: block_index,
1638                            preamble: PreambleAction::Take,
1639                            rows_to_skip: 0,
1640                            rows_to_take: 0,
1641                            // We still need to look at has_trailer to distinguish between "all preamble
1642                            // and row ends at end of chunk" and "all preamble and row bleeds into next
1643                            // chunk".  Both cases will have 0 rows available.
1644                            take_trailer: chunk.has_trailer,
1645                        });
1646                        // Only set need_preamble = false if the chunk has at least one row,
1647                        // Or we are reaching the last block,
1648                        // Otherwise, the chunk is entirely preamble and we need the next chunk's preamble too
1649                        if chunk.starts_including_trailer > 0
1650                            || block_index == rep_index.blocks.len() - 1
1651                        {
1652                            need_preamble = false;
1653                        }
1654                    }
1655                    // Move to next block
1656                    block_index += 1;
1657                    continue;
1658                }
1659
1660                // Edge case: if rows_avail == 0 but to_skip > 0
1661                // This theoretically shouldn't happen (binary search should avoid it)
1662                // but handle it for safety
1663                if rows_avail == 0 && to_skip > 0 {
1664                    // This block doesn't have enough rows to skip, move to next block
1665                    // Adjust to_skip by the number of rows in this block
1666                    to_skip -= chunk.starts_including_trailer;
1667                    block_index += 1;
1668                    continue;
1669                }
1670
1671                let rows_to_take = rows_avail.min(rows_needed);
1672                rows_needed -= rows_to_take;
1673
1674                let mut take_trailer = false;
1675                let preamble = if chunk.has_preamble {
1676                    if need_preamble {
1677                        PreambleAction::Take
1678                    } else {
1679                        PreambleAction::Skip
1680                    }
1681                } else {
1682                    PreambleAction::Absent
1683                };
1684
1685                // Are we taking the trailer?  If so, make sure we mark that we need the preamble
1686                if rows_to_take == rows_avail && chunk.has_trailer {
1687                    take_trailer = true;
1688                    need_preamble = true;
1689                } else {
1690                    need_preamble = false;
1691                };
1692
1693                chunk_instructions.push(Self {
1694                    preamble,
1695                    chunk_idx: block_index,
1696                    rows_to_skip: to_skip,
1697                    rows_to_take,
1698                    take_trailer,
1699                });
1700
1701                to_skip = 0;
1702                block_index += 1;
1703            }
1704        }
1705
1706        // If there were multiple ranges we may have multiple instructions for a single chunk.  Merge them now if they
1707        // are _adjacent_ (i.e. don't merge "take first row of chunk 0" and "take third row of chunk 0" into "take 2
1708        // rows of chunk 0 starting at 0")
1709        if user_ranges.len() > 1 {
1710            // TODO: Could probably optimize this allocation away
1711            let mut merged_instructions = Vec::with_capacity(chunk_instructions.len());
1712            let mut instructions_iter = chunk_instructions.into_iter();
1713            merged_instructions.push(instructions_iter.next().unwrap());
1714            for instruction in instructions_iter {
1715                let last = merged_instructions.last_mut().unwrap();
1716                if last.chunk_idx == instruction.chunk_idx
1717                    && last.rows_to_take + last.rows_to_skip == instruction.rows_to_skip
1718                {
1719                    last.rows_to_take += instruction.rows_to_take;
1720                    last.take_trailer |= instruction.take_trailer;
1721                } else {
1722                    merged_instructions.push(instruction);
1723                }
1724            }
1725            merged_instructions
1726        } else {
1727            chunk_instructions
1728        }
1729    }
1730
1731    fn drain_from_instruction(
1732        &self,
1733        rows_desired: &mut u64,
1734        need_preamble: &mut bool,
1735        skip_in_chunk: &mut u64,
1736    ) -> (ChunkDrainInstructions, bool) {
1737        // If we need the premable then we shouldn't be skipping anything
1738        debug_assert!(!*need_preamble || *skip_in_chunk == 0);
1739        let rows_avail = self.rows_to_take - *skip_in_chunk;
1740        let has_preamble = self.preamble != PreambleAction::Absent;
1741        let preamble_action = match (*need_preamble, has_preamble) {
1742            (true, true) => PreambleAction::Take,
1743            (true, false) => panic!("Need preamble but there isn't one"),
1744            (false, true) => PreambleAction::Skip,
1745            (false, false) => PreambleAction::Absent,
1746        };
1747
1748        // How many rows are we actually taking in this take step (including the preamble
1749        // and trailer both as individual rows)
1750        let rows_taking = if *rows_desired >= rows_avail {
1751            // We want all the rows.  If there is a trailer we are grabbing it and will need
1752            // the preamble of the next chunk
1753            // If there is a trailer and we are taking all the rows then we need the preamble
1754            // of the next chunk.
1755            //
1756            // Also, if this chunk is entirely preamble (rows_avail == 0 && !take_trailer) then we
1757            // need the preamble of the next chunk.
1758            *need_preamble = self.take_trailer;
1759            rows_avail
1760        } else {
1761            // We aren't taking all the rows.  Even if there is a trailer we aren't taking
1762            // it so we will not need the preamble
1763            *need_preamble = false;
1764            *rows_desired
1765        };
1766        let rows_skipped = *skip_in_chunk;
1767
1768        // Update the state for the next iteration
1769        let consumed_chunk = if *rows_desired >= rows_avail {
1770            *rows_desired -= rows_avail;
1771            *skip_in_chunk = 0;
1772            true
1773        } else {
1774            *skip_in_chunk += *rows_desired;
1775            *rows_desired = 0;
1776            false
1777        };
1778
1779        (
1780            ChunkDrainInstructions {
1781                chunk_instructions: self.clone(),
1782                rows_to_skip: rows_skipped,
1783                rows_to_take: rows_taking,
1784                preamble_action,
1785            },
1786            consumed_chunk,
1787        )
1788    }
1789}
1790
1791enum Words {
1792    U16(ScalarBuffer<u16>),
1793    U32(ScalarBuffer<u32>),
1794}
1795
1796struct WordsIter<'a> {
1797    iter: Box<dyn Iterator<Item = u32> + 'a>,
1798}
1799
1800impl Words {
1801    pub fn len(&self) -> usize {
1802        match self {
1803            Self::U16(b) => b.len(),
1804            Self::U32(b) => b.len(),
1805        }
1806    }
1807
1808    pub fn iter(&self) -> WordsIter<'_> {
1809        match self {
1810            Self::U16(buf) => WordsIter {
1811                iter: Box::new(buf.iter().map(|&x| x as u32)),
1812            },
1813            Self::U32(buf) => WordsIter {
1814                iter: Box::new(buf.iter().copied()),
1815            },
1816        }
1817    }
1818
1819    pub fn from_bytes(bytes: Bytes, has_large_chunk: bool) -> Result<Self> {
1820        let bytes_per_value = if has_large_chunk { 4 } else { 2 };
1821        assert_eq!(bytes.len() % bytes_per_value, 0);
1822        let buffer = LanceBuffer::from_bytes(bytes, bytes_per_value as u64);
1823        if has_large_chunk {
1824            Ok(Self::U32(buffer.borrow_to_typed_slice::<u32>()))
1825        } else {
1826            Ok(Self::U16(buffer.borrow_to_typed_slice::<u16>()))
1827        }
1828    }
1829}
1830
1831impl<'a> Iterator for WordsIter<'a> {
1832    type Item = u32;
1833
1834    fn next(&mut self) -> Option<Self::Item> {
1835        self.iter.next()
1836    }
1837}
1838
1839impl StructuralPageScheduler for MiniBlockScheduler {
1840    fn initialize<'a>(
1841        &'a mut self,
1842        io: &Arc<dyn EncodingsIo>,
1843    ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
1844        // We always need to fetch chunk metadata.  We may also need to fetch a dictionary and
1845        // we may also need to fetch the repetition index.  Here, we gather what buffers we
1846        // need.
1847        let (meta_buf_position, meta_buf_size) = self.buffer_offsets_and_sizes[0];
1848        let value_buf_position = self.buffer_offsets_and_sizes[1].0;
1849        let mut bufs_needed = 1;
1850        if self.dictionary.is_some() {
1851            bufs_needed += 1;
1852        }
1853        if self.repetition_index_depth > 0 {
1854            bufs_needed += 1;
1855        }
1856        let mut required_ranges = Vec::with_capacity(bufs_needed);
1857        required_ranges.push(meta_buf_position..meta_buf_position + meta_buf_size);
1858        if let Some(ref dictionary) = self.dictionary {
1859            required_ranges.push(
1860                dictionary.dictionary_buf_position_and_size.0
1861                    ..dictionary.dictionary_buf_position_and_size.0
1862                        + dictionary.dictionary_buf_position_and_size.1,
1863            );
1864        }
1865        if self.repetition_index_depth > 0 {
1866            let (rep_index_pos, rep_index_size) = self.buffer_offsets_and_sizes.last().unwrap();
1867            required_ranges.push(*rep_index_pos..*rep_index_pos + *rep_index_size);
1868        }
1869        let io_req = io.submit_request(required_ranges, 0);
1870
1871        async move {
1872            let mut buffers = io_req.await?.into_iter().fuse();
1873            let meta_bytes = buffers.next().unwrap();
1874            let dictionary_bytes = self.dictionary.as_ref().and_then(|_| buffers.next());
1875            let rep_index_bytes = buffers.next();
1876
1877            // Parse the metadata and build the chunk meta
1878            let words = Words::from_bytes(meta_bytes, self.has_large_chunk)?;
1879            let mut chunk_meta = Vec::with_capacity(words.len());
1880
1881            let mut rows_counter = 0;
1882            let mut offset_bytes = value_buf_position;
1883            for (word_idx, word) in words.iter().enumerate() {
1884                let log_num_values = word & 0x0F;
1885                let divided_bytes = word >> 4;
1886                let num_bytes = (divided_bytes as usize + 1) * MINIBLOCK_ALIGNMENT;
1887                debug_assert!(num_bytes > 0);
1888                let num_values = if word_idx < words.len() - 1 {
1889                    debug_assert!(log_num_values > 0);
1890                    1 << log_num_values
1891                } else {
1892                    debug_assert!(
1893                        log_num_values == 0
1894                            || (1 << log_num_values) == (self.items_in_page - rows_counter)
1895                    );
1896                    self.items_in_page - rows_counter
1897                };
1898                rows_counter += num_values;
1899
1900                chunk_meta.push(ChunkMeta {
1901                    num_values,
1902                    chunk_size_bytes: num_bytes as u64,
1903                    offset_bytes,
1904                });
1905                offset_bytes += num_bytes as u64;
1906            }
1907
1908            // Build the repetition index
1909            let rep_index = if let Some(rep_index_data) = rep_index_bytes {
1910                assert!(rep_index_data.len() % 8 == 0);
1911                let stride = self.repetition_index_depth as usize + 1;
1912                MiniBlockRepIndex::decode_from_bytes(&rep_index_data, stride)
1913            } else {
1914                MiniBlockRepIndex::default_from_chunks(&chunk_meta)
1915            };
1916
1917            let mut page_meta = MiniBlockCacheableState {
1918                chunk_meta,
1919                rep_index,
1920                dictionary: None,
1921            };
1922
1923            // decode dictionary
1924            if let Some(ref mut dictionary) = self.dictionary {
1925                let dictionary_data = dictionary_bytes.unwrap();
1926                page_meta.dictionary =
1927                    Some(Arc::new(dictionary.dictionary_decompressor.decompress(
1928                        LanceBuffer::from_bytes(
1929                            dictionary_data,
1930                            dictionary.dictionary_data_alignment,
1931                        ),
1932                        dictionary.num_dictionary_items,
1933                    )?));
1934            };
1935            let page_meta = Arc::new(page_meta);
1936            self.page_meta = Some(page_meta.clone());
1937            Ok(page_meta as Arc<dyn CachedPageData>)
1938        }
1939        .boxed()
1940    }
1941
1942    fn load(&mut self, data: &Arc<dyn CachedPageData>) {
1943        self.page_meta = Some(
1944            data.clone()
1945                .as_arc_any()
1946                .downcast::<MiniBlockCacheableState>()
1947                .unwrap(),
1948        );
1949    }
1950
1951    fn schedule_ranges(
1952        &self,
1953        ranges: &[Range<u64>],
1954        io: &Arc<dyn EncodingsIo>,
1955    ) -> Result<Vec<PageLoadTask>> {
1956        let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
1957
1958        let page_meta = self.page_meta.as_ref().unwrap();
1959
1960        let chunk_instructions =
1961            ChunkInstructions::schedule_instructions(&page_meta.rep_index, ranges);
1962
1963        debug_assert_eq!(
1964            num_rows,
1965            chunk_instructions
1966                .iter()
1967                .map(|ci| ci.rows_to_take)
1968                .sum::<u64>()
1969        );
1970
1971        let chunks_needed = chunk_instructions
1972            .iter()
1973            .map(|ci| ci.chunk_idx)
1974            .unique()
1975            .collect::<Vec<_>>();
1976
1977        let mut loaded_chunks = self.lookup_chunks(&chunks_needed);
1978        let chunk_ranges = loaded_chunks
1979            .iter()
1980            .map(|c| c.byte_range.clone())
1981            .collect::<Vec<_>>();
1982        let loaded_chunk_data = io.submit_request(chunk_ranges, self.priority);
1983
1984        let rep_decompressor = self.rep_decompressor.clone();
1985        let def_decompressor = self.def_decompressor.clone();
1986        let value_decompressor = self.value_decompressor.clone();
1987        let num_buffers = self.num_buffers;
1988        let has_large_chunk = self.has_large_chunk;
1989        let dictionary = page_meta
1990            .dictionary
1991            .as_ref()
1992            .map(|dictionary| dictionary.clone());
1993        let def_meaning = self.def_meaning.clone();
1994
1995        let res = async move {
1996            let loaded_chunk_data = loaded_chunk_data.await?;
1997            for (loaded_chunk, chunk_data) in loaded_chunks.iter_mut().zip(loaded_chunk_data) {
1998                loaded_chunk.data = LanceBuffer::from_bytes(chunk_data, 1);
1999            }
2000
2001            Ok(Box::new(MiniBlockDecoder {
2002                rep_decompressor,
2003                def_decompressor,
2004                value_decompressor,
2005                def_meaning,
2006                loaded_chunks: VecDeque::from_iter(loaded_chunks),
2007                instructions: VecDeque::from(chunk_instructions),
2008                offset_in_current_chunk: 0,
2009                dictionary,
2010                num_rows,
2011                num_buffers,
2012                has_large_chunk,
2013            }) as Box<dyn StructuralPageDecoder>)
2014        }
2015        .boxed();
2016        let page_load_task = PageLoadTask {
2017            decoder_fut: res,
2018            num_rows,
2019        };
2020        Ok(vec![page_load_task])
2021    }
2022}
2023
2024#[derive(Debug, Clone, Copy)]
2025struct FullZipRepIndexDetails {
2026    buf_position: u64,
2027    bytes_per_value: u64, // Will be 1, 2, 4, or 8
2028}
2029
2030#[derive(Debug)]
2031enum PerValueDecompressor {
2032    Fixed(Arc<dyn FixedPerValueDecompressor>),
2033    Variable(Arc<dyn VariablePerValueDecompressor>),
2034}
2035
2036#[derive(Debug)]
2037struct FullZipDecodeDetails {
2038    value_decompressor: PerValueDecompressor,
2039    def_meaning: Arc<[DefinitionInterpretation]>,
2040    ctrl_word_parser: ControlWordParser,
2041    max_rep: u16,
2042    max_visible_def: u16,
2043}
2044
2045/// Describes where FullZip byte ranges should be read from.
2046///
2047/// FullZip decoding always needs a list of byte ranges, but those bytes can come
2048/// from two different places:
2049/// - Remote I/O (normal path): ranges are fetched from the underlying `EncodingsIo`.
2050/// - A prefetched full page (full scan fast path): the entire page has already been
2051///   loaded once and ranges should be sliced from memory.
2052///
2053/// This abstraction keeps scheduling code focused on "which ranges are needed"
2054/// instead of "how bytes are fetched", and it lets full-page scans avoid the
2055/// two-stage rep-index -> data I/O pipeline.
2056#[derive(Debug, Clone)]
2057enum FullZipReadSource {
2058    /// Fetch ranges from the storage backend through the encoding I/O interface.
2059    Remote(Arc<dyn EncodingsIo>),
2060    /// Slice ranges from an already-loaded FullZip page buffer.
2061    PrefetchedPage { base_offset: u64, data: LanceBuffer },
2062}
2063
2064impl FullZipReadSource {
2065    /// Materialize the requested ranges as decode-ready `LanceBuffer`s.
2066    ///
2067    /// The returned buffers preserve the input range order.
2068    fn fetch(
2069        &self,
2070        ranges: &[Range<u64>],
2071        priority: u64,
2072    ) -> BoxFuture<'static, Result<VecDeque<LanceBuffer>>> {
2073        match self {
2074            Self::Remote(io) => {
2075                let io = io.clone();
2076                let ranges = ranges.to_vec();
2077                async move {
2078                    let data = io.submit_request(ranges, priority).await?;
2079                    Ok(data
2080                        .into_iter()
2081                        .map(|bytes| LanceBuffer::from_bytes(bytes, 1))
2082                        .collect::<VecDeque<_>>())
2083                }
2084                .boxed()
2085            }
2086            Self::PrefetchedPage { base_offset, data } => {
2087                let base_offset = *base_offset;
2088                let data = data.clone();
2089                let page_end = base_offset + data.len() as u64;
2090                std::future::ready(
2091                    ranges
2092                        .iter()
2093                        .map(|range| {
2094                            if range.start > range.end
2095                                || range.start < base_offset
2096                                || range.end > page_end
2097                            {
2098                                return Err(Error::internal(format!(
2099                                    "Requested range {:?} is outside page range {}..{}",
2100                                    range, base_offset, page_end
2101                                )));
2102                            }
2103                            let start = (range.start - base_offset) as usize;
2104                            let len = (range.end - range.start) as usize;
2105                            Ok(data.slice_with_length(start, len))
2106                        })
2107                        .collect::<Result<VecDeque<_>>>(),
2108                )
2109                .boxed()
2110            }
2111        }
2112    }
2113}
2114
2115/// A scheduler for full-zip encoded data
2116///
2117/// When the data type has a fixed-width then we simply need to map from
2118/// row ranges to byte ranges using the fixed-width of the data type.
2119///
2120/// When the data type is variable-width or has any repetition then a
2121/// repetition index is required.
2122#[derive(Debug)]
2123pub struct FullZipScheduler {
2124    data_buf_position: u64,
2125    data_buf_size: u64,
2126    rep_index: Option<FullZipRepIndexDetails>,
2127    priority: u64,
2128    rows_in_page: u64,
2129    bits_per_offset: u8,
2130    details: Arc<FullZipDecodeDetails>,
2131    /// Cached state containing the decoded repetition index
2132    cached_state: Option<Arc<FullZipCacheableState>>,
2133    /// Whether repetition index metadata should be cached during initialize.
2134    enable_cache: bool,
2135}
2136
2137impl FullZipScheduler {
2138    fn try_new(
2139        buffer_offsets_and_sizes: &[(u64, u64)],
2140        priority: u64,
2141        rows_in_page: u64,
2142        layout: &pb21::FullZipLayout,
2143        decompressors: &dyn DecompressionStrategy,
2144    ) -> Result<Self> {
2145        let (data_buf_position, data_buf_size) = buffer_offsets_and_sizes[0];
2146        let rep_index = buffer_offsets_and_sizes.get(1).map(|(pos, len)| {
2147            let num_reps = rows_in_page + 1;
2148            let bytes_per_rep = len / num_reps;
2149            debug_assert_eq!(len % num_reps, 0);
2150            debug_assert!(
2151                bytes_per_rep == 1
2152                    || bytes_per_rep == 2
2153                    || bytes_per_rep == 4
2154                    || bytes_per_rep == 8
2155            );
2156            FullZipRepIndexDetails {
2157                buf_position: *pos,
2158                bytes_per_value: bytes_per_rep,
2159            }
2160        });
2161
2162        let value_decompressor = match layout.details {
2163            Some(pb21::full_zip_layout::Details::BitsPerValue(_)) => {
2164                let decompressor = decompressors.create_fixed_per_value_decompressor(
2165                    layout.value_compression.as_ref().unwrap(),
2166                )?;
2167                PerValueDecompressor::Fixed(decompressor.into())
2168            }
2169            Some(pb21::full_zip_layout::Details::BitsPerOffset(_)) => {
2170                let decompressor = decompressors.create_variable_per_value_decompressor(
2171                    layout.value_compression.as_ref().unwrap(),
2172                )?;
2173                PerValueDecompressor::Variable(decompressor.into())
2174            }
2175            None => {
2176                panic!("Full-zip layout must have a `details` field");
2177            }
2178        };
2179        let ctrl_word_parser = ControlWordParser::new(
2180            layout.bits_rep.try_into().unwrap(),
2181            layout.bits_def.try_into().unwrap(),
2182        );
2183        let def_meaning = layout
2184            .layers
2185            .iter()
2186            .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
2187            .collect::<Vec<_>>();
2188
2189        let max_rep = def_meaning.iter().filter(|d| d.is_list()).count() as u16;
2190        let max_visible_def = def_meaning
2191            .iter()
2192            .filter(|d| !d.is_list())
2193            .map(|d| d.num_def_levels())
2194            .sum();
2195
2196        let bits_per_offset = match layout.details {
2197            Some(pb21::full_zip_layout::Details::BitsPerValue(_)) => 32,
2198            Some(pb21::full_zip_layout::Details::BitsPerOffset(bits_per_offset)) => {
2199                bits_per_offset as u8
2200            }
2201            None => panic!("Full-zip layout must have a `details` field"),
2202        };
2203
2204        let details = Arc::new(FullZipDecodeDetails {
2205            value_decompressor,
2206            def_meaning: def_meaning.into(),
2207            ctrl_word_parser,
2208            max_rep,
2209            max_visible_def,
2210        });
2211        Ok(Self {
2212            data_buf_position,
2213            data_buf_size,
2214            rep_index,
2215            details,
2216            priority,
2217            rows_in_page,
2218            bits_per_offset,
2219            cached_state: None,
2220            enable_cache: false,
2221        })
2222    }
2223
2224    fn covers_entire_page(ranges: &[Range<u64>], rows_in_page: u64) -> bool {
2225        if ranges.is_empty() {
2226            return false;
2227        }
2228        let mut expected_start = 0;
2229        for range in ranges {
2230            if range.start != expected_start || range.end > rows_in_page || range.end < range.start
2231            {
2232                return false;
2233            }
2234            expected_start = range.end;
2235        }
2236        expected_start == rows_in_page
2237    }
2238
2239    fn create_page_load_task(
2240        io_future: BoxFuture<'static, Result<Vec<Bytes>>>,
2241        num_rows: u64,
2242        details: Arc<FullZipDecodeDetails>,
2243        bits_per_offset: u8,
2244    ) -> PageLoadTask {
2245        let load_task = async move {
2246            let buffers = io_future.await?;
2247            let data = buffers
2248                .into_iter()
2249                .map(|bytes| LanceBuffer::from_bytes(bytes, 1))
2250                .collect::<VecDeque<_>>();
2251            Self::create_decoder(details, data, num_rows, bits_per_offset)
2252        }
2253        .boxed();
2254        PageLoadTask {
2255            decoder_fut: load_task,
2256            num_rows,
2257        }
2258    }
2259
2260    /// Creates a decoder from the loaded data
2261    fn create_decoder(
2262        details: Arc<FullZipDecodeDetails>,
2263        data: VecDeque<LanceBuffer>,
2264        num_rows: u64,
2265        bits_per_offset: u8,
2266    ) -> Result<Box<dyn StructuralPageDecoder>> {
2267        match &details.value_decompressor {
2268            PerValueDecompressor::Fixed(decompressor) => {
2269                let bits_per_value = decompressor.bits_per_value();
2270                if bits_per_value % 8 != 0 {
2271                    return Err(lance_core::Error::not_supported_source("Bit-packed full-zip encoding (non-byte-aligned values) is not yet implemented".into()));
2272                }
2273                let bytes_per_value = bits_per_value / 8;
2274                let total_bytes_per_value =
2275                    bytes_per_value as usize + details.ctrl_word_parser.bytes_per_word();
2276                if total_bytes_per_value == 0 {
2277                    return Err(lance_core::Error::internal(
2278                        "Invalid encoding: per-row byte width must be greater than 0",
2279                    ));
2280                }
2281                Ok(Box::new(FixedFullZipDecoder {
2282                    details,
2283                    data,
2284                    num_rows,
2285                    offset_in_current: 0,
2286                    bytes_per_value: bytes_per_value as usize,
2287                    total_bytes_per_value,
2288                }) as Box<dyn StructuralPageDecoder>)
2289            }
2290            PerValueDecompressor::Variable(_decompressor) => {
2291                Ok(Box::new(VariableFullZipDecoder::new(
2292                    details,
2293                    data,
2294                    num_rows,
2295                    bits_per_offset,
2296                    bits_per_offset,
2297                )))
2298            }
2299        }
2300    }
2301
2302    /// Extracts byte ranges from a repetition index buffer
2303    /// The buffer contains pairs of (start, end) values for each range
2304    fn extract_byte_ranges_from_pairs(
2305        buffer: LanceBuffer,
2306        bytes_per_value: u64,
2307        data_buf_position: u64,
2308    ) -> Vec<Range<u64>> {
2309        ByteUnpacker::new(buffer, bytes_per_value as usize)
2310            .chunks(2)
2311            .into_iter()
2312            .map(|mut c| {
2313                let start = c.next().unwrap() + data_buf_position;
2314                let end = c.next().unwrap() + data_buf_position;
2315                start..end
2316            })
2317            .collect::<Vec<_>>()
2318    }
2319
2320    /// Extracts byte ranges from a cached repetition index buffer
2321    /// The buffer contains all values and we need to extract specific ranges
2322    fn extract_byte_ranges_from_cached(
2323        buffer: &LanceBuffer,
2324        ranges: &[Range<u64>],
2325        bytes_per_value: u64,
2326        data_buf_position: u64,
2327    ) -> Vec<Range<u64>> {
2328        ranges
2329            .iter()
2330            .map(|r| {
2331                let start_offset = (r.start * bytes_per_value) as usize;
2332                let end_offset = (r.end * bytes_per_value) as usize;
2333
2334                let start_slice = &buffer[start_offset..start_offset + bytes_per_value as usize];
2335                let start_val =
2336                    ByteUnpacker::new(start_slice.iter().copied(), bytes_per_value as usize)
2337                        .next()
2338                        .unwrap();
2339
2340                let end_slice = &buffer[end_offset..end_offset + bytes_per_value as usize];
2341                let end_val =
2342                    ByteUnpacker::new(end_slice.iter().copied(), bytes_per_value as usize)
2343                        .next()
2344                        .unwrap();
2345
2346                (data_buf_position + start_val)..(data_buf_position + end_val)
2347            })
2348            .collect()
2349    }
2350
2351    /// Computes the ranges in the repetition index that need to be loaded
2352    fn compute_rep_index_ranges(
2353        ranges: &[Range<u64>],
2354        rep_index: &FullZipRepIndexDetails,
2355    ) -> Vec<Range<u64>> {
2356        ranges
2357            .iter()
2358            .flat_map(|r| {
2359                let first_val_start =
2360                    rep_index.buf_position + (r.start * rep_index.bytes_per_value);
2361                let first_val_end = first_val_start + rep_index.bytes_per_value;
2362                let last_val_start = rep_index.buf_position + (r.end * rep_index.bytes_per_value);
2363                let last_val_end = last_val_start + rep_index.bytes_per_value;
2364                [first_val_start..first_val_end, last_val_start..last_val_end]
2365            })
2366            .collect()
2367    }
2368
2369    /// Schedules ranges in the presence of a repetition index
2370    fn schedule_ranges_rep(
2371        &self,
2372        ranges: &[Range<u64>],
2373        io: &Arc<dyn EncodingsIo>,
2374        rep_index: FullZipRepIndexDetails,
2375    ) -> Result<Vec<PageLoadTask>> {
2376        let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
2377        let data_buf_position = self.data_buf_position;
2378        let priority = self.priority;
2379        let details = self.details.clone();
2380        let bits_per_offset = self.bits_per_offset;
2381
2382        if Self::covers_entire_page(ranges, self.rows_in_page) {
2383            let full_range = self.data_buf_position..(self.data_buf_position + self.data_buf_size);
2384            let page_data = io.submit_single(full_range.clone(), priority);
2385            let load_task = async move {
2386                let page_data = page_data.await?;
2387                let source = FullZipReadSource::PrefetchedPage {
2388                    base_offset: full_range.start,
2389                    data: LanceBuffer::from_bytes(page_data, 1),
2390                };
2391                let read_ranges = vec![full_range];
2392                let data = source.fetch(&read_ranges, priority).await?;
2393                Self::create_decoder(details, data, num_rows, bits_per_offset)
2394            }
2395            .boxed();
2396            let page_load_task = PageLoadTask {
2397                decoder_fut: load_task,
2398                num_rows,
2399            };
2400            return Ok(vec![page_load_task]);
2401        }
2402
2403        if let Some(cached_state) = &self.cached_state {
2404            let byte_ranges = Self::extract_byte_ranges_from_cached(
2405                &cached_state.rep_index_buffer,
2406                ranges,
2407                rep_index.bytes_per_value,
2408                data_buf_position,
2409            );
2410            let io_future = io.submit_request(byte_ranges, priority);
2411            let page_load_task =
2412                Self::create_page_load_task(io_future, num_rows, details, bits_per_offset);
2413            return Ok(vec![page_load_task]);
2414        }
2415
2416        let rep_ranges = Self::compute_rep_index_ranges(ranges, &rep_index);
2417        let rep_data = io.submit_request(rep_ranges, priority);
2418        let io_clone = io.clone();
2419        let load_task = async move {
2420            let rep_data = rep_data.await?;
2421            let rep_buffer = LanceBuffer::concat(
2422                &rep_data
2423                    .into_iter()
2424                    .map(|d| LanceBuffer::from_bytes(d, 1))
2425                    .collect::<Vec<_>>(),
2426            );
2427            let byte_ranges = Self::extract_byte_ranges_from_pairs(
2428                rep_buffer,
2429                rep_index.bytes_per_value,
2430                data_buf_position,
2431            );
2432            let source = FullZipReadSource::Remote(io_clone);
2433            let data = source.fetch(&byte_ranges, priority).await?;
2434            Self::create_decoder(details, data, num_rows, bits_per_offset)
2435        }
2436        .boxed();
2437        let page_load_task = PageLoadTask {
2438            decoder_fut: load_task,
2439            num_rows,
2440        };
2441        Ok(vec![page_load_task])
2442    }
2443
2444    // In the simple case there is no repetition and we just have large fixed-width
2445    // rows of data.  We can just map row ranges to byte ranges directly using the
2446    // fixed-width of the data type.
2447    fn schedule_ranges_simple(
2448        &self,
2449        ranges: &[Range<u64>],
2450        io: &Arc<dyn EncodingsIo>,
2451    ) -> Result<Vec<PageLoadTask>> {
2452        // Convert row ranges to item ranges (i.e. multiply by items per row)
2453        let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
2454
2455        let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor else {
2456            unreachable!()
2457        };
2458
2459        // Convert item ranges to byte ranges (i.e. multiply by bytes per item)
2460        let bits_per_value = decompressor.bits_per_value();
2461        if !bits_per_value.is_multiple_of(8) {
2462            return Err(Error::invalid_input_source(
2463                format!(
2464                    "Full-zip fixed-width values must be byte aligned, got {} bits per value",
2465                    bits_per_value
2466                )
2467                .into(),
2468            ));
2469        }
2470        let bytes_per_value = bits_per_value / 8;
2471        let bytes_per_cw = self.details.ctrl_word_parser.bytes_per_word();
2472        let total_bytes_per_value = bytes_per_value + bytes_per_cw as u64;
2473        let byte_ranges = ranges
2474            .iter()
2475            .map(|r| {
2476                debug_assert!(r.end <= self.rows_in_page);
2477                let start = self.data_buf_position + r.start * total_bytes_per_value;
2478                let end = self.data_buf_position + r.end * total_bytes_per_value;
2479                start..end
2480            })
2481            .collect::<Vec<_>>();
2482
2483        let io_future = io.submit_request(byte_ranges, self.priority);
2484        let page_load_task = Self::create_page_load_task(
2485            io_future,
2486            num_rows,
2487            self.details.clone(),
2488            self.bits_per_offset,
2489        );
2490        Ok(vec![page_load_task])
2491    }
2492}
2493
2494/// Cacheable state for FullZip encoding, storing the decoded repetition index
2495#[derive(Debug)]
2496struct FullZipCacheableState {
2497    /// The raw repetition index buffer for future decoding
2498    rep_index_buffer: LanceBuffer,
2499}
2500
2501impl DeepSizeOf for FullZipCacheableState {
2502    fn deep_size_of_children(&self, _context: &mut Context) -> usize {
2503        self.rep_index_buffer.len()
2504    }
2505}
2506
2507impl CachedPageData for FullZipCacheableState {
2508    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
2509        self
2510    }
2511}
2512
2513impl StructuralPageScheduler for FullZipScheduler {
2514    fn initialize<'a>(
2515        &'a mut self,
2516        io: &Arc<dyn EncodingsIo>,
2517    ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
2518        if self.enable_cache
2519            && let Some(rep_index) = self.rep_index
2520        {
2521            let total_size = (self.rows_in_page + 1) * rep_index.bytes_per_value;
2522            let rep_index_range = rep_index.buf_position..(rep_index.buf_position + total_size);
2523            let io_clone = io.clone();
2524            return async move {
2525                let rep_index_data = io_clone.submit_request(vec![rep_index_range], 0).await?;
2526                let state = Arc::new(FullZipCacheableState {
2527                    rep_index_buffer: LanceBuffer::from_bytes(rep_index_data[0].clone(), 1),
2528                });
2529                self.cached_state = Some(state.clone());
2530                Ok(state as Arc<dyn CachedPageData>)
2531            }
2532            .boxed();
2533        }
2534        std::future::ready(Ok(Arc::new(NoCachedPageData) as Arc<dyn CachedPageData>)).boxed()
2535    }
2536
2537    /// Loads previously cached repetition index data from the cache system.
2538    /// This method is called when a scheduler instance needs to use cached data
2539    /// that was initialized by another instance or in a previous operation.
2540    fn load(&mut self, cache: &Arc<dyn CachedPageData>) {
2541        // Try to downcast to our specific cache type
2542        if let Ok(cached_state) = cache
2543            .clone()
2544            .as_arc_any()
2545            .downcast::<FullZipCacheableState>()
2546        {
2547            // Store the cached state for use in schedule_ranges
2548            self.cached_state = Some(cached_state);
2549        }
2550    }
2551
2552    fn schedule_ranges(
2553        &self,
2554        ranges: &[Range<u64>],
2555        io: &Arc<dyn EncodingsIo>,
2556    ) -> Result<Vec<PageLoadTask>> {
2557        if let Some(rep_index) = self.rep_index {
2558            self.schedule_ranges_rep(ranges, io, rep_index)
2559        } else {
2560            self.schedule_ranges_simple(ranges, io)
2561        }
2562    }
2563}
2564
2565/// A decoder for full-zip encoded data when the data has a fixed-width
2566///
2567/// Here we need to unzip the control words from the values themselves and
2568/// then decompress the requested values.
2569///
2570/// We use a PerValueDecompressor because we will only be decompressing the
2571/// requested data.  This decoder / scheduler does not do any read amplification.
2572#[derive(Debug)]
2573struct FixedFullZipDecoder {
2574    details: Arc<FullZipDecodeDetails>,
2575    data: VecDeque<LanceBuffer>,
2576    offset_in_current: usize,
2577    bytes_per_value: usize,
2578    total_bytes_per_value: usize,
2579    num_rows: u64,
2580}
2581
2582impl FixedFullZipDecoder {
2583    fn slice_next_task(&mut self, num_rows: u64) -> FullZipDecodeTaskItem {
2584        debug_assert!(num_rows > 0);
2585        let cur_buf = self.data.front_mut().unwrap();
2586        let start = self.offset_in_current;
2587        if self.details.ctrl_word_parser.has_rep() {
2588            // This is a slightly slower path.  In order to figure out where to split we need to
2589            // examine the rep index so we can convert num_lists to num_rows
2590            let mut rows_started = 0;
2591            // We always need at least one value.  Now loop through until we have passed num_rows
2592            // values
2593            let mut num_items = 0;
2594            while self.offset_in_current < cur_buf.len() {
2595                let control = self.details.ctrl_word_parser.parse_desc(
2596                    &cur_buf[self.offset_in_current..],
2597                    self.details.max_rep,
2598                    self.details.max_visible_def,
2599                );
2600                if control.is_new_row {
2601                    if rows_started == num_rows {
2602                        break;
2603                    }
2604                    rows_started += 1;
2605                }
2606                num_items += 1;
2607                if control.is_visible {
2608                    self.offset_in_current += self.total_bytes_per_value;
2609                } else {
2610                    self.offset_in_current += self.details.ctrl_word_parser.bytes_per_word();
2611                }
2612            }
2613
2614            let task_slice = cur_buf.slice_with_length(start, self.offset_in_current - start);
2615            if self.offset_in_current == cur_buf.len() {
2616                self.data.pop_front();
2617                self.offset_in_current = 0;
2618            }
2619
2620            FullZipDecodeTaskItem {
2621                data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
2622                    data: task_slice,
2623                    bits_per_value: self.bytes_per_value as u64 * 8,
2624                    num_values: num_items,
2625                    block_info: BlockInfo::new(),
2626                }),
2627                rows_in_buf: rows_started,
2628            }
2629        } else {
2630            // If there's no repetition we can calculate the slicing point by just multiplying
2631            // the number of rows by the total bytes per value
2632            let cur_buf = self.data.front_mut().unwrap();
2633            let bytes_avail = cur_buf.len() - self.offset_in_current;
2634            let offset_in_cur = self.offset_in_current;
2635
2636            let bytes_needed = num_rows as usize * self.total_bytes_per_value;
2637            let mut rows_taken = num_rows;
2638            let task_slice = if bytes_needed >= bytes_avail {
2639                self.offset_in_current = 0;
2640                rows_taken = bytes_avail as u64 / self.total_bytes_per_value as u64;
2641                self.data
2642                    .pop_front()
2643                    .unwrap()
2644                    .slice_with_length(offset_in_cur, bytes_avail)
2645            } else {
2646                self.offset_in_current += bytes_needed;
2647                cur_buf.slice_with_length(offset_in_cur, bytes_needed)
2648            };
2649            FullZipDecodeTaskItem {
2650                data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
2651                    data: task_slice,
2652                    bits_per_value: self.bytes_per_value as u64 * 8,
2653                    num_values: rows_taken,
2654                    block_info: BlockInfo::new(),
2655                }),
2656                rows_in_buf: rows_taken,
2657            }
2658        }
2659    }
2660}
2661
2662impl StructuralPageDecoder for FixedFullZipDecoder {
2663    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
2664        let mut task_data = Vec::with_capacity(self.data.len());
2665        let mut remaining = num_rows;
2666        while remaining > 0 {
2667            let task_item = self.slice_next_task(remaining);
2668            remaining -= task_item.rows_in_buf;
2669            task_data.push(task_item);
2670        }
2671        Ok(Box::new(FixedFullZipDecodeTask {
2672            details: self.details.clone(),
2673            data: task_data,
2674            bytes_per_value: self.bytes_per_value,
2675            num_rows: num_rows as usize,
2676        }))
2677    }
2678
2679    fn num_rows(&self) -> u64 {
2680        self.num_rows
2681    }
2682}
2683
2684/// A decoder for full-zip encoded data when the data has a variable-width
2685///
2686/// Here we need to unzip the control words AND lengths from the values and
2687/// then decompress the requested values.
2688#[derive(Debug)]
2689struct VariableFullZipDecoder {
2690    details: Arc<FullZipDecodeDetails>,
2691    decompressor: Arc<dyn VariablePerValueDecompressor>,
2692    data: LanceBuffer,
2693    offsets: LanceBuffer,
2694    rep: ScalarBuffer<u16>,
2695    def: ScalarBuffer<u16>,
2696    repdef_starts: Vec<usize>,
2697    data_starts: Vec<usize>,
2698    offset_starts: Vec<usize>,
2699    visible_item_counts: Vec<u64>,
2700    bits_per_offset: u8,
2701    current_idx: usize,
2702    num_rows: u64,
2703}
2704
2705impl VariableFullZipDecoder {
2706    fn new(
2707        details: Arc<FullZipDecodeDetails>,
2708        data: VecDeque<LanceBuffer>,
2709        num_rows: u64,
2710        in_bits_per_length: u8,
2711        out_bits_per_offset: u8,
2712    ) -> Self {
2713        let decompressor = match details.value_decompressor {
2714            PerValueDecompressor::Variable(ref d) => d.clone(),
2715            _ => unreachable!(),
2716        };
2717
2718        assert_eq!(in_bits_per_length % 8, 0);
2719        assert!(out_bits_per_offset == 32 || out_bits_per_offset == 64);
2720
2721        let mut decoder = Self {
2722            details,
2723            decompressor,
2724            data: LanceBuffer::empty(),
2725            offsets: LanceBuffer::empty(),
2726            rep: LanceBuffer::empty().borrow_to_typed_slice(),
2727            def: LanceBuffer::empty().borrow_to_typed_slice(),
2728            bits_per_offset: out_bits_per_offset,
2729            repdef_starts: Vec::with_capacity(num_rows as usize + 1),
2730            data_starts: Vec::with_capacity(num_rows as usize + 1),
2731            offset_starts: Vec::with_capacity(num_rows as usize + 1),
2732            visible_item_counts: Vec::with_capacity(num_rows as usize + 1),
2733            current_idx: 0,
2734            num_rows,
2735        };
2736
2737        // There's no great time to do this and this is the least worst time.  If we don't unzip then
2738        // we can't slice the data during the decode phase.  This is because we need the offsets to be
2739        // unpacked to know where the values start and end.
2740        //
2741        // We don't want to unzip on the decode thread because that is a single-threaded path
2742        // We don't want to unzip on the scheduling thread because that is a single-threaded path
2743        //
2744        // Fortunately, we know variable length data will always be read indirectly and so we can do it
2745        // here, which should be on the indirect thread.  The primary disadvantage to doing it here is that
2746        // we load all the data into memory and then throw it away only to load it all into memory again during
2747        // the decode.
2748        //
2749        // There are some alternatives to investigate:
2750        //   - Instead of just reading the beginning and end of the rep index we could read the entire
2751        //     range in between.  This will give us the break points that we need for slicing and won't increase
2752        //     the number of IOPs but it will mean we are doing more total I/O and we need to load the rep index
2753        //     even when doing a full scan.
2754        //   - We could force each decode task to do a full unzip of all the data.  Each decode task now
2755        //     has to do more work but the work is all fused.
2756        //   - We could just try doing this work on the decode thread and see if it is a problem.
2757        decoder.unzip(data, in_bits_per_length, out_bits_per_offset, num_rows);
2758
2759        decoder
2760    }
2761
2762    fn slice_batch_data_and_rebase_offsets_typed<T>(
2763        data: &LanceBuffer,
2764        offsets: &LanceBuffer,
2765    ) -> Result<(LanceBuffer, LanceBuffer)>
2766    where
2767        T: arrow_buffer::ArrowNativeType
2768            + Copy
2769            + PartialOrd
2770            + std::ops::Sub<Output = T>
2771            + std::fmt::Display
2772            + TryInto<usize>,
2773    {
2774        let offsets_slice = offsets.borrow_to_typed_slice::<T>();
2775        let offsets_slice = offsets_slice.as_ref();
2776        if offsets_slice.is_empty() {
2777            return Err(Error::internal(
2778                "Variable offsets cannot be empty".to_string(),
2779            ));
2780        }
2781
2782        let base = offsets_slice[0];
2783        let end = *offsets_slice.last().unwrap();
2784        if end < base {
2785            return Err(Error::internal(format!(
2786                "Invalid variable offsets: end ({end}) is less than base ({base})"
2787            )));
2788        }
2789
2790        let data_start = base.try_into().map_err(|_| {
2791            Error::internal(format!("Variable offset ({base}) does not fit into usize"))
2792        })?;
2793        let data_end = end.try_into().map_err(|_| {
2794            Error::internal(format!("Variable offset ({end}) does not fit into usize"))
2795        })?;
2796        if data_end > data.len() {
2797            return Err(Error::internal(format!(
2798                "Invalid variable offsets: end ({data_end}) exceeds data len ({})",
2799                data.len()
2800            )));
2801        }
2802
2803        let mut rebased_offsets = Vec::with_capacity(offsets_slice.len());
2804        for &offset in offsets_slice {
2805            if offset < base {
2806                return Err(Error::internal(format!(
2807                    "Invalid variable offsets: offset ({offset}) is less than base ({base})"
2808                )));
2809            }
2810            rebased_offsets.push(offset - base);
2811        }
2812
2813        let sliced_data = data.slice_with_length(data_start, data_end - data_start);
2814        // Copy into a compact buffer so each output batch owns only what it references.
2815        let sliced_data = LanceBuffer::copy_slice(&sliced_data);
2816        let rebased_offsets = LanceBuffer::reinterpret_vec(rebased_offsets);
2817        Ok((sliced_data, rebased_offsets))
2818    }
2819
2820    fn slice_batch_data_and_rebase_offsets(
2821        data: &LanceBuffer,
2822        offsets: &LanceBuffer,
2823        bits_per_offset: u8,
2824    ) -> Result<(LanceBuffer, LanceBuffer)> {
2825        match bits_per_offset {
2826            32 => Self::slice_batch_data_and_rebase_offsets_typed::<u32>(data, offsets),
2827            64 => Self::slice_batch_data_and_rebase_offsets_typed::<u64>(data, offsets),
2828            _ => Err(Error::internal(format!(
2829                "Unsupported bits_per_offset={bits_per_offset}"
2830            ))),
2831        }
2832    }
2833
2834    unsafe fn parse_length(data: &[u8], bits_per_offset: u8) -> u64 {
2835        match bits_per_offset {
2836            8 => *data.get_unchecked(0) as u64,
2837            16 => u16::from_le_bytes([*data.get_unchecked(0), *data.get_unchecked(1)]) as u64,
2838            32 => u32::from_le_bytes([
2839                *data.get_unchecked(0),
2840                *data.get_unchecked(1),
2841                *data.get_unchecked(2),
2842                *data.get_unchecked(3),
2843            ]) as u64,
2844            64 => u64::from_le_bytes([
2845                *data.get_unchecked(0),
2846                *data.get_unchecked(1),
2847                *data.get_unchecked(2),
2848                *data.get_unchecked(3),
2849                *data.get_unchecked(4),
2850                *data.get_unchecked(5),
2851                *data.get_unchecked(6),
2852                *data.get_unchecked(7),
2853            ]),
2854            _ => unreachable!(),
2855        }
2856    }
2857
2858    fn unzip(
2859        &mut self,
2860        data: VecDeque<LanceBuffer>,
2861        in_bits_per_length: u8,
2862        out_bits_per_offset: u8,
2863        num_rows: u64,
2864    ) {
2865        // This undercounts if there are lists but, at this point, we don't really know how many items we have
2866        let mut rep = Vec::with_capacity(num_rows as usize);
2867        let mut def = Vec::with_capacity(num_rows as usize);
2868        let bytes_cw = self.details.ctrl_word_parser.bytes_per_word() * num_rows as usize;
2869
2870        // This undercounts if there are lists
2871        // It can also overcount if there are invisible items
2872        let bytes_per_offset = out_bits_per_offset as usize / 8;
2873        let bytes_offsets = bytes_per_offset * (num_rows as usize + 1);
2874        let mut offsets_data = Vec::with_capacity(bytes_offsets);
2875
2876        let bytes_per_length = in_bits_per_length as usize / 8;
2877        let bytes_lengths = bytes_per_length * num_rows as usize;
2878
2879        let bytes_data = data.iter().map(|d| d.len()).sum::<usize>();
2880        // This overcounts since bytes_lengths and bytes_cw are undercounts
2881        // It can also undercount if there are invisible items (hence the saturating_sub)
2882        let mut unzipped_data =
2883            Vec::with_capacity((bytes_data - bytes_cw).saturating_sub(bytes_lengths));
2884
2885        let mut current_offset = 0_u64;
2886        let mut visible_item_count = 0_u64;
2887        for databuf in data.into_iter() {
2888            let mut databuf = databuf.as_ref();
2889            while !databuf.is_empty() {
2890                let data_start = unzipped_data.len();
2891                let offset_start = offsets_data.len();
2892                // We might have only-rep or only-def, neither, or both.  They move at the same
2893                // speed though so we only need one index into it
2894                let repdef_start = rep.len().max(def.len());
2895                // TODO: Kind of inefficient we parse the control word twice here
2896                let ctrl_desc = self.details.ctrl_word_parser.parse_desc(
2897                    databuf,
2898                    self.details.max_rep,
2899                    self.details.max_visible_def,
2900                );
2901                self.details
2902                    .ctrl_word_parser
2903                    .parse(databuf, &mut rep, &mut def);
2904                databuf = &databuf[self.details.ctrl_word_parser.bytes_per_word()..];
2905
2906                if ctrl_desc.is_new_row {
2907                    self.repdef_starts.push(repdef_start);
2908                    self.data_starts.push(data_start);
2909                    self.offset_starts.push(offset_start);
2910                    self.visible_item_counts.push(visible_item_count);
2911                }
2912                if ctrl_desc.is_visible {
2913                    visible_item_count += 1;
2914                    if ctrl_desc.is_valid_item {
2915                        // Safety: Data should have at least bytes_per_length bytes remaining
2916                        debug_assert!(databuf.len() >= bytes_per_length);
2917                        let length = unsafe { Self::parse_length(databuf, in_bits_per_length) };
2918                        match out_bits_per_offset {
2919                            32 => offsets_data
2920                                .extend_from_slice(&(current_offset as u32).to_le_bytes()),
2921                            64 => offsets_data.extend_from_slice(&current_offset.to_le_bytes()),
2922                            _ => unreachable!(),
2923                        };
2924                        databuf = &databuf[bytes_per_offset..];
2925                        unzipped_data.extend_from_slice(&databuf[..length as usize]);
2926                        databuf = &databuf[length as usize..];
2927                        current_offset += length;
2928                    } else {
2929                        // Null items still get an offset
2930                        match out_bits_per_offset {
2931                            32 => offsets_data
2932                                .extend_from_slice(&(current_offset as u32).to_le_bytes()),
2933                            64 => offsets_data.extend_from_slice(&current_offset.to_le_bytes()),
2934                            _ => unreachable!(),
2935                        }
2936                    }
2937                }
2938            }
2939        }
2940        self.repdef_starts.push(rep.len().max(def.len()));
2941        self.data_starts.push(unzipped_data.len());
2942        self.offset_starts.push(offsets_data.len());
2943        self.visible_item_counts.push(visible_item_count);
2944        match out_bits_per_offset {
2945            32 => offsets_data.extend_from_slice(&(current_offset as u32).to_le_bytes()),
2946            64 => offsets_data.extend_from_slice(&current_offset.to_le_bytes()),
2947            _ => unreachable!(),
2948        };
2949        self.rep = ScalarBuffer::from(rep);
2950        self.def = ScalarBuffer::from(def);
2951        self.data = LanceBuffer::from(unzipped_data);
2952        self.offsets = LanceBuffer::from(offsets_data);
2953    }
2954}
2955
2956impl StructuralPageDecoder for VariableFullZipDecoder {
2957    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
2958        let start = self.current_idx;
2959        let end = start + num_rows as usize;
2960
2961        let offset_start = self.offset_starts[start];
2962        let offset_end = self.offset_starts[end] + (self.bits_per_offset as usize / 8);
2963        let offsets = self
2964            .offsets
2965            .slice_with_length(offset_start, offset_end - offset_start);
2966        // Keep each batch's variable data buffer bounded to the selected rows.
2967        let (data, offsets) =
2968            Self::slice_batch_data_and_rebase_offsets(&self.data, &offsets, self.bits_per_offset)?;
2969
2970        let repdef_start = self.repdef_starts[start];
2971        let repdef_end = self.repdef_starts[end];
2972        let rep = if self.rep.is_empty() {
2973            self.rep.clone()
2974        } else {
2975            self.rep.slice(repdef_start, repdef_end - repdef_start)
2976        };
2977        let def = if self.def.is_empty() {
2978            self.def.clone()
2979        } else {
2980            self.def.slice(repdef_start, repdef_end - repdef_start)
2981        };
2982
2983        let visible_item_counts_start = self.visible_item_counts[start];
2984        let visible_item_counts_end = self.visible_item_counts[end];
2985        let num_visible_items = visible_item_counts_end - visible_item_counts_start;
2986
2987        self.current_idx += num_rows as usize;
2988
2989        Ok(Box::new(VariableFullZipDecodeTask {
2990            details: self.details.clone(),
2991            decompressor: self.decompressor.clone(),
2992            data,
2993            offsets,
2994            bits_per_offset: self.bits_per_offset,
2995            num_visible_items,
2996            rep,
2997            def,
2998        }))
2999    }
3000
3001    fn num_rows(&self) -> u64 {
3002        self.num_rows
3003    }
3004}
3005
3006#[derive(Debug)]
3007struct VariableFullZipDecodeTask {
3008    details: Arc<FullZipDecodeDetails>,
3009    decompressor: Arc<dyn VariablePerValueDecompressor>,
3010    data: LanceBuffer,
3011    offsets: LanceBuffer,
3012    bits_per_offset: u8,
3013    num_visible_items: u64,
3014    rep: ScalarBuffer<u16>,
3015    def: ScalarBuffer<u16>,
3016}
3017
3018impl DecodePageTask for VariableFullZipDecodeTask {
3019    fn decode(self: Box<Self>) -> Result<DecodedPage> {
3020        let block = VariableWidthBlock {
3021            data: self.data,
3022            offsets: self.offsets,
3023            bits_per_offset: self.bits_per_offset,
3024            num_values: self.num_visible_items,
3025            block_info: BlockInfo::new(),
3026        };
3027        let decomopressed = self.decompressor.decompress(block)?;
3028        let rep = if self.rep.is_empty() {
3029            None
3030        } else {
3031            Some(self.rep.to_vec())
3032        };
3033        let def = if self.def.is_empty() {
3034            None
3035        } else {
3036            Some(self.def.to_vec())
3037        };
3038        let unraveler = RepDefUnraveler::new(
3039            rep,
3040            def,
3041            self.details.def_meaning.clone(),
3042            self.num_visible_items,
3043        );
3044        Ok(DecodedPage {
3045            data: decomopressed,
3046            repdef: unraveler,
3047        })
3048    }
3049}
3050
3051#[derive(Debug)]
3052struct FullZipDecodeTaskItem {
3053    data: PerValueDataBlock,
3054    rows_in_buf: u64,
3055}
3056
3057/// A task to unzip and decompress full-zip encoded data when that data
3058/// has a fixed-width.
3059#[derive(Debug)]
3060struct FixedFullZipDecodeTask {
3061    details: Arc<FullZipDecodeDetails>,
3062    data: Vec<FullZipDecodeTaskItem>,
3063    num_rows: usize,
3064    bytes_per_value: usize,
3065}
3066
3067impl DecodePageTask for FixedFullZipDecodeTask {
3068    fn decode(self: Box<Self>) -> Result<DecodedPage> {
3069        // Multiply by 2 to make a stab at the size of the output buffer (which will be decompressed and thus bigger)
3070        let estimated_size_bytes = self
3071            .data
3072            .iter()
3073            .map(|task_item| task_item.data.data_size() as usize)
3074            .sum::<usize>()
3075            * 2;
3076        let mut data_builder =
3077            DataBlockBuilder::with_capacity_estimate(estimated_size_bytes as u64);
3078
3079        if self.details.ctrl_word_parser.bytes_per_word() == 0 {
3080            // Fast path, no need to unzip because there is no rep/def
3081            //
3082            // We decompress each buffer and add it to our output buffer
3083            for task_item in self.data.into_iter() {
3084                let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
3085                    unreachable!()
3086                };
3087                let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
3088                else {
3089                    unreachable!()
3090                };
3091                debug_assert_eq!(fixed_data.num_values, task_item.rows_in_buf);
3092                let decompressed = decompressor.decompress(fixed_data, task_item.rows_in_buf)?;
3093                data_builder.append(&decompressed, 0..task_item.rows_in_buf);
3094            }
3095
3096            let unraveler = RepDefUnraveler::new(
3097                None,
3098                None,
3099                self.details.def_meaning.clone(),
3100                self.num_rows as u64,
3101            );
3102
3103            Ok(DecodedPage {
3104                data: data_builder.finish(),
3105                repdef: unraveler,
3106            })
3107        } else {
3108            // Slow path, unzipping needed
3109            let mut rep = Vec::with_capacity(self.num_rows);
3110            let mut def = Vec::with_capacity(self.num_rows);
3111
3112            for task_item in self.data.into_iter() {
3113                let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
3114                    unreachable!()
3115                };
3116                let mut buf_slice = fixed_data.data.as_ref();
3117                let num_values = fixed_data.num_values as usize;
3118                // We will be unzipping repdef in to `rep` and `def` and the
3119                // values into `values` (which contains the compressed values)
3120                let mut values = Vec::with_capacity(
3121                    fixed_data.data.len()
3122                        - (self.details.ctrl_word_parser.bytes_per_word() * num_values),
3123                );
3124                let mut visible_items = 0;
3125                for _ in 0..num_values {
3126                    // Extract rep/def
3127                    self.details
3128                        .ctrl_word_parser
3129                        .parse(buf_slice, &mut rep, &mut def);
3130                    buf_slice = &buf_slice[self.details.ctrl_word_parser.bytes_per_word()..];
3131
3132                    let is_visible = def
3133                        .last()
3134                        .map(|d| *d <= self.details.max_visible_def)
3135                        .unwrap_or(true);
3136                    if is_visible {
3137                        // Extract value
3138                        values.extend_from_slice(buf_slice[..self.bytes_per_value].as_ref());
3139                        buf_slice = &buf_slice[self.bytes_per_value..];
3140                        visible_items += 1;
3141                    }
3142                }
3143
3144                // Finally, we decompress the values and add them to our output buffer
3145                let values_buf = LanceBuffer::from(values);
3146                let fixed_data = FixedWidthDataBlock {
3147                    bits_per_value: self.bytes_per_value as u64 * 8,
3148                    block_info: BlockInfo::new(),
3149                    data: values_buf,
3150                    num_values: visible_items,
3151                };
3152                let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
3153                else {
3154                    unreachable!()
3155                };
3156                let decompressed = decompressor.decompress(fixed_data, visible_items)?;
3157                data_builder.append(&decompressed, 0..visible_items);
3158            }
3159
3160            let repetition = if rep.is_empty() { None } else { Some(rep) };
3161            let definition = if def.is_empty() { None } else { Some(def) };
3162
3163            let unraveler = RepDefUnraveler::new(
3164                repetition,
3165                definition,
3166                self.details.def_meaning.clone(),
3167                self.num_rows as u64,
3168            );
3169            let data = data_builder.finish();
3170
3171            Ok(DecodedPage {
3172                data,
3173                repdef: unraveler,
3174            })
3175        }
3176    }
3177}
3178
3179#[derive(Debug)]
3180struct StructuralPrimitiveFieldSchedulingJob<'a> {
3181    scheduler: &'a StructuralPrimitiveFieldScheduler,
3182    ranges: Vec<Range<u64>>,
3183    page_idx: usize,
3184    range_idx: usize,
3185    global_row_offset: u64,
3186}
3187
3188impl<'a> StructuralPrimitiveFieldSchedulingJob<'a> {
3189    pub fn new(scheduler: &'a StructuralPrimitiveFieldScheduler, ranges: Vec<Range<u64>>) -> Self {
3190        Self {
3191            scheduler,
3192            ranges,
3193            page_idx: 0,
3194            range_idx: 0,
3195            global_row_offset: 0,
3196        }
3197    }
3198}
3199
3200impl StructuralSchedulingJob for StructuralPrimitiveFieldSchedulingJob<'_> {
3201    fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result<Vec<ScheduledScanLine>> {
3202        if self.range_idx >= self.ranges.len() {
3203            return Ok(Vec::new());
3204        }
3205        // Get our current range
3206        let mut range = self.ranges[self.range_idx].clone();
3207        let priority = range.start;
3208
3209        let mut cur_page = &self.scheduler.page_schedulers[self.page_idx];
3210        trace!(
3211            "Current range is {:?} and current page has {} rows",
3212            range, cur_page.num_rows
3213        );
3214        // Skip entire pages until we have some overlap with our next range
3215        while cur_page.num_rows + self.global_row_offset <= range.start {
3216            self.global_row_offset += cur_page.num_rows;
3217            self.page_idx += 1;
3218            trace!("Skipping entire page of {} rows", cur_page.num_rows);
3219            cur_page = &self.scheduler.page_schedulers[self.page_idx];
3220        }
3221
3222        // Now the cur_page has overlap with range.  Continue looping through ranges
3223        // until we find a range that exceeds the current page
3224
3225        let mut ranges_in_page = Vec::new();
3226        while cur_page.num_rows + self.global_row_offset > range.start {
3227            range.start = range.start.max(self.global_row_offset);
3228            let start_in_page = range.start - self.global_row_offset;
3229            let end_in_page = start_in_page + (range.end - range.start);
3230            let end_in_page = end_in_page.min(cur_page.num_rows);
3231            let last_in_range = (end_in_page + self.global_row_offset) >= range.end;
3232
3233            ranges_in_page.push(start_in_page..end_in_page);
3234            if last_in_range {
3235                self.range_idx += 1;
3236                if self.range_idx == self.ranges.len() {
3237                    break;
3238                }
3239                range = self.ranges[self.range_idx].clone();
3240            } else {
3241                break;
3242            }
3243        }
3244
3245        trace!(
3246            "Scheduling {} rows across {} ranges from page with {} rows (priority={}, column_index={}, page_index={})",
3247            ranges_in_page.iter().map(|r| r.end - r.start).sum::<u64>(),
3248            ranges_in_page.len(),
3249            cur_page.num_rows,
3250            priority,
3251            self.scheduler.column_index,
3252            cur_page.page_index,
3253        );
3254
3255        self.global_row_offset += cur_page.num_rows;
3256        self.page_idx += 1;
3257
3258        let page_decoders = cur_page
3259            .scheduler
3260            .schedule_ranges(&ranges_in_page, context.io())?;
3261
3262        let cur_path = context.current_path();
3263        page_decoders
3264            .into_iter()
3265            .map(|page_load_task| {
3266                let cur_path = cur_path.clone();
3267                let page_decoder = page_load_task.decoder_fut;
3268                let unloaded_page = async move {
3269                    let page_decoder = page_decoder.await?;
3270                    Ok(LoadedPageShard {
3271                        decoder: page_decoder,
3272                        path: cur_path,
3273                    })
3274                }
3275                .boxed();
3276                Ok(ScheduledScanLine {
3277                    decoders: vec![MessageType::UnloadedPage(UnloadedPageShard(unloaded_page))],
3278                    rows_scheduled: page_load_task.num_rows,
3279                })
3280            })
3281            .collect::<Result<Vec<_>>>()
3282    }
3283}
3284
3285#[derive(Debug)]
3286struct PageInfoAndScheduler {
3287    page_index: usize,
3288    num_rows: u64,
3289    scheduler: Box<dyn StructuralPageScheduler>,
3290}
3291
3292/// A scheduler for a leaf node
3293///
3294/// Here we look at the layout of the various pages and delegate scheduling to a scheduler
3295/// appropriate for the layout of the page.
3296#[derive(Debug)]
3297pub struct StructuralPrimitiveFieldScheduler {
3298    page_schedulers: Vec<PageInfoAndScheduler>,
3299    column_index: u32,
3300    // Identifies the requested decode shape (e.g. blob descriptor struct vs
3301    // raw bytes). Blob columns can produce multiple page scheduler variants
3302    // for the same physical column depending on the target field's data type,
3303    // and the cached page state types differ per variant. The view tag is
3304    // mixed into the cache key so different variants do not collide.
3305    view_tag: String,
3306}
3307
3308impl StructuralPrimitiveFieldScheduler {
3309    pub fn try_new(
3310        column_info: &ColumnInfo,
3311        decompressors: &dyn DecompressionStrategy,
3312        cache_repetition_index: bool,
3313        target_field: &Field,
3314    ) -> Result<Self> {
3315        let page_schedulers = column_info
3316            .page_infos
3317            .iter()
3318            .enumerate()
3319            .map(|(page_index, page_info)| {
3320                Self::page_info_to_scheduler(
3321                    page_info,
3322                    page_index,
3323                    decompressors,
3324                    cache_repetition_index,
3325                    target_field,
3326                )
3327            })
3328            .collect::<Result<Vec<_>>>()?;
3329        Ok(Self {
3330            page_schedulers,
3331            column_index: column_info.index,
3332            view_tag: format!("{:?}", target_field.data_type()),
3333        })
3334    }
3335
3336    fn page_layout_to_scheduler(
3337        page_info: &PageInfo,
3338        page_layout: &PageLayout,
3339        decompressors: &dyn DecompressionStrategy,
3340        cache_repetition_index: bool,
3341        target_field: &Field,
3342    ) -> Result<Box<dyn StructuralPageScheduler>> {
3343        use pb21::page_layout::Layout;
3344        Ok(match page_layout.layout.as_ref().expect_ok()? {
3345            Layout::MiniBlockLayout(mini_block) => Box::new(MiniBlockScheduler::try_new(
3346                &page_info.buffer_offsets_and_sizes,
3347                page_info.priority,
3348                mini_block.num_items,
3349                mini_block,
3350                decompressors,
3351            )?),
3352            Layout::FullZipLayout(full_zip) => {
3353                let mut scheduler = FullZipScheduler::try_new(
3354                    &page_info.buffer_offsets_and_sizes,
3355                    page_info.priority,
3356                    page_info.num_rows,
3357                    full_zip,
3358                    decompressors,
3359                )?;
3360                scheduler.enable_cache = cache_repetition_index;
3361                Box::new(scheduler)
3362            }
3363            Layout::ConstantLayout(constant_layout) => {
3364                let def_meaning = constant_layout
3365                    .layers
3366                    .iter()
3367                    .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
3368                    .collect::<Vec<_>>();
3369                let has_scalar_value = constant_layout.inline_value.is_some()
3370                    || page_info.buffer_offsets_and_sizes.len() == 1
3371                    || page_info.buffer_offsets_and_sizes.len() == 3;
3372                if has_scalar_value {
3373                    Box::new(constant::ConstantPageScheduler::try_new(
3374                        page_info.buffer_offsets_and_sizes.clone(),
3375                        constant_layout.inline_value.clone(),
3376                        target_field.data_type(),
3377                        def_meaning.into(),
3378                    )?) as Box<dyn StructuralPageScheduler>
3379                } else if def_meaning.len() == 1
3380                    && def_meaning[0] == DefinitionInterpretation::NullableItem
3381                {
3382                    Box::new(SimpleAllNullScheduler::default()) as Box<dyn StructuralPageScheduler>
3383                } else {
3384                    let rep_decompressor = constant_layout
3385                        .rep_compression
3386                        .as_ref()
3387                        .map(|encoding| decompressors.create_block_decompressor(encoding))
3388                        .transpose()?
3389                        .map(Arc::from);
3390
3391                    let def_decompressor = constant_layout
3392                        .def_compression
3393                        .as_ref()
3394                        .map(|encoding| decompressors.create_block_decompressor(encoding))
3395                        .transpose()?
3396                        .map(Arc::from);
3397
3398                    Box::new(ComplexAllNullScheduler::new(
3399                        page_info.buffer_offsets_and_sizes.clone(),
3400                        def_meaning.into(),
3401                        rep_decompressor,
3402                        def_decompressor,
3403                        constant_layout.num_rep_values,
3404                        constant_layout.num_def_values,
3405                    )) as Box<dyn StructuralPageScheduler>
3406                }
3407            }
3408            Layout::BlobLayout(blob) => {
3409                let inner_scheduler = Self::page_layout_to_scheduler(
3410                    page_info,
3411                    blob.inner_layout.as_ref().expect_ok()?.as_ref(),
3412                    decompressors,
3413                    cache_repetition_index,
3414                    target_field,
3415                )?;
3416                let def_meaning = blob
3417                    .layers
3418                    .iter()
3419                    .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
3420                    .collect::<Vec<_>>();
3421                if matches!(target_field.data_type(), DataType::Struct(_)) {
3422                    // User wants to decode blob into struct
3423                    Box::new(BlobDescriptionPageScheduler::new(
3424                        inner_scheduler,
3425                        def_meaning.into(),
3426                    ))
3427                } else {
3428                    // User wants to decode blob into binary data
3429                    Box::new(BlobPageScheduler::new(
3430                        inner_scheduler,
3431                        page_info.priority,
3432                        page_info.num_rows,
3433                        def_meaning.into(),
3434                    ))
3435                }
3436            }
3437        })
3438    }
3439
3440    fn page_info_to_scheduler(
3441        page_info: &PageInfo,
3442        page_index: usize,
3443        decompressors: &dyn DecompressionStrategy,
3444        cache_repetition_index: bool,
3445        target_field: &Field,
3446    ) -> Result<PageInfoAndScheduler> {
3447        let page_layout = page_info.encoding.as_structural();
3448        let scheduler = Self::page_layout_to_scheduler(
3449            page_info,
3450            page_layout,
3451            decompressors,
3452            cache_repetition_index,
3453            target_field,
3454        )?;
3455        Ok(PageInfoAndScheduler {
3456            page_index,
3457            num_rows: page_info.num_rows,
3458            scheduler,
3459        })
3460    }
3461}
3462
3463pub trait CachedPageData: Any + Send + Sync + DeepSizeOf + 'static {
3464    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static>;
3465}
3466
3467pub struct NoCachedPageData;
3468
3469impl DeepSizeOf for NoCachedPageData {
3470    fn deep_size_of_children(&self, _ctx: &mut Context) -> usize {
3471        0
3472    }
3473}
3474impl CachedPageData for NoCachedPageData {
3475    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
3476        self
3477    }
3478}
3479
3480pub struct CachedFieldData {
3481    pages: Vec<Arc<dyn CachedPageData>>,
3482}
3483
3484impl DeepSizeOf for CachedFieldData {
3485    fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
3486        self.pages.deep_size_of_children(ctx)
3487    }
3488}
3489
3490// Cache key for field data
3491//
3492// Both `column_index` and `view_tag` are part of the key because a single
3493// physical column can be decoded under more than one shape — a blob column,
3494// for instance, materializes as a `Struct<position, size>` descriptor in one
3495// scheduler variant and as the raw `LargeBinary` bytes in another. Each
3496// variant builds different `CachedPageData` types per page, so two readers
3497// that hit the same `column_index` with different shapes used to collide and
3498// crash with a downcast failure when loading cached state.
3499#[derive(Debug, Clone)]
3500pub struct FieldDataCacheKey {
3501    pub column_index: u32,
3502    pub view_tag: String,
3503}
3504
3505impl CacheKey for FieldDataCacheKey {
3506    type ValueType = CachedFieldData;
3507
3508    fn key(&self) -> std::borrow::Cow<'_, str> {
3509        format!("{}:{}", self.column_index, self.view_tag).into()
3510    }
3511
3512    fn type_name() -> &'static str {
3513        "FieldData"
3514    }
3515}
3516
3517impl StructuralFieldScheduler for StructuralPrimitiveFieldScheduler {
3518    fn initialize<'a>(
3519        &'a mut self,
3520        _filter: &'a FilterExpression,
3521        context: &'a SchedulerContext,
3522    ) -> BoxFuture<'a, Result<()>> {
3523        let cache_key = FieldDataCacheKey {
3524            column_index: self.column_index,
3525            view_tag: self.view_tag.clone(),
3526        };
3527        let cache = context.cache().clone();
3528
3529        async move {
3530            if let Some(cached_data) = cache.get_with_key(&cache_key).await {
3531                self.page_schedulers
3532                    .iter_mut()
3533                    .zip(cached_data.pages.iter())
3534                    .for_each(|(page_scheduler, cached_data)| {
3535                        page_scheduler.scheduler.load(cached_data);
3536                    });
3537                return Ok(());
3538            }
3539
3540            let page_data = self
3541                .page_schedulers
3542                .iter_mut()
3543                .map(|s| s.scheduler.initialize(context.io()))
3544                .collect::<FuturesOrdered<_>>();
3545
3546            let page_data = page_data.try_collect::<Vec<_>>().await?;
3547            let cached_data = Arc::new(CachedFieldData { pages: page_data });
3548            cache.insert_with_key(&cache_key, cached_data).await;
3549            Ok(())
3550        }
3551        .boxed()
3552    }
3553
3554    fn schedule_ranges<'a>(
3555        &'a self,
3556        ranges: &[Range<u64>],
3557        _filter: &FilterExpression,
3558    ) -> Result<Box<dyn StructuralSchedulingJob + 'a>> {
3559        let ranges = ranges.to_vec();
3560        Ok(Box::new(StructuralPrimitiveFieldSchedulingJob::new(
3561            self, ranges,
3562        )))
3563    }
3564}
3565
3566/// Takes the output from several pages decoders and
3567/// concatenates them.
3568#[derive(Debug)]
3569pub struct StructuralCompositeDecodeArrayTask {
3570    tasks: Vec<Box<dyn DecodePageTask>>,
3571    should_validate: bool,
3572    data_type: DataType,
3573}
3574
3575impl StructuralCompositeDecodeArrayTask {
3576    fn restore_validity(
3577        array: Arc<dyn Array>,
3578        unraveler: &mut CompositeRepDefUnraveler,
3579    ) -> Arc<dyn Array> {
3580        let validity = unraveler.unravel_validity(array.len());
3581        let Some(validity) = validity else {
3582            return array;
3583        };
3584        if array.data_type() == &DataType::Null {
3585            // We unravel from a null array but we don't add the null buffer because arrow-rs doesn't like it
3586            return array;
3587        }
3588        assert_eq!(validity.len(), array.len());
3589        // SAFETY: We've should have already asserted the buffers are all valid, we are just
3590        // adding null buffers to the array here
3591        make_array(unsafe {
3592            array
3593                .to_data()
3594                .into_builder()
3595                .nulls(Some(validity))
3596                .build_unchecked()
3597        })
3598    }
3599}
3600
3601impl StructuralDecodeArrayTask for StructuralCompositeDecodeArrayTask {
3602    fn decode(self: Box<Self>) -> Result<DecodedArray> {
3603        let mut arrays = Vec::with_capacity(self.tasks.len());
3604        let mut unravelers = Vec::with_capacity(self.tasks.len());
3605        let mut data_size = 0u64;
3606        for task in self.tasks {
3607            let decoded = task.decode()?;
3608            data_size += decoded.data.data_size();
3609            unravelers.push(decoded.repdef);
3610
3611            let array = make_array(
3612                decoded
3613                    .data
3614                    .into_arrow(self.data_type.clone(), self.should_validate)?,
3615            );
3616
3617            arrays.push(array);
3618        }
3619        let array_refs = arrays.iter().map(|arr| arr.as_ref()).collect::<Vec<_>>();
3620        let array = arrow_select::concat::concat(&array_refs)?;
3621        let mut repdef = CompositeRepDefUnraveler::new(unravelers);
3622
3623        let array = Self::restore_validity(array, &mut repdef);
3624
3625        Ok(DecodedArray {
3626            array,
3627            repdef,
3628            data_size,
3629        })
3630    }
3631}
3632
3633#[derive(Debug)]
3634pub struct StructuralPrimitiveFieldDecoder {
3635    field: Arc<ArrowField>,
3636    page_decoders: VecDeque<Box<dyn StructuralPageDecoder>>,
3637    should_validate: bool,
3638    rows_drained_in_current: u64,
3639}
3640
3641impl StructuralPrimitiveFieldDecoder {
3642    pub fn new(field: &Arc<ArrowField>, should_validate: bool) -> Self {
3643        Self {
3644            field: field.clone(),
3645            page_decoders: VecDeque::new(),
3646            should_validate,
3647            rows_drained_in_current: 0,
3648        }
3649    }
3650}
3651
3652impl StructuralFieldDecoder for StructuralPrimitiveFieldDecoder {
3653    fn accept_page(&mut self, child: LoadedPageShard) -> Result<()> {
3654        assert!(child.path.is_empty());
3655        self.page_decoders.push_back(child.decoder);
3656        Ok(())
3657    }
3658
3659    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn StructuralDecodeArrayTask>> {
3660        let mut remaining = num_rows;
3661        let mut tasks = Vec::new();
3662        while remaining > 0 {
3663            let queued_pages = self.page_decoders.len();
3664            let Some(cur_page) = self.page_decoders.front_mut() else {
3665                return Err(Error::internal(format!(
3666                    "Primitive decoder missing page decoder while draining field '{}' (data_type={:?}, requested_rows={}, remaining_rows={}, rows_drained_in_current={}, queued_pages={})",
3667                    self.field.name(),
3668                    self.field.data_type(),
3669                    num_rows,
3670                    remaining,
3671                    self.rows_drained_in_current,
3672                    queued_pages
3673                )));
3674            };
3675            let num_in_page = cur_page.num_rows() - self.rows_drained_in_current;
3676            let to_take = num_in_page.min(remaining);
3677
3678            let task = cur_page.drain(to_take)?;
3679            tasks.push(task);
3680
3681            if to_take == num_in_page {
3682                self.page_decoders.pop_front();
3683                self.rows_drained_in_current = 0;
3684            } else {
3685                self.rows_drained_in_current += to_take;
3686            }
3687
3688            remaining -= to_take;
3689        }
3690        Ok(Box::new(StructuralCompositeDecodeArrayTask {
3691            tasks,
3692            should_validate: self.should_validate,
3693            data_type: self.field.data_type().clone(),
3694        }))
3695    }
3696
3697    fn data_type(&self) -> &DataType {
3698        self.field.data_type()
3699    }
3700}
3701
3702/// The serialized representation of full-zip data
3703struct SerializedFullZip {
3704    /// The zipped values buffer
3705    values: LanceBuffer,
3706    /// The repetition index (only present if there is repetition)
3707    repetition_index: Option<LanceBuffer>,
3708}
3709
3710// We align and pad mini-blocks to 8 byte boundaries for two reasons.  First,
3711// to allow us to store a chunk size in 12 bits.
3712//
3713// If we directly record the size in bytes with 12 bits we would be limited to
3714// 4KiB which is too small.  Since we know each mini-block consists of 8 byte
3715// words we can store the # of words instead which gives us 32KiB.
3716//
3717// Second, each chunk in a mini-block is aligned to 8 bytes.  This allows multi-byte
3718// values like offsets to be stored in a mini-block and safely read back out.  It also
3719// helps ensure zero-copy reads in cases where zero-copy is possible (e.g. no decoding
3720// needed).
3721//
3722// Note: by "aligned to 8 bytes" we mean BOTH "aligned to 8 bytes from the start of
3723// the page" and "aligned to 8 bytes from the start of the file."
3724const MINIBLOCK_ALIGNMENT: usize = 8;
3725
3726/// An encoder for primitive (leaf) arrays
3727///
3728/// This encoder is fairly complicated and follows a number of paths depending
3729/// on the data.
3730///
3731/// First, we convert the validity & offsets information into repetition and
3732/// definition levels.  Then we compress the data itself into a single buffer.
3733///
3734/// If the data is narrow then we encode the data in small chunks (each chunk
3735/// should be a few disk sectors and contains a buffer of repetition, a buffer
3736/// of definition, and a buffer of value data).  This approach is called
3737/// "mini-block".  These mini-blocks are stored into a single data buffer.
3738///
3739/// If the data is wide then we zip together the repetition and definition value
3740/// with the value data into a single buffer.  This approach is called "zipped".
3741///
3742/// If there is any repetition information then we create a repetition index
3743///
3744/// In addition, the compression process may create zero or more metadata buffers.
3745/// For example, a dictionary compression will create dictionary metadata.  Any
3746/// mini-block approach has a metadata buffer of block sizes.  This metadata is
3747/// stored in a separate buffer on disk and read at initialization time.
3748///
3749/// TODO: We should concatenate metadata buffers from all pages into a single buffer
3750/// at (roughly) the end of the file so there is, at most, one read per column of
3751/// metadata per file.
3752pub struct PrimitiveStructuralEncoder {
3753    // Accumulates arrays until we have enough data to justify a disk page
3754    accumulation_queue: AccumulationQueue,
3755
3756    keep_original_array: bool,
3757    support_large_chunk: bool,
3758    accumulated_repdefs: Vec<RepDefBuilder>,
3759    // The compression strategy we will use to compress the data
3760    compression_strategy: Arc<dyn CompressionStrategy>,
3761    column_index: u32,
3762    field: Field,
3763    encoding_metadata: Arc<HashMap<String, String>>,
3764    version: LanceFileVersion,
3765}
3766
3767struct CompressedLevelsChunk {
3768    data: LanceBuffer,
3769    num_levels: u16,
3770}
3771
3772struct CompressedLevels {
3773    data: Vec<CompressedLevelsChunk>,
3774    compression: CompressiveEncoding,
3775    rep_index: Option<LanceBuffer>,
3776}
3777
3778struct SerializedMiniBlockPage {
3779    num_buffers: u64,
3780    data: LanceBuffer,
3781    metadata: LanceBuffer,
3782}
3783
3784#[derive(Debug, Clone, Copy)]
3785struct DictEncodingBudget {
3786    max_dict_entries: u32,
3787    max_encoded_size: usize,
3788}
3789
3790// A primitive page after applying the dense mini-block rep/def budget.
3791struct PrimitivePageData {
3792    // Arrow leaf arrays that contain this page's visible values.
3793    arrays: Vec<ArrayRef>,
3794    // Repetition / definition levels aligned to this page.
3795    repdef: SerializedRepDefs,
3796    // Top-level row number of the first row in this page.
3797    row_number: u64,
3798    // Number of top-level rows in this page.
3799    num_rows: u64,
3800    // Present when one top-level row is too large for one mini-block rep/def page.
3801    single_row_miniblock_repdef_levels: Option<u64>,
3802}
3803
3804// Immutable encoder state shared by per-page encode tasks.
3805//
3806// Cloning this only clones Arc-backed configuration and field metadata.  Page data
3807// stays in PrimitivePageData and is moved into exactly one task.
3808#[derive(Clone)]
3809struct PrimitiveEncodeContext {
3810    // Column being encoded.
3811    column_idx: u32,
3812    // Logical field metadata for compression/layout selection.
3813    field: Field,
3814    // Compression strategy shared across pages.
3815    compression_strategy: Arc<dyn CompressionStrategy>,
3816    // Field-level encoding metadata such as structural encoding overrides.
3817    encoding_metadata: Arc<HashMap<String, String>>,
3818    // Whether miniblock chunks may use the v2.2 large-chunk metadata.
3819    support_large_chunk: bool,
3820    // Lance file version selected by the writer.
3821    version: LanceFileVersion,
3822    // True when the only rep/def information is simple nullable validity.
3823    is_simple_validity: bool,
3824    // True when the field has any non-empty rep/def information.
3825    has_repdef_info: bool,
3826}
3827
3828impl PrimitiveStructuralEncoder {
3829    pub fn try_new(
3830        options: &EncodingOptions,
3831        compression_strategy: Arc<dyn CompressionStrategy>,
3832        column_index: u32,
3833        field: Field,
3834        encoding_metadata: Arc<HashMap<String, String>>,
3835    ) -> Result<Self> {
3836        Ok(Self {
3837            accumulation_queue: AccumulationQueue::new(
3838                options.cache_bytes_per_column,
3839                column_index,
3840                options.keep_original_array,
3841            ),
3842            support_large_chunk: options.support_large_chunk(),
3843            keep_original_array: options.keep_original_array,
3844            accumulated_repdefs: Vec::new(),
3845            column_index,
3846            compression_strategy,
3847            field,
3848            encoding_metadata,
3849            version: options.version,
3850        })
3851    }
3852
3853    // TODO: This is a heuristic we may need to tune at some point
3854    //
3855    // As data gets narrow then the "zipping" process gets too expensive
3856    //   and we prefer mini-block
3857    // As data gets wide then the # of values per block shrinks (very wide)
3858    //   data doesn't even fit in a mini-block and the block overhead gets
3859    //   too large and we prefer zipped.
3860    fn is_narrow(data_block: &DataBlock) -> bool {
3861        const MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE: u64 = 256;
3862
3863        if let Some(max_len_array) = data_block.get_stat(Stat::MaxLength) {
3864            let max_len_array = max_len_array
3865                .as_any()
3866                .downcast_ref::<PrimitiveArray<UInt64Type>>()
3867                .unwrap();
3868            if max_len_array.value(0) < MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE {
3869                return true;
3870            }
3871        }
3872        false
3873    }
3874
3875    fn prefers_miniblock(
3876        data_block: &DataBlock,
3877        encoding_metadata: &HashMap<String, String>,
3878    ) -> bool {
3879        // If the user specifically requested miniblock then use it
3880        if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
3881            return user_requested.to_lowercase() == STRUCTURAL_ENCODING_MINIBLOCK;
3882        }
3883        // Otherwise only use miniblock if it is narrow
3884        Self::is_narrow(data_block)
3885    }
3886
3887    fn prefers_fullzip(encoding_metadata: &HashMap<String, String>) -> bool {
3888        // Fullzip is the backup option so the only reason we wouldn't use it is if the
3889        // user specifically requested not to use it (in which case we're probably going
3890        // to emit an error)
3891        if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
3892            return user_requested.to_lowercase() == STRUCTURAL_ENCODING_FULLZIP;
3893        }
3894        true
3895    }
3896
3897    // Converts value data, repetition levels, and definition levels into a single
3898    // buffer of mini-blocks.  In addition, creates a buffer of mini-block metadata
3899    // which tells us the size of each block.  Finally, if repetition is present then
3900    // we also create a buffer for the repetition index.
3901    //
3902    // Each chunk is serialized as:
3903    // | num_bufs (1 byte) | buf_lens (2 bytes per buffer) | P | buf0 | P | buf1 | ... | bufN | P |
3904    //
3905    // P - Padding inserted to ensure each buffer is 8-byte aligned and the buffer size is a multiple
3906    //     of 8 bytes (so that the next chunk is 8-byte aligned).
3907    //
3908    // Each block has a u16 word of metadata.  The upper 12 bits contain the
3909    // # of 8-byte words in the block (if the block does not fill the final word
3910    // then up to 7 bytes of padding are added).  The lower 4 bits describe the log_2
3911    // number of values (e.g. if there are 1024 then the lower 4 bits will be
3912    // 0xA)  All blocks except the last must have power-of-two number of values.
3913    // This not only makes metadata smaller but it makes decoding easier since
3914    // batch sizes are typically a power of 2.  4 bits would allow us to express
3915    // up to 32Ki values.
3916    //
3917    // This means blocks can have 1 to 32Ki values and 8 - 32Ki bytes.
3918    //
3919    // All metadata words are serialized (as little endian) into a single buffer
3920    // of metadata values.
3921    //
3922    // If there is repetition then we also create a repetition index.  This is a
3923    // single buffer of integer vectors (stored in row major order).  There is one
3924    // entry for each chunk.  The size of the vector is based on the depth of random
3925    // access we want to support.
3926    //
3927    // A vector of size 2 is the minimum and will support row-based random access (e.g.
3928    // "take the 57th row").  A vector of size 3 will support 1 level of nested access
3929    // (e.g. "take the 3rd item in the 57th row").  A vector of size 4 will support 2
3930    // levels of nested access and so on.
3931    //
3932    // The first number in the vector is the number of top-level rows that complete in
3933    // the chunk.  The second number is the number of second-level rows that complete
3934    // after the final top-level row completed (or beginning of the chunk if no top-level
3935    // row completes in the chunk).  And so on.  The final number in the vector is always
3936    // the number of leftover items not covered by earlier entries in the vector.
3937    //
3938    // Currently we are limited to 0 levels of nested access but that will change in the
3939    // future.
3940    //
3941    // The repetition index and the chunk metadata are read at initialization time and
3942    // cached in memory.
3943    fn serialize_miniblocks(
3944        miniblocks: MiniBlockCompressed,
3945        rep: Option<Vec<CompressedLevelsChunk>>,
3946        def: Option<Vec<CompressedLevelsChunk>>,
3947        support_large_chunk: bool,
3948    ) -> Result<SerializedMiniBlockPage> {
3949        let bytes_rep = rep
3950            .as_ref()
3951            .map(|rep| rep.iter().map(|r| r.data.len()).sum::<usize>())
3952            .unwrap_or(0);
3953        let bytes_def = def
3954            .as_ref()
3955            .map(|def| def.iter().map(|d| d.data.len()).sum::<usize>())
3956            .unwrap_or(0);
3957        let bytes_data = miniblocks.data.iter().map(|d| d.len()).sum::<usize>();
3958        let mut num_buffers = miniblocks.data.len();
3959        if rep.is_some() {
3960            num_buffers += 1;
3961        }
3962        if def.is_some() {
3963            num_buffers += 1;
3964        }
3965        // 2 bytes for the length of each buffer and up to 7 bytes of padding per buffer
3966        let max_extra = 9 * num_buffers;
3967        let mut data_buffer = Vec::with_capacity(bytes_rep + bytes_def + bytes_data + max_extra);
3968        let chunk_size_bytes = if support_large_chunk { 4 } else { 2 };
3969        let mut meta_buffer = Vec::with_capacity(miniblocks.chunks.len() * chunk_size_bytes);
3970
3971        let mut rep_iter = rep.map(|r| r.into_iter());
3972        let mut def_iter = def.map(|d| d.into_iter());
3973
3974        let mut buffer_offsets = vec![0; miniblocks.data.len()];
3975        for chunk in miniblocks.chunks {
3976            let start_pos = data_buffer.len();
3977            // Start of chunk should be aligned
3978            debug_assert_eq!(start_pos % MINIBLOCK_ALIGNMENT, 0);
3979
3980            let rep = rep_iter.as_mut().map(|r| r.next().unwrap());
3981            let def = def_iter.as_mut().map(|d| d.next().unwrap());
3982
3983            // Write the number of levels, or 0 if there is no rep/def
3984            let num_levels = rep
3985                .as_ref()
3986                .map(|r| r.num_levels)
3987                .unwrap_or(def.as_ref().map(|d| d.num_levels).unwrap_or(0));
3988            data_buffer.extend_from_slice(&num_levels.to_le_bytes());
3989
3990            // Write the buffer lengths
3991            if let Some(rep) = rep.as_ref() {
3992                let bytes_rep = u16::try_from(rep.data.len()).map_err(|_| {
3993                    Error::internal(format!(
3994                        "Repetition buffer size ({} bytes) too large",
3995                        rep.data.len()
3996                    ))
3997                })?;
3998                data_buffer.extend_from_slice(&bytes_rep.to_le_bytes());
3999            }
4000            if let Some(def) = def.as_ref() {
4001                let bytes_def = u16::try_from(def.data.len()).map_err(|_| {
4002                    Error::internal(format!(
4003                        "Definition buffer size ({} bytes) too large",
4004                        def.data.len()
4005                    ))
4006                })?;
4007                data_buffer.extend_from_slice(&bytes_def.to_le_bytes());
4008            }
4009
4010            if support_large_chunk {
4011                for &buffer_size in &chunk.buffer_sizes {
4012                    data_buffer.extend_from_slice(&buffer_size.to_le_bytes());
4013                }
4014            } else {
4015                for &buffer_size in &chunk.buffer_sizes {
4016                    let buffer_size = u16::try_from(buffer_size).map_err(|_| {
4017                        Error::internal(format!(
4018                            "Mini-block buffer size ({} bytes) too large for 16-bit metadata",
4019                            buffer_size
4020                        ))
4021                    })?;
4022                    data_buffer.extend_from_slice(&buffer_size.to_le_bytes());
4023                }
4024            }
4025
4026            // Pad
4027            let add_padding = |data_buffer: &mut Vec<u8>| {
4028                let pad = pad_bytes::<MINIBLOCK_ALIGNMENT>(data_buffer.len());
4029                data_buffer.extend(iter::repeat_n(FILL_BYTE, pad));
4030            };
4031            add_padding(&mut data_buffer);
4032
4033            // Write the buffers themselves
4034            if let Some(rep) = rep.as_ref() {
4035                data_buffer.extend_from_slice(&rep.data);
4036                add_padding(&mut data_buffer);
4037            }
4038            if let Some(def) = def.as_ref() {
4039                data_buffer.extend_from_slice(&def.data);
4040                add_padding(&mut data_buffer);
4041            }
4042            for (buffer_size, (buffer, buffer_offset)) in chunk
4043                .buffer_sizes
4044                .iter()
4045                .zip(miniblocks.data.iter().zip(buffer_offsets.iter_mut()))
4046            {
4047                let start = *buffer_offset;
4048                let end = start + *buffer_size as usize;
4049                *buffer_offset += *buffer_size as usize;
4050                data_buffer.extend_from_slice(&buffer[start..end]);
4051                add_padding(&mut data_buffer);
4052            }
4053
4054            let chunk_bytes = data_buffer.len() - start_pos;
4055            let max_chunk_size = if support_large_chunk {
4056                1_u64 << 31 // 28 bits of 8-byte words in u32 metadata
4057            } else {
4058                32 * 1024 // 32KiB limit with u16 metadata
4059            };
4060            if chunk_bytes == 0 || chunk_bytes as u64 > max_chunk_size {
4061                return Err(Error::internal(format!(
4062                    "Mini-block chunk size {} bytes exceeds the {} byte metadata limit",
4063                    chunk_bytes, max_chunk_size
4064                )));
4065            }
4066            if chunk_bytes % MINIBLOCK_ALIGNMENT != 0 {
4067                return Err(Error::internal(format!(
4068                    "Mini-block chunk size {} bytes is not aligned to {} bytes",
4069                    chunk_bytes, MINIBLOCK_ALIGNMENT
4070                )));
4071            }
4072            if chunk.log_num_values > 15 {
4073                return Err(Error::internal(format!(
4074                    "Mini-block log_num_values {} exceeds the 4-bit metadata limit",
4075                    chunk.log_num_values
4076                )));
4077            }
4078            // We subtract 1 here from chunk_bytes because we want to be able to express
4079            // a size of 32KiB and not (32Ki - 8)B which is what we'd get otherwise with
4080            // 0xFFF
4081            let divided_bytes = chunk_bytes / MINIBLOCK_ALIGNMENT;
4082            let divided_bytes_minus_one = (divided_bytes - 1) as u64;
4083
4084            let metadata = (divided_bytes_minus_one << 4) | chunk.log_num_values as u64;
4085            if support_large_chunk {
4086                meta_buffer.extend_from_slice(&(metadata as u32).to_le_bytes());
4087            } else {
4088                meta_buffer.extend_from_slice(&(metadata as u16).to_le_bytes());
4089            }
4090        }
4091
4092        let data_buffer = LanceBuffer::from(data_buffer);
4093        let metadata_buffer = LanceBuffer::from(meta_buffer);
4094
4095        Ok(SerializedMiniBlockPage {
4096            num_buffers: miniblocks.data.len() as u64,
4097            data: data_buffer,
4098            metadata: metadata_buffer,
4099        })
4100    }
4101
4102    /// Compresses a buffer of levels into chunks
4103    ///
4104    /// If these are repetition levels then we also calculate the repetition index here (that
4105    /// is the third return value)
4106    fn compress_levels(
4107        mut levels: RepDefSlicer<'_>,
4108        num_elements: u64,
4109        compression_strategy: &dyn CompressionStrategy,
4110        chunks: &[MiniBlockChunk],
4111        // This will be 0 if we are compressing def levels
4112        max_rep: u16,
4113    ) -> Result<CompressedLevels> {
4114        let mut rep_index = if max_rep > 0 {
4115            Vec::with_capacity(chunks.len())
4116        } else {
4117            vec![]
4118        };
4119        // Make the levels into a FixedWidth data block
4120        let num_levels = levels.num_levels() as u64;
4121        let levels_buf = levels.all_levels().clone();
4122
4123        let mut fixed_width_block = FixedWidthDataBlock {
4124            data: levels_buf,
4125            bits_per_value: 16,
4126            num_values: num_levels,
4127            block_info: BlockInfo::new(),
4128        };
4129        // Compute statistics to enable optimal compression for rep/def levels
4130        fixed_width_block.compute_stat();
4131
4132        let levels_block = DataBlock::FixedWidth(fixed_width_block);
4133        let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
4134        // Pick a block compressor
4135        let (compressor, compressor_desc) =
4136            compression_strategy.create_block_compressor(&levels_field, &levels_block)?;
4137        // Compress blocks of levels (sized according to the chunks)
4138        let mut level_chunks = Vec::with_capacity(chunks.len());
4139        let mut values_counter = 0;
4140        for (chunk_idx, chunk) in chunks.iter().enumerate() {
4141            let chunk_num_values = chunk.num_values(values_counter, num_elements);
4142            debug_assert!(chunk_num_values > 0);
4143            values_counter += chunk_num_values;
4144            let chunk_levels = if chunk_idx < chunks.len() - 1 {
4145                levels.slice_next(chunk_num_values as usize)
4146            } else {
4147                levels.slice_rest()
4148            };
4149            let num_chunk_levels = (chunk_levels.len() / 2) as u64;
4150            if max_rep > 0 {
4151                // If max_rep > 0 then we are working with rep levels and we need
4152                // to calculate the repetition index.  The repetition index for a
4153                // chunk is currently 2 values (in the future it may be more).
4154                //
4155                // The first value is the number of rows that _finish_ in the
4156                // chunk.
4157                //
4158                // The second value is the number of "leftovers" after the last
4159                // finished row in the chunk.
4160                let rep_values = chunk_levels.borrow_to_typed_slice::<u16>();
4161                let rep_values = rep_values.as_ref();
4162
4163                // We skip 1 here because a max_rep at spot 0 doesn't count as a finished list (we
4164                // will count it in the previous chunk)
4165                let mut num_rows = rep_values.iter().skip(1).filter(|v| **v == max_rep).count();
4166                let num_leftovers = if chunk_idx < chunks.len() - 1 {
4167                    rep_values
4168                        .iter()
4169                        .rev()
4170                        .position(|v| *v == max_rep)
4171                        // # of leftovers includes the max_rep spot
4172                        .map(|pos| pos + 1)
4173                        .unwrap_or(rep_values.len())
4174                } else {
4175                    // Last chunk can't have leftovers
4176                    0
4177                };
4178
4179                if chunk_idx != 0 && rep_values.first() == Some(&max_rep) {
4180                    // This chunk starts with a new row and so, if we thought we had leftovers
4181                    // in the previous chunk, we were mistaken
4182                    // TODO: Can use unchecked here
4183                    let rep_len = rep_index.len();
4184                    if rep_index[rep_len - 1] != 0 {
4185                        // We thought we had leftovers but that was actually a full row
4186                        rep_index[rep_len - 2] += 1;
4187                        rep_index[rep_len - 1] = 0;
4188                    }
4189                }
4190
4191                if chunk_idx == chunks.len() - 1 {
4192                    // The final list
4193                    num_rows += 1;
4194                }
4195                rep_index.push(num_rows as u64);
4196                rep_index.push(num_leftovers as u64);
4197            }
4198            let mut chunk_fixed_width = FixedWidthDataBlock {
4199                data: chunk_levels,
4200                bits_per_value: 16,
4201                num_values: num_chunk_levels,
4202                block_info: BlockInfo::new(),
4203            };
4204            chunk_fixed_width.compute_stat();
4205            let chunk_levels_block = DataBlock::FixedWidth(chunk_fixed_width);
4206            let compressed_levels = compressor.compress(chunk_levels_block)?;
4207            let num_levels = u16::try_from(num_chunk_levels).map_err(|_| {
4208                Error::invalid_input_source(
4209                    format!(
4210                        "Mini-block cannot encode {} rep/def levels in one chunk. \
4211                         This usually means a top-level row contains too much nested structure \
4212                         for the current layout.",
4213                        num_chunk_levels
4214                    )
4215                    .into(),
4216                )
4217            })?;
4218            level_chunks.push(CompressedLevelsChunk {
4219                data: compressed_levels,
4220                num_levels,
4221            });
4222        }
4223        debug_assert_eq!(levels.num_levels_remaining(), 0);
4224        let rep_index = if rep_index.is_empty() {
4225            None
4226        } else {
4227            Some(LanceBuffer::reinterpret_vec(rep_index))
4228        };
4229        Ok(CompressedLevels {
4230            data: level_chunks,
4231            compression: compressor_desc,
4232            rep_index,
4233        })
4234    }
4235
4236    fn encode_simple_all_null(
4237        column_idx: u32,
4238        num_rows: u64,
4239        row_number: u64,
4240    ) -> Result<EncodedPage> {
4241        let description =
4242            ProtobufUtils21::constant_layout(&[DefinitionInterpretation::NullableItem], None);
4243        Ok(EncodedPage {
4244            column_idx,
4245            data: vec![],
4246            description: PageEncoding::Structural(description),
4247            num_rows,
4248            row_number,
4249        })
4250    }
4251
4252    fn encode_complex_all_null_vals(
4253        data: &Arc<[u16]>,
4254        compression_strategy: &dyn CompressionStrategy,
4255    ) -> Result<(LanceBuffer, pb21::CompressiveEncoding)> {
4256        let buffer = LanceBuffer::reinterpret_slice(data.clone());
4257        let mut fixed_width_block = FixedWidthDataBlock {
4258            data: buffer,
4259            bits_per_value: 16,
4260            num_values: data.len() as u64,
4261            block_info: BlockInfo::new(),
4262        };
4263        fixed_width_block.compute_stat();
4264
4265        let levels_block = DataBlock::FixedWidth(fixed_width_block);
4266        let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
4267        let (compressor, encoding) =
4268            compression_strategy.create_block_compressor(&levels_field, &levels_block)?;
4269        let compressed_buffer = compressor.compress(levels_block)?;
4270        Ok((compressed_buffer, encoding))
4271    }
4272
4273    // Encodes a page where all values are null but we have rep/def
4274    // information that we need to store (e.g. to distinguish between
4275    // different kinds of null)
4276    fn encode_complex_all_null(
4277        column_idx: u32,
4278        repdef: crate::repdef::SerializedRepDefs,
4279        row_number: u64,
4280        num_rows: u64,
4281        version: LanceFileVersion,
4282        compression_strategy: &dyn CompressionStrategy,
4283    ) -> Result<EncodedPage> {
4284        if version.resolve() < LanceFileVersion::V2_2 {
4285            let rep_bytes = if let Some(rep) = repdef.repetition_levels.as_ref() {
4286                LanceBuffer::reinterpret_slice(rep.clone())
4287            } else {
4288                LanceBuffer::empty()
4289            };
4290
4291            let def_bytes = if let Some(def) = repdef.definition_levels.as_ref() {
4292                LanceBuffer::reinterpret_slice(def.clone())
4293            } else {
4294                LanceBuffer::empty()
4295            };
4296
4297            let description = ProtobufUtils21::constant_layout(&repdef.def_meaning, None);
4298            return Ok(EncodedPage {
4299                column_idx,
4300                data: vec![rep_bytes, def_bytes],
4301                description: PageEncoding::Structural(description),
4302                num_rows,
4303                row_number,
4304            });
4305        }
4306
4307        let (rep_bytes, rep_encoding, num_rep_values) = if let Some(rep) =
4308            repdef.repetition_levels.as_ref()
4309        {
4310            let num_values = rep.len() as u64;
4311            let (buffer, encoding) = Self::encode_complex_all_null_vals(rep, compression_strategy)?;
4312            (buffer, Some(encoding), num_values)
4313        } else {
4314            (LanceBuffer::empty(), None, 0)
4315        };
4316
4317        let (def_bytes, def_encoding, num_def_values) = if let Some(def) =
4318            repdef.definition_levels.as_ref()
4319        {
4320            let num_values = def.len() as u64;
4321            let (buffer, encoding) = Self::encode_complex_all_null_vals(def, compression_strategy)?;
4322            (buffer, Some(encoding), num_values)
4323        } else {
4324            (LanceBuffer::empty(), None, 0)
4325        };
4326
4327        let description = ProtobufUtils21::compressed_all_null_constant_layout(
4328            &repdef.def_meaning,
4329            rep_encoding,
4330            def_encoding,
4331            num_rep_values,
4332            num_def_values,
4333        );
4334        Ok(EncodedPage {
4335            column_idx,
4336            data: vec![rep_bytes, def_bytes],
4337            description: PageEncoding::Structural(description),
4338            num_rows,
4339            row_number,
4340        })
4341    }
4342
4343    fn leaf_validity(
4344        repdef: &crate::repdef::SerializedRepDefs,
4345        num_values: usize,
4346    ) -> Result<Option<BooleanBuffer>> {
4347        let rep = repdef
4348            .repetition_levels
4349            .as_ref()
4350            .map(|rep| rep.as_ref().to_vec());
4351        let def = repdef
4352            .definition_levels
4353            .as_ref()
4354            .map(|def| def.as_ref().to_vec());
4355        let mut unraveler = RepDefUnraveler::new(
4356            rep,
4357            def,
4358            repdef.def_meaning.clone().into(),
4359            num_values as u64,
4360        );
4361        if unraveler.is_all_valid() {
4362            return Ok(None);
4363        }
4364        let mut validity = BooleanBufferBuilder::new(num_values);
4365        unraveler.unravel_validity(&mut validity);
4366        Ok(Some(validity.finish()))
4367    }
4368
4369    fn is_constant_values(
4370        arrays: &[ArrayRef],
4371        scalar: &ArrayRef,
4372        validity: Option<&BooleanBuffer>,
4373    ) -> Result<bool> {
4374        debug_assert_eq!(scalar.len(), 1);
4375        debug_assert_eq!(scalar.null_count(), 0);
4376
4377        match scalar.data_type() {
4378            DataType::Boolean => {
4379                let mut global_idx = 0usize;
4380                let scalar_val = scalar.as_boolean().value(0);
4381                for arr in arrays {
4382                    let bool_arr = arr.as_boolean();
4383                    for i in 0..arr.len() {
4384                        let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4385                        global_idx += 1;
4386                        if !is_valid {
4387                            continue;
4388                        }
4389                        if bool_arr.value(i) != scalar_val {
4390                            return Ok(false);
4391                        }
4392                    }
4393                }
4394                Ok(true)
4395            }
4396            DataType::Utf8 => Self::is_constant_utf8::<i32>(arrays, scalar, validity),
4397            DataType::LargeUtf8 => Self::is_constant_utf8::<i64>(arrays, scalar, validity),
4398            DataType::Binary => Self::is_constant_binary::<i32>(arrays, scalar, validity),
4399            DataType::LargeBinary => Self::is_constant_binary::<i64>(arrays, scalar, validity),
4400            data_type => {
4401                let mut global_idx = 0usize;
4402                let Some(byte_width) = data_type.byte_width_opt() else {
4403                    return Ok(false);
4404                };
4405                let scalar_data = scalar.to_data();
4406                if scalar_data.buffers().len() != 1 || !scalar_data.child_data().is_empty() {
4407                    return Ok(false);
4408                }
4409                let scalar_bytes = scalar_data.buffers()[0].as_slice();
4410                if scalar_bytes.len() != byte_width {
4411                    return Ok(false);
4412                }
4413
4414                for arr in arrays {
4415                    let data = arr.to_data();
4416                    if data.buffers().is_empty() {
4417                        return Ok(false);
4418                    }
4419                    let buf = data.buffers()[0].as_slice();
4420                    let base = data.offset();
4421                    for i in 0..arr.len() {
4422                        let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4423                        global_idx += 1;
4424                        if !is_valid {
4425                            continue;
4426                        }
4427                        let start = (base + i) * byte_width;
4428                        if buf[start..start + byte_width] != scalar_bytes[..] {
4429                            return Ok(false);
4430                        }
4431                    }
4432                }
4433                Ok(true)
4434            }
4435        }
4436    }
4437
4438    fn is_constant_utf8<O: arrow_array::OffsetSizeTrait>(
4439        arrays: &[ArrayRef],
4440        scalar: &ArrayRef,
4441        validity: Option<&BooleanBuffer>,
4442    ) -> Result<bool> {
4443        debug_assert_eq!(scalar.len(), 1);
4444        let scalar_val = scalar.as_string::<O>().value(0).as_bytes();
4445        let mut global_idx = 0usize;
4446        for arr in arrays {
4447            let str_arr = arr.as_string::<O>();
4448            for i in 0..arr.len() {
4449                let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4450                global_idx += 1;
4451                if !is_valid {
4452                    continue;
4453                }
4454                if str_arr.value(i).as_bytes() != scalar_val {
4455                    return Ok(false);
4456                }
4457            }
4458        }
4459        Ok(true)
4460    }
4461
4462    fn is_constant_binary<O: arrow_array::OffsetSizeTrait>(
4463        arrays: &[ArrayRef],
4464        scalar: &ArrayRef,
4465        validity: Option<&BooleanBuffer>,
4466    ) -> Result<bool> {
4467        debug_assert_eq!(scalar.len(), 1);
4468        let scalar_val = scalar.as_binary::<O>().value(0);
4469        let mut global_idx = 0usize;
4470        for arr in arrays {
4471            let bin_arr = arr.as_binary::<O>();
4472            for i in 0..arr.len() {
4473                let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4474                global_idx += 1;
4475                if !is_valid {
4476                    continue;
4477                }
4478                if bin_arr.value(i) != scalar_val {
4479                    return Ok(false);
4480                }
4481            }
4482        }
4483        Ok(true)
4484    }
4485
4486    fn find_constant_scalar(
4487        arrays: &[ArrayRef],
4488        validity: Option<&BooleanBuffer>,
4489    ) -> Result<Option<ArrayRef>> {
4490        if arrays.is_empty() {
4491            return Ok(None);
4492        }
4493
4494        let global_scalar_idx = if let Some(validity) = validity {
4495            let Some(idx) = (0..validity.len()).find(|&i| validity.value(i)) else {
4496                return Ok(None);
4497            };
4498            idx
4499        } else {
4500            0
4501        };
4502
4503        let mut idx_remaining = global_scalar_idx;
4504        let mut scalar_arr_idx = 0usize;
4505        while scalar_arr_idx < arrays.len() {
4506            let len = arrays[scalar_arr_idx].len();
4507            if idx_remaining < len {
4508                break;
4509            }
4510            idx_remaining -= len;
4511            scalar_arr_idx += 1;
4512        }
4513
4514        if scalar_arr_idx >= arrays.len() {
4515            return Ok(None);
4516        }
4517
4518        let scalar =
4519            lance_arrow::scalar::extract_scalar_value(&arrays[scalar_arr_idx], idx_remaining)?;
4520        if scalar.null_count() != 0 {
4521            return Ok(None);
4522        }
4523        if !Self::is_constant_values(arrays, &scalar, validity)? {
4524            return Ok(None);
4525        }
4526        Ok(Some(scalar))
4527    }
4528
4529    fn resolve_dict_values_compression_metadata(
4530        field_metadata: &HashMap<String, String>,
4531        env_compression: Option<String>,
4532        env_compression_level: Option<String>,
4533    ) -> HashMap<String, String> {
4534        let mut metadata = HashMap::new();
4535
4536        let compression = field_metadata
4537            .get(DICT_VALUES_COMPRESSION_META_KEY)
4538            .cloned()
4539            .or(env_compression)
4540            .unwrap_or_else(|| DEFAULT_DICT_VALUES_COMPRESSION.to_string());
4541        metadata.insert(COMPRESSION_META_KEY.to_string(), compression);
4542
4543        if let Some(compression_level) = field_metadata
4544            .get(DICT_VALUES_COMPRESSION_LEVEL_META_KEY)
4545            .cloned()
4546            .or(env_compression_level)
4547        {
4548            metadata.insert(COMPRESSION_LEVEL_META_KEY.to_string(), compression_level);
4549        }
4550
4551        metadata
4552    }
4553
4554    fn build_dict_values_compressor_field(field: &Field) -> Result<Field> {
4555        // This is an internal synthetic field used only to feed metadata into
4556        // `create_block_compressor` for dictionary values. The concrete type/name here
4557        // are not semantically meaningful; we rely on explicit metadata below to control
4558        // general compression selection for dictionary values.
4559        let mut dict_values_field = Field::new_arrow("", DataType::UInt16, false)?;
4560        dict_values_field.metadata = Self::resolve_dict_values_compression_metadata(
4561            &field.metadata,
4562            env::var(DICT_VALUES_COMPRESSION_ENV_VAR).ok(),
4563            env::var(DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR).ok(),
4564        );
4565        Ok(dict_values_field)
4566    }
4567
4568    #[allow(clippy::too_many_arguments)]
4569    fn encode_miniblock(
4570        column_idx: u32,
4571        field: &Field,
4572        compression_strategy: &dyn CompressionStrategy,
4573        data: DataBlock,
4574        repdef: crate::repdef::SerializedRepDefs,
4575        row_number: u64,
4576        dictionary_data: Option<DataBlock>,
4577        num_rows: u64,
4578        support_large_chunk: bool,
4579    ) -> Result<EncodedPage> {
4580        if let DataBlock::AllNull(_null_block) = data {
4581            // We should not be using mini-block for all-null.  There are other structural
4582            // encodings for that.
4583            unreachable!()
4584        }
4585
4586        let num_items = data.num_values();
4587
4588        let compressor = compression_strategy.create_miniblock_compressor(field, &data)?;
4589        let (compressed_data, value_encoding) = compressor.compress(data)?;
4590
4591        let max_rep = repdef.def_meaning.iter().filter(|l| l.is_list()).count() as u16;
4592
4593        let mut compressed_rep = repdef
4594            .rep_slicer()
4595            .map(|rep_slicer| {
4596                Self::compress_levels(
4597                    rep_slicer,
4598                    num_items,
4599                    compression_strategy,
4600                    &compressed_data.chunks,
4601                    max_rep,
4602                )
4603            })
4604            .transpose()?;
4605
4606        let (rep_index, rep_index_depth) =
4607            match compressed_rep.as_mut().and_then(|cr| cr.rep_index.as_mut()) {
4608                Some(rep_index) => (Some(rep_index.clone()), 1),
4609                None => (None, 0),
4610            };
4611
4612        let mut compressed_def = repdef
4613            .def_slicer()
4614            .map(|def_slicer| {
4615                Self::compress_levels(
4616                    def_slicer,
4617                    num_items,
4618                    compression_strategy,
4619                    &compressed_data.chunks,
4620                    /*max_rep=*/ 0,
4621                )
4622            })
4623            .transpose()?;
4624
4625        // TODO: Parquet sparsely encodes values here.  We could do the same but
4626        // then we won't have log2 values per chunk.  This means more metadata
4627        // and potentially more decoder asymmetry.  However, it may be worth
4628        // investigating at some point
4629
4630        let rep_data = compressed_rep
4631            .as_mut()
4632            .map(|cr| std::mem::take(&mut cr.data));
4633        let def_data = compressed_def
4634            .as_mut()
4635            .map(|cd| std::mem::take(&mut cd.data));
4636
4637        let serialized =
4638            Self::serialize_miniblocks(compressed_data, rep_data, def_data, support_large_chunk)?;
4639
4640        // Metadata, Data, Dictionary, (maybe) Repetition Index
4641        let mut data = Vec::with_capacity(4);
4642        data.push(serialized.metadata);
4643        data.push(serialized.data);
4644
4645        if let Some(dictionary_data) = dictionary_data {
4646            let num_dictionary_items = dictionary_data.num_values();
4647            let dict_values_field = Self::build_dict_values_compressor_field(field)?;
4648
4649            let (compressor, dictionary_encoding) = compression_strategy
4650                .create_block_compressor(&dict_values_field, &dictionary_data)?;
4651            let dictionary_buffer = compressor.compress(dictionary_data)?;
4652
4653            data.push(dictionary_buffer);
4654            if let Some(rep_index) = rep_index {
4655                data.push(rep_index);
4656            }
4657
4658            let description = ProtobufUtils21::miniblock_layout(
4659                compressed_rep.map(|cr| cr.compression),
4660                compressed_def.map(|cd| cd.compression),
4661                value_encoding,
4662                rep_index_depth,
4663                serialized.num_buffers,
4664                Some((dictionary_encoding, num_dictionary_items)),
4665                &repdef.def_meaning,
4666                num_items,
4667                support_large_chunk,
4668            );
4669            Ok(EncodedPage {
4670                num_rows,
4671                column_idx,
4672                data,
4673                description: PageEncoding::Structural(description),
4674                row_number,
4675            })
4676        } else {
4677            let description = ProtobufUtils21::miniblock_layout(
4678                compressed_rep.map(|cr| cr.compression),
4679                compressed_def.map(|cd| cd.compression),
4680                value_encoding,
4681                rep_index_depth,
4682                serialized.num_buffers,
4683                None,
4684                &repdef.def_meaning,
4685                num_items,
4686                support_large_chunk,
4687            );
4688
4689            if let Some(rep_index) = rep_index {
4690                let view = rep_index.borrow_to_typed_slice::<u64>();
4691                let total = view.chunks_exact(2).map(|c| c[0]).sum::<u64>();
4692                debug_assert_eq!(total, num_rows);
4693
4694                data.push(rep_index);
4695            }
4696
4697            Ok(EncodedPage {
4698                num_rows,
4699                column_idx,
4700                data,
4701                description: PageEncoding::Structural(description),
4702                row_number,
4703            })
4704        }
4705    }
4706
4707    // For fixed-size data we encode < control word | data > for each value
4708    fn serialize_full_zip_fixed(
4709        fixed: FixedWidthDataBlock,
4710        mut repdef: ControlWordIterator,
4711        num_values: u64,
4712    ) -> Result<SerializedFullZip> {
4713        if !fixed.bits_per_value.is_multiple_of(8) {
4714            return Err(Error::invalid_input_source(
4715                format!(
4716                    "Full-zip fixed-width values must be byte aligned, got {} bits per value",
4717                    fixed.bits_per_value
4718                )
4719                .into(),
4720            ));
4721        }
4722
4723        let len = fixed.data.len() + repdef.bytes_per_word() * num_values as usize;
4724        let mut zipped_data = Vec::with_capacity(len);
4725
4726        let max_rep_index_val = if repdef.has_repetition() {
4727            len as u64
4728        } else {
4729            // Setting this to 0 means we won't write a repetition index
4730            0
4731        };
4732        let mut rep_index_builder =
4733            BytepackedIntegerEncoder::with_capacity(num_values as usize + 1, max_rep_index_val);
4734
4735        let bytes_per_value = fixed.bits_per_value as usize / 8;
4736        let mut offset = 0;
4737
4738        if bytes_per_value == 0 {
4739            // No data, just dump the repdef into the buffer
4740            while let Some(control) = repdef.append_next(&mut zipped_data) {
4741                if control.is_new_row {
4742                    // We have finished a row
4743                    debug_assert!(offset <= len);
4744                    // SAFETY: We know that `start <= len`
4745                    unsafe { rep_index_builder.append(offset as u64) };
4746                }
4747                offset = zipped_data.len();
4748            }
4749        } else {
4750            // We have data, zip it with the repdef
4751            let mut data_iter = fixed.data.chunks_exact(bytes_per_value);
4752            while let Some(control) = repdef.append_next(&mut zipped_data) {
4753                if control.is_new_row {
4754                    // We have finished a row
4755                    debug_assert!(offset <= len);
4756                    // SAFETY: We know that `start <= len`
4757                    unsafe { rep_index_builder.append(offset as u64) };
4758                }
4759                if control.is_visible {
4760                    let value = data_iter.next().unwrap();
4761                    zipped_data.extend_from_slice(value);
4762                }
4763                offset = zipped_data.len();
4764            }
4765        }
4766
4767        debug_assert_eq!(zipped_data.len(), len);
4768        // Put the final value in the rep index
4769        // SAFETY: `zipped_data.len() == len`
4770        unsafe {
4771            rep_index_builder.append(zipped_data.len() as u64);
4772        }
4773
4774        let zipped_data = LanceBuffer::from(zipped_data);
4775        let rep_index = rep_index_builder.into_data();
4776        let rep_index = if rep_index.is_empty() {
4777            None
4778        } else {
4779            Some(LanceBuffer::from(rep_index))
4780        };
4781        Ok(SerializedFullZip {
4782            values: zipped_data,
4783            repetition_index: rep_index,
4784        })
4785    }
4786
4787    // For variable-size data we encode < control word | length | data > for each value
4788    //
4789    // In addition, we create a second buffer, the repetition index
4790    fn serialize_full_zip_variable(
4791        variable: VariableWidthBlock,
4792        mut repdef: ControlWordIterator,
4793        num_items: u64,
4794    ) -> Result<SerializedFullZip> {
4795        let bytes_per_offset = variable.bits_per_offset as usize / 8;
4796        if !variable.bits_per_offset.is_multiple_of(8) {
4797            return Err(Error::invalid_input_source(
4798                format!(
4799                    "Full-zip variable-width offsets must be byte aligned, got {} bits per offset",
4800                    variable.bits_per_offset
4801                )
4802                .into(),
4803            ));
4804        }
4805        let len = variable.data.len()
4806            + repdef.bytes_per_word() * num_items as usize
4807            + bytes_per_offset * variable.num_values as usize;
4808        let mut buf = Vec::with_capacity(len);
4809
4810        let max_rep_index_val = len as u64;
4811        let mut rep_index_builder =
4812            BytepackedIntegerEncoder::with_capacity(num_items as usize + 1, max_rep_index_val);
4813
4814        // TODO: byte pack the item lengths with varint encoding
4815        match bytes_per_offset {
4816            4 => {
4817                let offs = variable.offsets.borrow_to_typed_slice::<u32>();
4818                let mut rep_offset = 0;
4819                let mut windows_iter = offs.as_ref().windows(2);
4820                while let Some(control) = repdef.append_next(&mut buf) {
4821                    if control.is_new_row {
4822                        // We have finished a row
4823                        debug_assert!(rep_offset <= len);
4824                        // SAFETY: We know that `buf.len() <= len`
4825                        unsafe { rep_index_builder.append(rep_offset as u64) };
4826                    }
4827                    if control.is_visible {
4828                        let window = windows_iter.next().unwrap();
4829                        if control.is_valid_item {
4830                            buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
4831                            buf.extend_from_slice(
4832                                &variable.data[window[0] as usize..window[1] as usize],
4833                            );
4834                        }
4835                    }
4836                    rep_offset = buf.len();
4837                }
4838            }
4839            8 => {
4840                let offs = variable.offsets.borrow_to_typed_slice::<u64>();
4841                let mut rep_offset = 0;
4842                let mut windows_iter = offs.as_ref().windows(2);
4843                while let Some(control) = repdef.append_next(&mut buf) {
4844                    if control.is_new_row {
4845                        // We have finished a row
4846                        debug_assert!(rep_offset <= len);
4847                        // SAFETY: We know that `buf.len() <= len`
4848                        unsafe { rep_index_builder.append(rep_offset as u64) };
4849                    }
4850                    if control.is_visible {
4851                        let window = windows_iter.next().unwrap();
4852                        if control.is_valid_item {
4853                            buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
4854                            buf.extend_from_slice(
4855                                &variable.data[window[0] as usize..window[1] as usize],
4856                            );
4857                        }
4858                    }
4859                    rep_offset = buf.len();
4860                }
4861            }
4862            _ => {
4863                return Err(Error::invalid_input_source(
4864                    format!(
4865                        "Full-zip variable-width offsets must be 32 or 64 bits, got {} bits",
4866                        variable.bits_per_offset
4867                    )
4868                    .into(),
4869                ));
4870            }
4871        }
4872
4873        // We might have saved a few bytes by not copying lengths when the length was zero.  However,
4874        // if we are over `len` then we have a bug.
4875        debug_assert!(buf.len() <= len);
4876        // Put the final value in the rep index
4877        // SAFETY: `zipped_data.len() == len`
4878        unsafe {
4879            rep_index_builder.append(buf.len() as u64);
4880        }
4881
4882        let zipped_data = LanceBuffer::from(buf);
4883        let rep_index = rep_index_builder.into_data();
4884        debug_assert!(!rep_index.is_empty());
4885        let rep_index = Some(LanceBuffer::from(rep_index));
4886        Ok(SerializedFullZip {
4887            values: zipped_data,
4888            repetition_index: rep_index,
4889        })
4890    }
4891
4892    /// Serializes data into a single buffer according to the full-zip format which zips
4893    /// together the repetition, definition, and value data into a single buffer.
4894    fn serialize_full_zip(
4895        compressed_data: PerValueDataBlock,
4896        repdef: ControlWordIterator,
4897        num_items: u64,
4898    ) -> Result<SerializedFullZip> {
4899        match compressed_data {
4900            PerValueDataBlock::Fixed(fixed) => {
4901                Self::serialize_full_zip_fixed(fixed, repdef, num_items)
4902            }
4903            PerValueDataBlock::Variable(var) => {
4904                Self::serialize_full_zip_variable(var, repdef, num_items)
4905            }
4906        }
4907    }
4908
4909    fn expand_boolean_to_bytes(fixed: FixedWidthDataBlock) -> FixedWidthDataBlock {
4910        debug_assert_eq!(fixed.bits_per_value, 1);
4911        let num_values = fixed.num_values as usize;
4912        let bool_buf = BooleanBuffer::new(fixed.data.into_buffer(), 0, num_values);
4913        let expanded: Vec<u8> = (0..num_values).map(|i| bool_buf.value(i) as u8).collect();
4914        FixedWidthDataBlock {
4915            data: LanceBuffer::from(expanded),
4916            bits_per_value: 8,
4917            num_values: fixed.num_values,
4918            block_info: BlockInfo::new(),
4919        }
4920    }
4921
4922    fn encode_full_zip(
4923        column_idx: u32,
4924        field: &Field,
4925        compression_strategy: &dyn CompressionStrategy,
4926        data: DataBlock,
4927        repdef: crate::repdef::SerializedRepDefs,
4928        row_number: u64,
4929        num_lists: u64,
4930    ) -> Result<EncodedPage> {
4931        let max_rep = repdef
4932            .repetition_levels
4933            .as_ref()
4934            .map_or(0, |r| r.iter().max().copied().unwrap_or(0));
4935        let max_def = repdef
4936            .definition_levels
4937            .as_ref()
4938            .map_or(0, |d| d.iter().max().copied().unwrap_or(0));
4939
4940        // To handle FSL we just flatten
4941        // let data = data.flatten();
4942
4943        let (num_items, num_visible_items) =
4944            if let Some(rep_levels) = repdef.repetition_levels.as_ref() {
4945                // If there are rep levels there may be "invisible" items and we need to encode
4946                // rep_levels.len() things which might be larger than data.num_values()
4947                (rep_levels.len() as u64, data.num_values())
4948            } else {
4949                // If there are no rep levels then we encode data.num_values() things
4950                (data.num_values(), data.num_values())
4951            };
4952
4953        let max_visible_def = repdef.max_visible_level.unwrap_or(u16::MAX);
4954
4955        let repdef_iter = build_control_word_iterator(
4956            repdef.repetition_levels.as_deref(),
4957            max_rep,
4958            repdef.definition_levels.as_deref(),
4959            max_def,
4960            max_visible_def,
4961            num_items as usize,
4962        );
4963        let bits_rep = repdef_iter.bits_rep();
4964        let bits_def = repdef_iter.bits_def();
4965
4966        // Full-zip requires byte-aligned values; expand 1-bit booleans to 1 byte each.
4967        let data = match data {
4968            DataBlock::FixedWidth(fixed) if fixed.bits_per_value == 1 => {
4969                DataBlock::FixedWidth(Self::expand_boolean_to_bytes(fixed))
4970            }
4971            other => other,
4972        };
4973
4974        let compressor = compression_strategy.create_per_value(field, &data)?;
4975        let (compressed_data, value_encoding) = compressor.compress(data)?;
4976
4977        let description = match &compressed_data {
4978            PerValueDataBlock::Fixed(fixed) => ProtobufUtils21::fixed_full_zip_layout(
4979                bits_rep,
4980                bits_def,
4981                fixed.bits_per_value as u32,
4982                value_encoding,
4983                &repdef.def_meaning,
4984                num_items as u32,
4985                num_visible_items as u32,
4986            ),
4987            PerValueDataBlock::Variable(variable) => ProtobufUtils21::variable_full_zip_layout(
4988                bits_rep,
4989                bits_def,
4990                variable.bits_per_offset as u32,
4991                value_encoding,
4992                &repdef.def_meaning,
4993                num_items as u32,
4994                num_visible_items as u32,
4995            ),
4996        };
4997
4998        let zipped = Self::serialize_full_zip(compressed_data, repdef_iter, num_items)?;
4999
5000        let data = if let Some(repindex) = zipped.repetition_index {
5001            vec![zipped.values, repindex]
5002        } else {
5003            vec![zipped.values]
5004        };
5005
5006        Ok(EncodedPage {
5007            num_rows: num_lists,
5008            column_idx,
5009            data,
5010            description: PageEncoding::Structural(description),
5011            row_number,
5012        })
5013    }
5014
5015    fn should_dictionary_encode(
5016        data_block: &DataBlock,
5017        field: &Field,
5018        version: LanceFileVersion,
5019    ) -> Option<DictEncodingBudget> {
5020        const DEFAULT_SAMPLE_SIZE: usize = 4096;
5021        const DEFAULT_SAMPLE_UNIQUE_RATIO: f64 = 0.98;
5022
5023        // Since we only dictionary encode FixedWidth and VariableWidth blocks for now, we skip
5024        // estimating the size for other types.
5025        match data_block {
5026            DataBlock::FixedWidth(fixed) => {
5027                if fixed.bits_per_value == 64 && version < LanceFileVersion::V2_2 {
5028                    return None;
5029                }
5030                if fixed.bits_per_value != 64 && fixed.bits_per_value != 128 {
5031                    return None;
5032                }
5033                if fixed.bits_per_value % 8 != 0 {
5034                    return None;
5035                }
5036            }
5037            DataBlock::VariableWidth(var) => {
5038                if var.bits_per_offset != 32 && var.bits_per_offset != 64 {
5039                    return None;
5040                }
5041            }
5042            _ => return None,
5043        }
5044
5045        // Don't dictionary encode tiny arrays.
5046        let too_small = env::var("LANCE_ENCODING_DICT_TOO_SMALL")
5047            .ok()
5048            .and_then(|val| val.parse().ok())
5049            .unwrap_or(100);
5050        if data_block.num_values() < too_small {
5051            return None;
5052        }
5053
5054        let num_values = data_block.num_values();
5055
5056        // Apply divisor threshold and cap. This is intentionally conservative: the goal is to
5057        // avoid spending too much CPU trying to estimate very high cardinalities.
5058        let divisor: u64 = field
5059            .metadata
5060            .get(DICT_DIVISOR_META_KEY)
5061            .and_then(|val| val.parse().ok())
5062            .or_else(|| {
5063                env::var("LANCE_ENCODING_DICT_DIVISOR")
5064                    .ok()
5065                    .and_then(|val| val.parse().ok())
5066            })
5067            .unwrap_or(DEFAULT_DICT_DIVISOR);
5068
5069        let max_cardinality: u64 = env::var("LANCE_ENCODING_DICT_MAX_CARDINALITY")
5070            .ok()
5071            .and_then(|val| val.parse().ok())
5072            .unwrap_or(DEFAULT_DICT_MAX_CARDINALITY);
5073
5074        let threshold_cardinality = num_values
5075            .checked_div(divisor.max(1))
5076            .unwrap_or(0)
5077            .min(max_cardinality);
5078        if threshold_cardinality == 0 {
5079            return None;
5080        }
5081
5082        // Get size ratio from metadata or env var.
5083        let threshold_ratio = field
5084            .metadata
5085            .get(DICT_SIZE_RATIO_META_KEY)
5086            .and_then(|val| val.parse::<f64>().ok())
5087            .or_else(|| {
5088                env::var("LANCE_ENCODING_DICT_SIZE_RATIO")
5089                    .ok()
5090                    .and_then(|val| val.parse().ok())
5091            })
5092            .unwrap_or(DEFAULT_DICT_SIZE_RATIO);
5093
5094        if threshold_ratio <= 0.0 || threshold_ratio > 1.0 {
5095            panic!(
5096                "Invalid parameter: dict-size-ratio is {} which is not in the range (0, 1].",
5097                threshold_ratio
5098            );
5099        }
5100
5101        let data_size = data_block.data_size();
5102        if data_size == 0 {
5103            return None;
5104        }
5105
5106        let max_encoded_size = (data_size as f64 * threshold_ratio) as u64;
5107        let max_encoded_size = usize::try_from(max_encoded_size).ok()?;
5108
5109        // Avoid probing dictionary encoding on data that appears to be near-unique
5110        // or likely to exceed the dictionary budget.
5111        if let Some(sample_unique_ratio) =
5112            Self::sample_unique_ratio(data_block, DEFAULT_SAMPLE_SIZE)?
5113        {
5114            if sample_unique_ratio >= DEFAULT_SAMPLE_UNIQUE_RATIO {
5115                return None;
5116            }
5117
5118            let projected_cardinality = (sample_unique_ratio * num_values as f64).ceil() as u64;
5119            if projected_cardinality > threshold_cardinality {
5120                return None;
5121            }
5122        }
5123
5124        let max_dict_entries = u32::try_from(threshold_cardinality.min(i32::MAX as u64)).ok()?;
5125        Some(DictEncodingBudget {
5126            max_dict_entries,
5127            max_encoded_size,
5128        })
5129    }
5130
5131    /// Samples whether a page looks near-unique before attempting dictionary encoding.
5132    ///
5133    /// The probe uses deterministic block sampling (not RNG sampling), which keeps
5134    /// the check cheap and reproducible across runs. The result is only a gate for
5135    /// whether we try dictionary encoding, not a cardinality statistic.
5136    /// Returns `Some(None)` when there are too few reliable samples or the block type does not
5137    /// support dictionary encoding. Returns `None` for malformed data.
5138    fn sample_unique_ratio(data_block: &DataBlock, max_samples: usize) -> Option<Option<f64>> {
5139        use std::collections::HashSet;
5140
5141        const NUM_SAMPLE_BLOCKS: usize = 32;
5142        const MIN_RELIABLE_SAMPLES: usize = 1024;
5143
5144        let num_values = usize::try_from(data_block.num_values()).ok()?;
5145        if num_values == 0 {
5146            return Some(None);
5147        }
5148
5149        let sample_count = num_values.min(max_samples).max(1);
5150        if sample_count < MIN_RELIABLE_SAMPLES {
5151            return Some(None);
5152        }
5153
5154        let block_count = NUM_SAMPLE_BLOCKS.min(sample_count).min(num_values).max(1);
5155        let samples_per_block = (sample_count / block_count).max(1);
5156        let mut indices = Vec::with_capacity(sample_count);
5157        for block_idx in 0..block_count {
5158            let block_start = block_idx * num_values / block_count;
5159            let next_block_start = ((block_idx + 1) * num_values / block_count).min(num_values);
5160            let block_len = next_block_start.saturating_sub(block_start);
5161            let samples_in_block = samples_per_block.min(block_len);
5162            indices.extend((0..samples_in_block).map(|offset| block_start + offset));
5163        }
5164
5165        if indices.len() < MIN_RELIABLE_SAMPLES {
5166            return Some(None);
5167        }
5168
5169        let ratio = match data_block {
5170            DataBlock::FixedWidth(fixed) => match fixed.bits_per_value {
5171                64 => {
5172                    let values = fixed.data.borrow_to_typed_slice::<u64>();
5173                    let values = values.as_ref();
5174                    let mut unique: HashSet<u64> =
5175                        HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
5176                    for idx in indices.iter().copied() {
5177                        unique.insert(values.get(idx).copied()?);
5178                    }
5179                    unique.len() as f64 / indices.len() as f64
5180                }
5181                128 => {
5182                    let values = fixed.data.borrow_to_typed_slice::<u128>();
5183                    let values = values.as_ref();
5184                    let mut unique: HashSet<u128> =
5185                        HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
5186                    for idx in indices.iter().copied() {
5187                        unique.insert(values.get(idx).copied()?);
5188                    }
5189                    unique.len() as f64 / indices.len() as f64
5190                }
5191                _ => return Some(None),
5192            },
5193            DataBlock::VariableWidth(var) => {
5194                use xxhash_rust::xxh3::xxh3_64;
5195
5196                // Hash variable-width slices instead of storing borrowed slice keys.
5197                let mut unique: HashSet<u64> =
5198                    HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
5199                match var.bits_per_offset {
5200                    32 => {
5201                        let offsets_ref = var.offsets.borrow_to_typed_slice::<u32>();
5202                        let offsets: &[u32] = offsets_ref.as_ref();
5203                        for i in indices.iter().copied() {
5204                            let start = usize::try_from(*offsets.get(i)?).ok()?;
5205                            let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
5206                            if start > end || end > var.data.len() {
5207                                return None;
5208                            }
5209                            unique.insert(xxh3_64(&var.data[start..end]));
5210                        }
5211                    }
5212                    64 => {
5213                        let offsets_ref = var.offsets.borrow_to_typed_slice::<u64>();
5214                        let offsets: &[u64] = offsets_ref.as_ref();
5215                        for i in indices.iter().copied() {
5216                            let start = usize::try_from(*offsets.get(i)?).ok()?;
5217                            let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
5218                            if start > end || end > var.data.len() {
5219                                return None;
5220                            }
5221                            unique.insert(xxh3_64(&var.data[start..end]));
5222                        }
5223                    }
5224                    _ => return Some(None),
5225                }
5226                unique.len() as f64 / indices.len() as f64
5227            }
5228            _ => return Some(None),
5229        };
5230
5231        Some(Some(ratio))
5232    }
5233
5234    fn slice_repdef(repdef: &SerializedRepDefs, range: Range<usize>) -> SerializedRepDefs {
5235        let repetition_levels = repdef
5236            .repetition_levels
5237            .as_ref()
5238            .map(|levels| levels[range.clone()].to_vec());
5239        let definition_levels = repdef
5240            .definition_levels
5241            .as_ref()
5242            .map(|levels| levels[range].to_vec());
5243        SerializedRepDefs::new_with_fixed_size_list_levels(
5244            repetition_levels,
5245            definition_levels,
5246            repdef.def_meaning.clone(),
5247            repdef.has_fixed_size_list_levels(),
5248        )
5249    }
5250
5251    fn slice_arrays(
5252        arrays: &[ArrayRef],
5253        value_start: u64,
5254        num_values: u64,
5255    ) -> Result<Vec<ArrayRef>> {
5256        if num_values == 0 {
5257            return Ok(Vec::new());
5258        }
5259
5260        let mut values_to_skip = usize::try_from(value_start).map_err(|_| {
5261            Error::invalid_input(format!("Value start {} is too large", value_start))
5262        })?;
5263        let mut values_remaining = usize::try_from(num_values).map_err(|_| {
5264            Error::invalid_input(format!("Value count {} is too large", num_values))
5265        })?;
5266        let mut sliced = Vec::new();
5267
5268        for array in arrays {
5269            if values_to_skip >= array.len() {
5270                values_to_skip -= array.len();
5271                continue;
5272            }
5273
5274            let offset = values_to_skip;
5275            let len = (array.len() - offset).min(values_remaining);
5276            sliced.push(array.slice(offset, len));
5277            values_remaining -= len;
5278            values_to_skip = 0;
5279
5280            if values_remaining == 0 {
5281                break;
5282            }
5283        }
5284
5285        if values_remaining != 0 {
5286            return Err(Error::internal(format!(
5287                "Page split requested {} values starting at {}, but the page did not contain enough values",
5288                num_values, value_start
5289            )));
5290        }
5291
5292        Ok(sliced)
5293    }
5294
5295    fn split_pages_for_miniblock_repdef_budget(
5296        arrays: Vec<ArrayRef>,
5297        repdef: SerializedRepDefs,
5298        budget: MiniBlockRepDefBudget,
5299        row_number: u64,
5300        num_rows: u64,
5301    ) -> Result<Vec<PrimitivePageData>> {
5302        if budget == MiniBlockRepDefBudget::WithinBudget {
5303            return Ok(vec![PrimitivePageData {
5304                arrays,
5305                repdef,
5306                row_number,
5307                num_rows,
5308                single_row_miniblock_repdef_levels: None,
5309            }]);
5310        }
5311        if let MiniBlockRepDefBudget::SingleRowOverBudget(num_levels) = budget {
5312            return Ok(vec![PrimitivePageData {
5313                arrays,
5314                repdef,
5315                row_number,
5316                num_rows,
5317                single_row_miniblock_repdef_levels: Some(num_levels),
5318            }]);
5319        }
5320
5321        let MiniBlockRepDefBudget::RequiresPageSplit(splits) = budget else {
5322            unreachable!();
5323        };
5324
5325        let mut pages = Vec::with_capacity(splits.len());
5326        for split in splits {
5327            let arrays = Self::slice_arrays(&arrays, split.value_start, split.num_values)?;
5328            let repdef = Self::slice_repdef(&repdef, split.level_range);
5329            pages.push(PrimitivePageData {
5330                arrays,
5331                repdef,
5332                row_number: row_number + split.row_start,
5333                num_rows: split.num_rows,
5334                single_row_miniblock_repdef_levels: None,
5335            });
5336        }
5337        Ok(pages)
5338    }
5339
5340    fn encode_page(ctx: PrimitiveEncodeContext, page: PrimitivePageData) -> Result<EncodedPage> {
5341        let PrimitiveEncodeContext {
5342            column_idx,
5343            field,
5344            compression_strategy,
5345            encoding_metadata,
5346            support_large_chunk,
5347            version,
5348            is_simple_validity,
5349            has_repdef_info,
5350        } = ctx;
5351        let PrimitivePageData {
5352            arrays,
5353            repdef,
5354            row_number,
5355            num_rows,
5356            single_row_miniblock_repdef_levels,
5357        } = page;
5358        let num_values = arrays.iter().map(|arr| arr.len() as u64).sum();
5359
5360        if num_values == 0 {
5361            // This page contains only structural events, such as empty/null list rows.
5362            // The existing complex-null layout stores the rep/def stream without value buffers.
5363            log::debug!(
5364                "Encoding column {} with {} items ({} rows) using complex-null layout",
5365                column_idx,
5366                num_values,
5367                num_rows
5368            );
5369            return Self::encode_complex_all_null(
5370                column_idx,
5371                repdef,
5372                row_number,
5373                num_rows,
5374                version,
5375                compression_strategy.as_ref(),
5376            );
5377        }
5378
5379        let leaf_validity = Self::leaf_validity(&repdef, num_values as usize)?;
5380        let all_null = leaf_validity
5381            .as_ref()
5382            .map(|validity| validity.count_set_bits() == 0)
5383            .unwrap_or(false);
5384
5385        if all_null {
5386            return if is_simple_validity {
5387                log::debug!(
5388                    "Encoding column {} with {} items ({} rows) using simple-null layout",
5389                    column_idx,
5390                    num_values,
5391                    num_rows
5392                );
5393                Self::encode_simple_all_null(column_idx, num_values, row_number)
5394            } else {
5395                log::debug!(
5396                    "Encoding column {} with {} items ({} rows) using complex-null layout",
5397                    column_idx,
5398                    num_values,
5399                    num_rows
5400                );
5401                Self::encode_complex_all_null(
5402                    column_idx,
5403                    repdef,
5404                    row_number,
5405                    num_rows,
5406                    version,
5407                    compression_strategy.as_ref(),
5408                )
5409            };
5410        }
5411
5412        if let DataType::Struct(fields) = &field.data_type()
5413            && fields.is_empty()
5414        {
5415            if has_repdef_info {
5416                return Err(Error::invalid_input_source(format!("Empty structs with rep/def information are not yet supported.  The field {} is an empty struct that either has nulls or is in a list.", field.name).into()));
5417            }
5418            // This is maybe a little confusing but the reader should never look at this anyways and it
5419            // seems like overkill to invent a new layout just for "empty structs".
5420            return Self::encode_simple_all_null(column_idx, num_values, row_number);
5421        }
5422
5423        let data_block = DataBlock::from_arrays(&arrays, num_values);
5424
5425        if version.resolve() >= LanceFileVersion::V2_2
5426            && let Some(scalar) = Self::find_constant_scalar(&arrays, leaf_validity.as_ref())?
5427        {
5428            log::debug!(
5429                "Encoding column {} with {} items ({} rows) using constant layout",
5430                column_idx,
5431                num_values,
5432                num_rows
5433            );
5434            return constant::encode_constant_page(
5435                column_idx, scalar, repdef, row_number, num_rows,
5436            );
5437        }
5438
5439        if let Some(num_levels) = single_row_miniblock_repdef_levels {
5440            let requested_encoding = encoding_metadata
5441                .get(STRUCTURAL_ENCODING_META_KEY)
5442                .map(|requested| requested.to_lowercase());
5443            let fullzip_error = match &data_block {
5444                DataBlock::FixedWidth(fixed) if !fixed.bits_per_value.is_multiple_of(8) => {
5445                    Some(format!(
5446                        "Full-zip fixed-width values must be byte aligned, got {} bits per value",
5447                        fixed.bits_per_value
5448                    ))
5449                }
5450                DataBlock::VariableWidth(variable)
5451                    if !variable.bits_per_offset.is_multiple_of(8) =>
5452                {
5453                    Some(format!(
5454                        "Full-zip variable-width offsets must be byte aligned, got {} bits per offset",
5455                        variable.bits_per_offset
5456                    ))
5457                }
5458                DataBlock::VariableWidth(variable)
5459                    if variable.bits_per_offset != 32 && variable.bits_per_offset != 64 =>
5460                {
5461                    Some(format!(
5462                        "Full-zip variable-width offsets must be 32 or 64 bits, got {} bits",
5463                        variable.bits_per_offset
5464                    ))
5465                }
5466                DataBlock::Struct(struct_data_block)
5467                    if !struct_data_block.has_variable_width_child() =>
5468                {
5469                    Some(
5470                        "Full-zip packed struct requires at least one variable-width child"
5471                            .to_string(),
5472                    )
5473                }
5474                DataBlock::Dictionary(_) => {
5475                    Some("Full-zip does not encode dictionary data blocks directly".to_string())
5476                }
5477                DataBlock::FixedSizeList(fsl) => match fsl.clone().try_into_flat() {
5478                    Some(flat) if flat.bits_per_value.is_multiple_of(8) => None,
5479                    Some(flat) => Some(format!(
5480                        "Full-zip fixed-size-list values must be byte aligned after flattening, got {} bits per value",
5481                        flat.bits_per_value
5482                    )),
5483                    None => Some(
5484                        "Full-zip fixed-size-list capability requires a flat fixed-width child"
5485                            .to_string(),
5486                    ),
5487                },
5488                DataBlock::FixedWidth(_) | DataBlock::VariableWidth(_) | DataBlock::Struct(_) => {
5489                    None
5490                }
5491                other => Some(format!(
5492                    "Full-zip does not support value block type {}",
5493                    other.name()
5494                )),
5495            };
5496            match requested_encoding.as_deref() {
5497                Some(STRUCTURAL_ENCODING_FULLZIP) => {
5498                    if let Some(reason) = fullzip_error {
5499                        return Err(Error::invalid_input_source(reason.into()));
5500                    }
5501                    return Self::encode_full_zip(
5502                        column_idx,
5503                        &field,
5504                        compression_strategy.as_ref(),
5505                        data_block,
5506                        repdef,
5507                        row_number,
5508                        num_rows,
5509                    );
5510                }
5511                Some(STRUCTURAL_ENCODING_MINIBLOCK) | None => {
5512                    if requested_encoding.is_none() && fullzip_error.is_none() {
5513                        log::debug!(
5514                            "Encoding column {} with {} items using full-zip layout because mini-block cannot split the structural page",
5515                            column_idx,
5516                            num_values
5517                        );
5518                        return Self::encode_full_zip(
5519                            column_idx,
5520                            &field,
5521                            compression_strategy.as_ref(),
5522                            data_block,
5523                            repdef,
5524                            row_number,
5525                            num_rows,
5526                        );
5527                    }
5528                    return Err(Error::invalid_input_source(
5529                        format!(
5530                            "Mini-block cannot encode {} rep/def levels in one top-level row. \
5531                             This usually means the row contains too much nested structure \
5532                             for the current layout.",
5533                            num_levels
5534                        )
5535                        .into(),
5536                    ));
5537                }
5538                _ => {}
5539            }
5540        }
5541
5542        let requires_full_zip_packed_struct =
5543            if let DataBlock::Struct(ref struct_data_block) = data_block {
5544                struct_data_block.has_variable_width_child()
5545            } else {
5546                false
5547            };
5548
5549        if requires_full_zip_packed_struct {
5550            log::debug!(
5551                "Encoding column {} with {} items using full-zip packed struct layout",
5552                column_idx,
5553                num_values
5554            );
5555            return Self::encode_full_zip(
5556                column_idx,
5557                &field,
5558                compression_strategy.as_ref(),
5559                data_block,
5560                repdef,
5561                row_number,
5562                num_rows,
5563            );
5564        }
5565
5566        if let DataBlock::Dictionary(dict) = data_block {
5567            log::debug!(
5568                "Encoding column {} with {} items using dictionary encoding (already dictionary encoded)",
5569                column_idx,
5570                num_values
5571            );
5572            let (mut indices_data_block, dictionary_data_block) = dict.into_parts();
5573            // TODO: https://github.com/lancedb/lance/issues/4809
5574            // If we compute stats on dictionary_data_block => panic.
5575            // If we don't compute stats on indices_data_block => panic.
5576            // This is messy.  Don't make me call compute_stat ever.
5577            indices_data_block.compute_stat();
5578            return Self::encode_miniblock(
5579                column_idx,
5580                &field,
5581                compression_strategy.as_ref(),
5582                indices_data_block,
5583                repdef,
5584                row_number,
5585                Some(dictionary_data_block),
5586                num_rows,
5587                support_large_chunk,
5588            );
5589        }
5590
5591        // Try dictionary encoding first if applicable. If encoding aborts, fall back to the
5592        // preferred structural encoding.
5593        let dict_result = Self::should_dictionary_encode(&data_block, &field, version).and_then(|budget| {
5594                log::debug!(
5595                    "Encoding column {} with {} items using dictionary encoding (mini-block layout)",
5596                    column_idx,
5597                    num_values
5598                );
5599                dict::dictionary_encode(
5600                    &data_block,
5601                    budget.max_dict_entries,
5602                    budget.max_encoded_size,
5603                )
5604            });
5605
5606        if let Some((indices_data_block, dictionary_data_block)) = dict_result {
5607            Self::encode_miniblock(
5608                column_idx,
5609                &field,
5610                compression_strategy.as_ref(),
5611                indices_data_block,
5612                repdef,
5613                row_number,
5614                Some(dictionary_data_block),
5615                num_rows,
5616                support_large_chunk,
5617            )
5618        } else if Self::prefers_miniblock(&data_block, encoding_metadata.as_ref()) {
5619            log::debug!(
5620                "Encoding column {} with {} items using mini-block layout",
5621                column_idx,
5622                num_values
5623            );
5624            Self::encode_miniblock(
5625                column_idx,
5626                &field,
5627                compression_strategy.as_ref(),
5628                data_block,
5629                repdef,
5630                row_number,
5631                None,
5632                num_rows,
5633                support_large_chunk,
5634            )
5635        } else if Self::prefers_fullzip(encoding_metadata.as_ref()) {
5636            log::debug!(
5637                "Encoding column {} with {} items using full-zip layout",
5638                column_idx,
5639                num_values
5640            );
5641            Self::encode_full_zip(
5642                column_idx,
5643                &field,
5644                compression_strategy.as_ref(),
5645                data_block,
5646                repdef,
5647                row_number,
5648                num_rows,
5649            )
5650        } else {
5651            Err(Error::invalid_input_source(format!("Cannot determine structural encoding for field {}.  This typically indicates an invalid value of the field metadata key {}", field.name, STRUCTURAL_ENCODING_META_KEY).into()))
5652        }
5653    }
5654
5655    // Creates encode tasks, consuming all buffered data
5656    fn do_flush(
5657        &mut self,
5658        arrays: Vec<ArrayRef>,
5659        repdefs: Vec<RepDefBuilder>,
5660        row_number: u64,
5661        num_rows: u64,
5662    ) -> Result<Vec<EncodeTask>> {
5663        let num_values = arrays.iter().map(|arr| arr.len() as u64).sum();
5664        let is_simple_validity = repdefs.iter().all(|rd| rd.is_simple_validity());
5665        let has_repdef_info = repdefs.iter().any(|rd| !rd.is_empty());
5666        let (repdef, miniblock_repdef_budget) =
5667            RepDefBuilder::serialize_with_miniblock_repdef_budget(
5668                repdefs,
5669                miniblock::max_repdef_levels_per_chunk,
5670                num_rows,
5671                num_values,
5672            )?;
5673        let pages = Self::split_pages_for_miniblock_repdef_budget(
5674            arrays,
5675            repdef,
5676            miniblock_repdef_budget,
5677            row_number,
5678            num_rows,
5679        )?;
5680
5681        let mut tasks = Vec::with_capacity(pages.len());
5682        let ctx = PrimitiveEncodeContext {
5683            column_idx: self.column_index,
5684            field: self.field.clone(),
5685            compression_strategy: self.compression_strategy.clone(),
5686            encoding_metadata: self.encoding_metadata.clone(),
5687            support_large_chunk: self.support_large_chunk,
5688            version: self.version,
5689            is_simple_validity,
5690            has_repdef_info,
5691        };
5692        for page in pages {
5693            let ctx = ctx.clone();
5694            let task = spawn_cpu(move || Self::encode_page(ctx, page)).boxed();
5695            tasks.push(task);
5696        }
5697        Ok(tasks)
5698    }
5699
5700    fn extract_validity_buf(
5701        array: Arc<dyn Array>,
5702        repdef: &mut RepDefBuilder,
5703        keep_original_array: bool,
5704    ) -> Result<Arc<dyn Array>> {
5705        if let Some(validity) = array.nulls() {
5706            if keep_original_array {
5707                repdef.add_validity_bitmap(validity.clone());
5708            } else {
5709                repdef.add_validity_bitmap(deep_copy_nulls(Some(validity)).unwrap());
5710            }
5711            let data_no_nulls = array.to_data().into_builder().nulls(None).build()?;
5712            Ok(make_array(data_no_nulls))
5713        } else {
5714            repdef.add_no_null(array.len());
5715            Ok(array)
5716        }
5717    }
5718
5719    fn extract_validity(
5720        mut array: Arc<dyn Array>,
5721        repdef: &mut RepDefBuilder,
5722        keep_original_array: bool,
5723    ) -> Result<Arc<dyn Array>> {
5724        match array.data_type() {
5725            DataType::Null => {
5726                repdef.add_validity_bitmap(NullBuffer::new(BooleanBuffer::new_unset(array.len())));
5727                Ok(array)
5728            }
5729            DataType::Dictionary(_, _) => {
5730                array = dict::normalize_dict_nulls(array)?;
5731                Self::extract_validity_buf(array, repdef, keep_original_array)
5732            }
5733            // Extract our validity buf but NOT any child validity bufs. (they will be encoded in
5734            // as part of the values).  Note: for FSL we do not use repdef.add_fsl because we do
5735            // NOT want to increase the repdef depth.
5736            //
5737            // This would be quite catasrophic for something like vector embeddings.  Imagine we
5738            // had thousands of vectors and some were null but no vector contained null items.  If
5739            // we treated the vectors (primitive FSL) like we treat structural FSL we would end up
5740            // with a rep/def value for every single item in the vector.
5741            _ => Self::extract_validity_buf(array, repdef, keep_original_array),
5742        }
5743    }
5744}
5745
5746impl FieldEncoder for PrimitiveStructuralEncoder {
5747    // Buffers data, if there is enough to write a page then we create an encode task
5748    fn maybe_encode(
5749        &mut self,
5750        array: ArrayRef,
5751        _external_buffers: &mut OutOfLineBuffers,
5752        mut repdef: RepDefBuilder,
5753        row_number: u64,
5754        num_rows: u64,
5755    ) -> Result<Vec<EncodeTask>> {
5756        let array = Self::extract_validity(array, &mut repdef, self.keep_original_array)?;
5757        self.accumulated_repdefs.push(repdef);
5758
5759        if let Some((arrays, row_number, num_rows)) =
5760            self.accumulation_queue.insert(array, row_number, num_rows)
5761        {
5762            let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
5763            Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
5764        } else {
5765            Ok(vec![])
5766        }
5767    }
5768
5769    // If there is any data left in the buffer then create an encode task from it
5770    fn flush(&mut self, _external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
5771        if let Some((arrays, row_number, num_rows)) = self.accumulation_queue.flush() {
5772            let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
5773            Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
5774        } else {
5775            Ok(vec![])
5776        }
5777    }
5778
5779    fn num_columns(&self) -> u32 {
5780        1
5781    }
5782
5783    fn finish(
5784        &mut self,
5785        _external_buffers: &mut OutOfLineBuffers,
5786    ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
5787        std::future::ready(Ok(vec![EncodedColumn::default()])).boxed()
5788    }
5789}
5790
5791#[cfg(test)]
5792#[allow(clippy::single_range_in_vec_init)]
5793mod tests {
5794    use super::{
5795        ChunkInstructions, DataBlock, DecodeMiniBlockTask, FixedPerValueDecompressor,
5796        FixedWidthDataBlock, FullZipCacheableState, FullZipDecodeDetails, FullZipReadSource,
5797        FullZipRepIndexDetails, FullZipScheduler, MiniBlockChunk, MiniBlockCompressed,
5798        MiniBlockRepIndex, PerValueDecompressor, PreambleAction, StructuralPageScheduler,
5799        VariableFullZipDecoder,
5800    };
5801    use crate::buffer::LanceBuffer;
5802    use crate::compression::DefaultDecompressionStrategy;
5803    use crate::constants::{
5804        COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_VALUES_COMPRESSION_LEVEL_META_KEY,
5805        DICT_VALUES_COMPRESSION_META_KEY, STRUCTURAL_ENCODING_META_KEY,
5806        STRUCTURAL_ENCODING_MINIBLOCK,
5807    };
5808    use crate::data::BlockInfo;
5809    use crate::decoder::{PageEncoding, StructuralFieldDecoder};
5810    use crate::encodings::logical::primitive::{
5811        ChunkDrainInstructions, PrimitiveStructuralEncoder, StructuralPrimitiveFieldDecoder,
5812    };
5813    use crate::format::ProtobufUtils21;
5814    use crate::format::pb21;
5815    use crate::format::pb21::compressive_encoding::Compression;
5816    use crate::repdef::build_control_word_iterator;
5817    use crate::testing::{TestCases, check_round_trip_encoding_of_data};
5818    use crate::version::LanceFileVersion;
5819    use arrow_array::{ArrayRef, Int8Array, StringArray};
5820    use arrow_schema::{DataType, Field as ArrowField};
5821    use std::collections::HashMap;
5822    use std::{collections::VecDeque, sync::Arc};
5823
5824    #[test]
5825    fn test_is_narrow() {
5826        let int8_array = Int8Array::from(vec![1, 2, 3]);
5827        let array_ref: ArrayRef = Arc::new(int8_array);
5828        let block = DataBlock::from_array(array_ref);
5829
5830        assert!(PrimitiveStructuralEncoder::is_narrow(&block));
5831
5832        let string_array = StringArray::from(vec![Some("hello"), Some("world")]);
5833        let block = DataBlock::from_array(string_array);
5834        assert!(PrimitiveStructuralEncoder::is_narrow(&block));
5835
5836        let string_array = StringArray::from(vec![
5837            Some("hello world".repeat(100)),
5838            Some("world".to_string()),
5839        ]);
5840        let block = DataBlock::from_array(string_array);
5841        assert!((!PrimitiveStructuralEncoder::is_narrow(&block)));
5842    }
5843
5844    #[test]
5845    fn test_primitive_decoder_empty_page_queue_returns_error() {
5846        let field = Arc::new(ArrowField::new("vector", DataType::Float32, true));
5847        let mut decoder = StructuralPrimitiveFieldDecoder::new(&field, false);
5848
5849        let err = decoder.drain(1).unwrap_err();
5850        assert!(
5851            matches!(&err, lance_core::Error::Internal { .. }),
5852            "expected internal error, got: {err:?}"
5853        );
5854        let message = err.to_string();
5855        for expected in [
5856            "Primitive decoder missing page decoder",
5857            "field 'vector'",
5858            "data_type=Float32",
5859            "requested_rows=1",
5860            "remaining_rows=1",
5861            "rows_drained_in_current=0",
5862            "queued_pages=0",
5863        ] {
5864            assert!(
5865                message.contains(expected),
5866                "expected error to contain {expected:?}, got: {message}"
5867            );
5868        }
5869    }
5870
5871    #[test]
5872    fn test_fullzip_fixed_rejects_non_byte_aligned_values() {
5873        let fixed = FixedWidthDataBlock {
5874            data: LanceBuffer::from(vec![0_u8]),
5875            bits_per_value: 1,
5876            num_values: 8,
5877            block_info: BlockInfo::new(),
5878        };
5879        let repdef = build_control_word_iterator(None, 0, None, 0, u16::MAX, 8);
5880
5881        let Err(err) = PrimitiveStructuralEncoder::serialize_full_zip_fixed(fixed, repdef, 8)
5882        else {
5883            panic!("expected full-zip to reject 1-bit fixed-width values");
5884        };
5885        assert!(
5886            err.to_string().contains("byte aligned"),
5887            "unexpected error: {err}"
5888        );
5889    }
5890
5891    #[test]
5892    fn test_map_range() {
5893        // Null in the middle
5894        // [[A, B, C], [D, E], NULL, [F, G, H]]
5895        let rep = Some(vec![1, 0, 0, 1, 0, 1, 1, 0, 0]);
5896        let def = Some(vec![0, 0, 0, 0, 0, 1, 0, 0, 0]);
5897        let max_visible_def = 0;
5898        let total_items = 8;
5899        let max_rep = 1;
5900
5901        let check = |range, expected_item_range, expected_level_range| {
5902            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5903                range,
5904                rep.as_ref(),
5905                def.as_ref(),
5906                max_rep,
5907                max_visible_def,
5908                total_items,
5909                PreambleAction::Absent,
5910            );
5911            assert_eq!(item_range, expected_item_range);
5912            assert_eq!(level_range, expected_level_range);
5913        };
5914
5915        check(0..1, 0..3, 0..3);
5916        check(1..2, 3..5, 3..5);
5917        check(2..3, 5..5, 5..6);
5918        check(3..4, 5..8, 6..9);
5919        check(0..2, 0..5, 0..5);
5920        check(1..3, 3..5, 3..6);
5921        check(2..4, 5..8, 5..9);
5922        check(0..3, 0..5, 0..6);
5923        check(1..4, 3..8, 3..9);
5924        check(0..4, 0..8, 0..9);
5925
5926        // Null at start
5927        // [NULL, [A, B], [C]]
5928        let rep = Some(vec![1, 1, 0, 1]);
5929        let def = Some(vec![1, 0, 0, 0]);
5930        let max_visible_def = 0;
5931        let total_items = 3;
5932
5933        let check = |range, expected_item_range, expected_level_range| {
5934            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5935                range,
5936                rep.as_ref(),
5937                def.as_ref(),
5938                max_rep,
5939                max_visible_def,
5940                total_items,
5941                PreambleAction::Absent,
5942            );
5943            assert_eq!(item_range, expected_item_range);
5944            assert_eq!(level_range, expected_level_range);
5945        };
5946
5947        check(0..1, 0..0, 0..1);
5948        check(1..2, 0..2, 1..3);
5949        check(2..3, 2..3, 3..4);
5950        check(0..2, 0..2, 0..3);
5951        check(1..3, 0..3, 1..4);
5952        check(0..3, 0..3, 0..4);
5953
5954        // Null at end
5955        // [[A], [B, C], NULL]
5956        let rep = Some(vec![1, 1, 0, 1]);
5957        let def = Some(vec![0, 0, 0, 1]);
5958        let max_visible_def = 0;
5959        let total_items = 3;
5960
5961        let check = |range, expected_item_range, expected_level_range| {
5962            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5963                range,
5964                rep.as_ref(),
5965                def.as_ref(),
5966                max_rep,
5967                max_visible_def,
5968                total_items,
5969                PreambleAction::Absent,
5970            );
5971            assert_eq!(item_range, expected_item_range);
5972            assert_eq!(level_range, expected_level_range);
5973        };
5974
5975        check(0..1, 0..1, 0..1);
5976        check(1..2, 1..3, 1..3);
5977        check(2..3, 3..3, 3..4);
5978        check(0..2, 0..3, 0..3);
5979        check(1..3, 1..3, 1..4);
5980        check(0..3, 0..3, 0..4);
5981
5982        // No nulls, with repetition
5983        // [[A, B], [C, D], [E, F]]
5984        let rep = Some(vec![1, 0, 1, 0, 1, 0]);
5985        let def: Option<&[u16]> = None;
5986        let max_visible_def = 0;
5987        let total_items = 6;
5988
5989        let check = |range, expected_item_range, expected_level_range| {
5990            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5991                range,
5992                rep.as_ref(),
5993                def.as_ref(),
5994                max_rep,
5995                max_visible_def,
5996                total_items,
5997                PreambleAction::Absent,
5998            );
5999            assert_eq!(item_range, expected_item_range);
6000            assert_eq!(level_range, expected_level_range);
6001        };
6002
6003        check(0..1, 0..2, 0..2);
6004        check(1..2, 2..4, 2..4);
6005        check(2..3, 4..6, 4..6);
6006        check(0..2, 0..4, 0..4);
6007        check(1..3, 2..6, 2..6);
6008        check(0..3, 0..6, 0..6);
6009
6010        // No repetition, with nulls (this case is trivial)
6011        // [A, B, NULL, C]
6012        let rep: Option<&[u16]> = None;
6013        let def = Some(vec![0, 0, 1, 0]);
6014        let max_visible_def = 1;
6015        let total_items = 4;
6016
6017        let check = |range, expected_item_range, expected_level_range| {
6018            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6019                range,
6020                rep.as_ref(),
6021                def.as_ref(),
6022                max_rep,
6023                max_visible_def,
6024                total_items,
6025                PreambleAction::Absent,
6026            );
6027            assert_eq!(item_range, expected_item_range);
6028            assert_eq!(level_range, expected_level_range);
6029        };
6030
6031        check(0..1, 0..1, 0..1);
6032        check(1..2, 1..2, 1..2);
6033        check(2..3, 2..3, 2..3);
6034        check(0..2, 0..2, 0..2);
6035        check(1..3, 1..3, 1..3);
6036        check(0..3, 0..3, 0..3);
6037
6038        // Tricky case, this chunk is a continuation and starts with a rep-index = 0
6039        // [[..., A] [B, C], NULL]
6040        //
6041        // What we do will depend on the preamble action
6042        let rep = Some(vec![0, 1, 0, 1]);
6043        let def = Some(vec![0, 0, 0, 1]);
6044        let max_visible_def = 0;
6045        let total_items = 3;
6046
6047        let check = |range, expected_item_range, expected_level_range| {
6048            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6049                range,
6050                rep.as_ref(),
6051                def.as_ref(),
6052                max_rep,
6053                max_visible_def,
6054                total_items,
6055                PreambleAction::Take,
6056            );
6057            assert_eq!(item_range, expected_item_range);
6058            assert_eq!(level_range, expected_level_range);
6059        };
6060
6061        // If we are taking the preamble then the range must start at 0
6062        check(0..1, 0..3, 0..3);
6063        check(0..2, 0..3, 0..4);
6064
6065        let check = |range, expected_item_range, expected_level_range| {
6066            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6067                range,
6068                rep.as_ref(),
6069                def.as_ref(),
6070                max_rep,
6071                max_visible_def,
6072                total_items,
6073                PreambleAction::Skip,
6074            );
6075            assert_eq!(item_range, expected_item_range);
6076            assert_eq!(level_range, expected_level_range);
6077        };
6078
6079        check(0..1, 1..3, 1..3);
6080        check(1..2, 3..3, 3..4);
6081        check(0..2, 1..3, 1..4);
6082
6083        // Another preamble case but now it doesn't end with a new list
6084        // [[..., A], NULL, [D, E]]
6085        //
6086        // What we do will depend on the preamble action
6087        let rep = Some(vec![0, 1, 1, 0]);
6088        let def = Some(vec![0, 1, 0, 0]);
6089        let max_visible_def = 0;
6090        let total_items = 4;
6091
6092        let check = |range, expected_item_range, expected_level_range| {
6093            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6094                range,
6095                rep.as_ref(),
6096                def.as_ref(),
6097                max_rep,
6098                max_visible_def,
6099                total_items,
6100                PreambleAction::Take,
6101            );
6102            assert_eq!(item_range, expected_item_range);
6103            assert_eq!(level_range, expected_level_range);
6104        };
6105
6106        // If we are taking the preamble then the range must start at 0
6107        check(0..1, 0..1, 0..2);
6108        check(0..2, 0..3, 0..4);
6109
6110        let check = |range, expected_item_range, expected_level_range| {
6111            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6112                range,
6113                rep.as_ref(),
6114                def.as_ref(),
6115                max_rep,
6116                max_visible_def,
6117                total_items,
6118                PreambleAction::Skip,
6119            );
6120            assert_eq!(item_range, expected_item_range);
6121            assert_eq!(level_range, expected_level_range);
6122        };
6123
6124        // If we are taking the preamble then the range must start at 0
6125        check(0..1, 1..1, 1..2);
6126        check(1..2, 1..3, 2..4);
6127        check(0..2, 1..3, 1..4);
6128
6129        // Now a preamble case without any definition levels
6130        // [[..., A] [B, C], [D]]
6131        let rep = Some(vec![0, 1, 0, 1]);
6132        let def: Option<Vec<u16>> = None;
6133        let max_visible_def = 0;
6134        let total_items = 4;
6135
6136        let check = |range, expected_item_range, expected_level_range| {
6137            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6138                range,
6139                rep.as_ref(),
6140                def.as_ref(),
6141                max_rep,
6142                max_visible_def,
6143                total_items,
6144                PreambleAction::Take,
6145            );
6146            assert_eq!(item_range, expected_item_range);
6147            assert_eq!(level_range, expected_level_range);
6148        };
6149
6150        // If we are taking the preamble then the range must start at 0
6151        check(0..1, 0..3, 0..3);
6152        check(0..2, 0..4, 0..4);
6153
6154        let check = |range, expected_item_range, expected_level_range| {
6155            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6156                range,
6157                rep.as_ref(),
6158                def.as_ref(),
6159                max_rep,
6160                max_visible_def,
6161                total_items,
6162                PreambleAction::Skip,
6163            );
6164            assert_eq!(item_range, expected_item_range);
6165            assert_eq!(level_range, expected_level_range);
6166        };
6167
6168        check(0..1, 1..3, 1..3);
6169        check(1..2, 3..4, 3..4);
6170        check(0..2, 1..4, 1..4);
6171
6172        // If we have nested lists then non-top level lists may be empty/null
6173        // and we need to make sure we still handle them as invisible items (we
6174        // failed to do this previously)
6175        let rep = Some(vec![2, 1, 2, 0, 1, 2]);
6176        let def = Some(vec![0, 1, 2, 0, 0, 0]);
6177        let max_rep = 2;
6178        let max_visible_def = 0;
6179        let total_items = 4;
6180
6181        let check = |range, expected_item_range, expected_level_range| {
6182            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6183                range,
6184                rep.as_ref(),
6185                def.as_ref(),
6186                max_rep,
6187                max_visible_def,
6188                total_items,
6189                PreambleAction::Absent,
6190            );
6191            assert_eq!(item_range, expected_item_range);
6192            assert_eq!(level_range, expected_level_range);
6193        };
6194
6195        check(0..3, 0..4, 0..6);
6196        check(0..1, 0..1, 0..2);
6197        check(1..2, 1..3, 2..5);
6198        check(2..3, 3..4, 5..6);
6199
6200        // Invisible items in a preamble that we are taking (regressing a previous failure)
6201        let rep = Some(vec![0, 0, 1, 0, 1, 1]);
6202        let def = Some(vec![0, 1, 0, 0, 0, 0]);
6203        let max_rep = 1;
6204        let max_visible_def = 0;
6205        let total_items = 5;
6206
6207        let check = |range, expected_item_range, expected_level_range| {
6208            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6209                range,
6210                rep.as_ref(),
6211                def.as_ref(),
6212                max_rep,
6213                max_visible_def,
6214                total_items,
6215                PreambleAction::Take,
6216            );
6217            assert_eq!(item_range, expected_item_range);
6218            assert_eq!(level_range, expected_level_range);
6219        };
6220
6221        check(0..0, 0..1, 0..2);
6222        check(0..1, 0..3, 0..4);
6223        check(0..2, 0..4, 0..5);
6224
6225        // Skip preamble (with invis items) and skip a few rows (with invis items)
6226        // and then take a few rows but not all the rows
6227        let rep = Some(vec![0, 1, 0, 1, 0, 1, 0, 1]);
6228        let def = Some(vec![1, 0, 1, 1, 0, 0, 0, 0]);
6229        let max_rep = 1;
6230        let max_visible_def = 0;
6231        let total_items = 5;
6232
6233        let check = |range, expected_item_range, expected_level_range| {
6234            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6235                range,
6236                rep.as_ref(),
6237                def.as_ref(),
6238                max_rep,
6239                max_visible_def,
6240                total_items,
6241                PreambleAction::Skip,
6242            );
6243            assert_eq!(item_range, expected_item_range);
6244            assert_eq!(level_range, expected_level_range);
6245        };
6246
6247        check(2..3, 2..4, 5..7);
6248    }
6249
6250    #[test]
6251    fn test_slice_batch_data_and_rebase_offsets_u32() {
6252        let data = LanceBuffer::copy_slice(b"0123456789abcdefghij");
6253        let offsets = LanceBuffer::reinterpret_vec(vec![6_u32, 8_u32, 8_u32, 12_u32]);
6254
6255        let (sliced_data, normalized_offsets) =
6256            VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
6257                .unwrap();
6258
6259        assert_eq!(sliced_data.as_ref(), b"6789ab");
6260        let normalized = normalized_offsets.borrow_to_typed_slice::<u32>();
6261        assert_eq!(normalized.as_ref(), &[0, 2, 2, 6]);
6262    }
6263
6264    #[test]
6265    fn test_slice_batch_data_and_rebase_offsets_u64() {
6266        let data = LanceBuffer::copy_slice(b"abcdefghijklmnopqrstuvwxyz");
6267        let offsets = LanceBuffer::reinterpret_vec(vec![10_u64, 12_u64, 16_u64, 20_u64]);
6268
6269        let (sliced_data, normalized_offsets) =
6270            VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 64)
6271                .unwrap();
6272
6273        assert_eq!(sliced_data.as_ref(), b"klmnopqrst");
6274        let normalized = normalized_offsets.borrow_to_typed_slice::<u64>();
6275        assert_eq!(normalized.as_ref(), &[0, 2, 6, 10]);
6276    }
6277
6278    #[test]
6279    fn test_slice_batch_data_and_rebase_offsets_rejects_invalid_offsets() {
6280        let data = LanceBuffer::copy_slice(b"abcd");
6281        let offsets = LanceBuffer::reinterpret_vec(vec![3_u32, 2_u32]);
6282
6283        let err = VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
6284            .expect_err("offset end before start should error");
6285        assert!(err.to_string().contains("less than base"));
6286    }
6287
6288    #[test]
6289    fn test_schedule_instructions() {
6290        // Convert repetition index to bytes for testing
6291        let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
6292        let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
6293        let repetition_index = MiniBlockRepIndex::decode_from_bytes(&rep_bytes, 2);
6294
6295        let check = |user_ranges, expected_instructions| {
6296            let instructions =
6297                ChunkInstructions::schedule_instructions(&repetition_index, user_ranges);
6298            assert_eq!(instructions, expected_instructions);
6299        };
6300
6301        // The instructions we expect if we're grabbing the whole range
6302        let expected_take_all = vec![
6303            ChunkInstructions {
6304                chunk_idx: 0,
6305                preamble: PreambleAction::Absent,
6306                rows_to_skip: 0,
6307                rows_to_take: 6,
6308                take_trailer: true,
6309            },
6310            ChunkInstructions {
6311                chunk_idx: 1,
6312                preamble: PreambleAction::Take,
6313                rows_to_skip: 0,
6314                rows_to_take: 2,
6315                take_trailer: false,
6316            },
6317            ChunkInstructions {
6318                chunk_idx: 2,
6319                preamble: PreambleAction::Absent,
6320                rows_to_skip: 0,
6321                rows_to_take: 5,
6322                take_trailer: true,
6323            },
6324            ChunkInstructions {
6325                chunk_idx: 3,
6326                preamble: PreambleAction::Take,
6327                rows_to_skip: 0,
6328                rows_to_take: 1,
6329                take_trailer: false,
6330            },
6331        ];
6332
6333        // Take all as 1 range
6334        check(&[0..14], expected_take_all.clone());
6335
6336        // Take all a individual rows
6337        check(
6338            &[
6339                0..1,
6340                1..2,
6341                2..3,
6342                3..4,
6343                4..5,
6344                5..6,
6345                6..7,
6346                7..8,
6347                8..9,
6348                9..10,
6349                10..11,
6350                11..12,
6351                12..13,
6352                13..14,
6353            ],
6354            expected_take_all,
6355        );
6356
6357        // Test some partial takes
6358
6359        // 2 rows in the same chunk but not contiguous
6360        check(
6361            &[0..1, 3..4],
6362            vec![
6363                ChunkInstructions {
6364                    chunk_idx: 0,
6365                    preamble: PreambleAction::Absent,
6366                    rows_to_skip: 0,
6367                    rows_to_take: 1,
6368                    take_trailer: false,
6369                },
6370                ChunkInstructions {
6371                    chunk_idx: 0,
6372                    preamble: PreambleAction::Absent,
6373                    rows_to_skip: 3,
6374                    rows_to_take: 1,
6375                    take_trailer: false,
6376                },
6377            ],
6378        );
6379
6380        // Taking just a trailer/preamble
6381        check(
6382            &[5..6],
6383            vec![
6384                ChunkInstructions {
6385                    chunk_idx: 0,
6386                    preamble: PreambleAction::Absent,
6387                    rows_to_skip: 5,
6388                    rows_to_take: 1,
6389                    take_trailer: true,
6390                },
6391                ChunkInstructions {
6392                    chunk_idx: 1,
6393                    preamble: PreambleAction::Take,
6394                    rows_to_skip: 0,
6395                    rows_to_take: 0,
6396                    take_trailer: false,
6397                },
6398            ],
6399        );
6400
6401        // Skipping an entire chunk
6402        check(
6403            &[7..10],
6404            vec![
6405                ChunkInstructions {
6406                    chunk_idx: 1,
6407                    preamble: PreambleAction::Skip,
6408                    rows_to_skip: 1,
6409                    rows_to_take: 1,
6410                    take_trailer: false,
6411                },
6412                ChunkInstructions {
6413                    chunk_idx: 2,
6414                    preamble: PreambleAction::Absent,
6415                    rows_to_skip: 0,
6416                    rows_to_take: 2,
6417                    take_trailer: false,
6418                },
6419            ],
6420        );
6421    }
6422
6423    #[test]
6424    fn test_drain_instructions() {
6425        fn drain_from_instructions(
6426            instructions: &mut VecDeque<ChunkInstructions>,
6427            mut rows_desired: u64,
6428            need_preamble: &mut bool,
6429            skip_in_chunk: &mut u64,
6430        ) -> Vec<ChunkDrainInstructions> {
6431            // Note: instructions.len() is an upper bound, we typically take much fewer
6432            let mut drain_instructions = Vec::with_capacity(instructions.len());
6433            while rows_desired > 0 || *need_preamble {
6434                let (next_instructions, consumed_chunk) = instructions
6435                    .front()
6436                    .unwrap()
6437                    .drain_from_instruction(&mut rows_desired, need_preamble, skip_in_chunk);
6438                if consumed_chunk {
6439                    instructions.pop_front();
6440                }
6441                drain_instructions.push(next_instructions);
6442            }
6443            drain_instructions
6444        }
6445
6446        // Convert repetition index to bytes for testing
6447        let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
6448        let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
6449        let repetition_index = MiniBlockRepIndex::decode_from_bytes(&rep_bytes, 2);
6450        let user_ranges = vec![1..7, 10..14];
6451
6452        // First, schedule the ranges
6453        let scheduled = ChunkInstructions::schedule_instructions(&repetition_index, &user_ranges);
6454
6455        let mut to_drain = VecDeque::from(scheduled.clone());
6456
6457        // Now we drain in batches of 4
6458
6459        let mut need_preamble = false;
6460        let mut skip_in_chunk = 0;
6461
6462        let next_batch =
6463            drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
6464
6465        assert!(!need_preamble);
6466        assert_eq!(skip_in_chunk, 4);
6467        assert_eq!(
6468            next_batch,
6469            vec![ChunkDrainInstructions {
6470                chunk_instructions: scheduled[0].clone(),
6471                rows_to_take: 4,
6472                rows_to_skip: 0,
6473                preamble_action: PreambleAction::Absent,
6474            }]
6475        );
6476
6477        let next_batch =
6478            drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
6479
6480        assert!(!need_preamble);
6481        assert_eq!(skip_in_chunk, 2);
6482
6483        assert_eq!(
6484            next_batch,
6485            vec![
6486                ChunkDrainInstructions {
6487                    chunk_instructions: scheduled[0].clone(),
6488                    rows_to_take: 1,
6489                    rows_to_skip: 4,
6490                    preamble_action: PreambleAction::Absent,
6491                },
6492                ChunkDrainInstructions {
6493                    chunk_instructions: scheduled[1].clone(),
6494                    rows_to_take: 1,
6495                    rows_to_skip: 0,
6496                    preamble_action: PreambleAction::Take,
6497                },
6498                ChunkDrainInstructions {
6499                    chunk_instructions: scheduled[2].clone(),
6500                    rows_to_take: 2,
6501                    rows_to_skip: 0,
6502                    preamble_action: PreambleAction::Absent,
6503                }
6504            ]
6505        );
6506
6507        let next_batch =
6508            drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
6509
6510        assert!(!need_preamble);
6511        assert_eq!(skip_in_chunk, 0);
6512
6513        assert_eq!(
6514            next_batch,
6515            vec![
6516                ChunkDrainInstructions {
6517                    chunk_instructions: scheduled[2].clone(),
6518                    rows_to_take: 1,
6519                    rows_to_skip: 2,
6520                    preamble_action: PreambleAction::Absent,
6521                },
6522                ChunkDrainInstructions {
6523                    chunk_instructions: scheduled[3].clone(),
6524                    rows_to_take: 1,
6525                    rows_to_skip: 0,
6526                    preamble_action: PreambleAction::Take,
6527                },
6528            ]
6529        );
6530
6531        // Regression case.  Need a chunk with preamble, rows, and trailer (the middle chunk here)
6532        let rep_data: Vec<u64> = vec![5, 2, 3, 3, 20, 0];
6533        let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
6534        let repetition_index = MiniBlockRepIndex::decode_from_bytes(&rep_bytes, 2);
6535        let user_ranges = vec![0..28];
6536
6537        // First, schedule the ranges
6538        let scheduled = ChunkInstructions::schedule_instructions(&repetition_index, &user_ranges);
6539
6540        let mut to_drain = VecDeque::from(scheduled.clone());
6541
6542        // Drain first chunk and some of second chunk
6543
6544        let mut need_preamble = false;
6545        let mut skip_in_chunk = 0;
6546
6547        let next_batch =
6548            drain_from_instructions(&mut to_drain, 7, &mut need_preamble, &mut skip_in_chunk);
6549
6550        assert_eq!(
6551            next_batch,
6552            vec![
6553                ChunkDrainInstructions {
6554                    chunk_instructions: scheduled[0].clone(),
6555                    rows_to_take: 6,
6556                    rows_to_skip: 0,
6557                    preamble_action: PreambleAction::Absent,
6558                },
6559                ChunkDrainInstructions {
6560                    chunk_instructions: scheduled[1].clone(),
6561                    rows_to_take: 1,
6562                    rows_to_skip: 0,
6563                    preamble_action: PreambleAction::Take,
6564                },
6565            ]
6566        );
6567
6568        assert!(!need_preamble);
6569        assert_eq!(skip_in_chunk, 1);
6570
6571        // Now, the tricky part.  We drain the second chunk, including the trailer, and need to make sure
6572        // we get a drain task to take the preamble of the third chunk (and nothing else)
6573        let next_batch =
6574            drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
6575
6576        assert_eq!(
6577            next_batch,
6578            vec![
6579                ChunkDrainInstructions {
6580                    chunk_instructions: scheduled[1].clone(),
6581                    rows_to_take: 2,
6582                    rows_to_skip: 1,
6583                    preamble_action: PreambleAction::Skip,
6584                },
6585                ChunkDrainInstructions {
6586                    chunk_instructions: scheduled[2].clone(),
6587                    rows_to_take: 0,
6588                    rows_to_skip: 0,
6589                    preamble_action: PreambleAction::Take,
6590                },
6591            ]
6592        );
6593
6594        assert!(!need_preamble);
6595        assert_eq!(skip_in_chunk, 0);
6596    }
6597
6598    #[tokio::test]
6599    async fn test_fullzip_initialize_is_lazy() {
6600        use futures::{FutureExt, future::BoxFuture};
6601        use std::ops::Range;
6602        use std::sync::Mutex;
6603
6604        #[derive(Debug, Clone)]
6605        struct RecordingScheduler {
6606            data: bytes::Bytes,
6607            requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
6608        }
6609
6610        impl RecordingScheduler {
6611            fn new(data: bytes::Bytes) -> Self {
6612                Self {
6613                    data,
6614                    requests: Arc::new(Mutex::new(Vec::new())),
6615                }
6616            }
6617
6618            fn requests(&self) -> Vec<Vec<Range<u64>>> {
6619                self.requests.lock().unwrap().clone()
6620            }
6621        }
6622
6623        impl crate::EncodingsIo for RecordingScheduler {
6624            fn submit_request(
6625                &self,
6626                ranges: Vec<Range<u64>>,
6627                _priority: u64,
6628            ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
6629                self.requests.lock().unwrap().push(ranges.clone());
6630                let data = ranges
6631                    .into_iter()
6632                    .map(|range| self.data.slice(range.start as usize..range.end as usize))
6633                    .collect::<Vec<_>>();
6634                std::future::ready(Ok(data)).boxed()
6635            }
6636        }
6637
6638        #[derive(Debug)]
6639        struct TestFixedDecompressor;
6640
6641        impl FixedPerValueDecompressor for TestFixedDecompressor {
6642            fn decompress(
6643                &self,
6644                _data: FixedWidthDataBlock,
6645                _num_rows: u64,
6646            ) -> crate::Result<DataBlock> {
6647                unimplemented!("Test decompressor")
6648            }
6649
6650            fn bits_per_value(&self) -> u64 {
6651                32
6652            }
6653        }
6654
6655        let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(vec![
6656            0;
6657            16 * 1024
6658        ])));
6659        let mut scheduler = FullZipScheduler {
6660            data_buf_position: 0,
6661            data_buf_size: 4096,
6662            rep_index: Some(FullZipRepIndexDetails {
6663                buf_position: 1000,
6664                bytes_per_value: 4,
6665            }),
6666            priority: 0,
6667            rows_in_page: 100,
6668            bits_per_offset: 32,
6669            details: Arc::new(FullZipDecodeDetails {
6670                value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
6671                def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
6672                ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
6673                max_rep: 0,
6674                max_visible_def: 0,
6675            }),
6676            cached_state: None,
6677            enable_cache: false,
6678        };
6679
6680        let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
6681        let cached_data = scheduler.initialize(&io_dyn).await.unwrap();
6682
6683        assert!(
6684            cached_data
6685                .as_arc_any()
6686                .downcast_ref::<super::NoCachedPageData>()
6687                .is_some(),
6688            "FullZip initialize should not eagerly load repetition index data"
6689        );
6690        assert!(scheduler.cached_state.is_none());
6691        assert!(
6692            io.requests().is_empty(),
6693            "FullZip initialize should not issue any I/O"
6694        );
6695    }
6696
6697    #[tokio::test]
6698    async fn test_fullzip_read_source_slices_prefetched_page() {
6699        let page_start = 200_u64;
6700        let page_data = LanceBuffer::copy_slice(&[0, 1, 2, 3, 4, 5, 6, 7]);
6701        let source = FullZipReadSource::PrefetchedPage {
6702            base_offset: page_start,
6703            data: page_data,
6704        };
6705        let ranges = vec![
6706            page_start..(page_start + 3),
6707            (page_start + 4)..(page_start + 8),
6708        ];
6709        let mut data = source.fetch(&ranges, 0).await.unwrap();
6710        assert_eq!(data.pop_front().unwrap().as_ref(), &[0, 1, 2]);
6711        assert_eq!(data.pop_front().unwrap().as_ref(), &[4, 5, 6, 7]);
6712    }
6713
6714    #[tokio::test]
6715    async fn test_fullzip_initialize_caches_rep_index_when_enabled() {
6716        use futures::{FutureExt, future::BoxFuture};
6717        use std::ops::Range;
6718        use std::sync::Mutex;
6719
6720        #[derive(Debug, Clone)]
6721        struct RecordingScheduler {
6722            data: bytes::Bytes,
6723            requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
6724        }
6725
6726        impl RecordingScheduler {
6727            fn new(data: bytes::Bytes) -> Self {
6728                Self {
6729                    data,
6730                    requests: Arc::new(Mutex::new(Vec::new())),
6731                }
6732            }
6733
6734            fn requests(&self) -> Vec<Vec<Range<u64>>> {
6735                self.requests.lock().unwrap().clone()
6736            }
6737        }
6738
6739        impl crate::EncodingsIo for RecordingScheduler {
6740            fn submit_request(
6741                &self,
6742                ranges: Vec<Range<u64>>,
6743                _priority: u64,
6744            ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
6745                self.requests.lock().unwrap().push(ranges.clone());
6746                let data = ranges
6747                    .into_iter()
6748                    .map(|range| self.data.slice(range.start as usize..range.end as usize))
6749                    .collect::<Vec<_>>();
6750                std::future::ready(Ok(data)).boxed()
6751            }
6752        }
6753
6754        #[derive(Debug)]
6755        struct TestFixedDecompressor;
6756
6757        impl FixedPerValueDecompressor for TestFixedDecompressor {
6758            fn decompress(
6759                &self,
6760                _data: FixedWidthDataBlock,
6761                _num_rows: u64,
6762            ) -> crate::Result<DataBlock> {
6763                unimplemented!("Test decompressor")
6764            }
6765
6766            fn bits_per_value(&self) -> u64 {
6767                32
6768            }
6769        }
6770
6771        let rows_in_page = 100_u64;
6772        let bytes_per_value = 4_u64;
6773        let rep_start = 1000_u64;
6774        let rep_size = ((rows_in_page + 1) * bytes_per_value) as usize;
6775        let mut data = vec![0_u8; 16 * 1024];
6776        data[rep_start as usize..rep_start as usize + rep_size].fill(7);
6777        let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(data)));
6778
6779        let mut scheduler = FullZipScheduler {
6780            data_buf_position: 0,
6781            data_buf_size: 4096,
6782            rep_index: Some(FullZipRepIndexDetails {
6783                buf_position: rep_start,
6784                bytes_per_value,
6785            }),
6786            priority: 0,
6787            rows_in_page,
6788            bits_per_offset: 32,
6789            details: Arc::new(FullZipDecodeDetails {
6790                value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
6791                def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
6792                ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
6793                max_rep: 0,
6794                max_visible_def: 0,
6795            }),
6796            cached_state: None,
6797            enable_cache: true,
6798        };
6799
6800        let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
6801        let cached_data = scheduler.initialize(&io_dyn).await.unwrap();
6802        assert!(
6803            cached_data
6804                .as_arc_any()
6805                .downcast_ref::<FullZipCacheableState>()
6806                .is_some()
6807        );
6808        assert!(scheduler.cached_state.is_some());
6809        assert_eq!(
6810            io.requests(),
6811            vec![vec![
6812                rep_start..(rep_start + (rows_in_page + 1) * bytes_per_value)
6813            ]]
6814        );
6815    }
6816
6817    #[tokio::test]
6818    async fn test_fullzip_full_page_bypasses_rep_index_io() {
6819        use futures::{FutureExt, future::BoxFuture};
6820        use std::ops::Range;
6821        use std::sync::Mutex;
6822
6823        #[derive(Debug, Clone)]
6824        struct RecordingScheduler {
6825            data: bytes::Bytes,
6826            requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
6827        }
6828
6829        impl RecordingScheduler {
6830            fn new(data: bytes::Bytes) -> Self {
6831                Self {
6832                    data,
6833                    requests: Arc::new(Mutex::new(Vec::new())),
6834                }
6835            }
6836
6837            fn requests(&self) -> Vec<Vec<Range<u64>>> {
6838                self.requests.lock().unwrap().clone()
6839            }
6840        }
6841
6842        impl crate::EncodingsIo for RecordingScheduler {
6843            fn submit_request(
6844                &self,
6845                ranges: Vec<Range<u64>>,
6846                _priority: u64,
6847            ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
6848                self.requests.lock().unwrap().push(ranges.clone());
6849                let data = ranges
6850                    .into_iter()
6851                    .map(|range| self.data.slice(range.start as usize..range.end as usize))
6852                    .collect::<Vec<_>>();
6853                std::future::ready(Ok(data)).boxed()
6854            }
6855        }
6856
6857        #[derive(Debug)]
6858        struct TestFixedDecompressor;
6859
6860        impl FixedPerValueDecompressor for TestFixedDecompressor {
6861            fn decompress(
6862                &self,
6863                _data: FixedWidthDataBlock,
6864                _num_rows: u64,
6865            ) -> crate::Result<DataBlock> {
6866                unimplemented!("Test decompressor")
6867            }
6868
6869            fn bits_per_value(&self) -> u64 {
6870                32
6871            }
6872        }
6873
6874        let rows_in_page = 100_u64;
6875        let data_start = 256_u64;
6876        let data_size = 500_u64;
6877        let rep_start = 4096_u64;
6878        let bytes_per_value = 4_u64;
6879
6880        let mut bytes = vec![0_u8; 16 * 1024];
6881        for i in 0..=rows_in_page {
6882            let offset = (i * 5) as u32;
6883            let pos = rep_start as usize + (i * bytes_per_value) as usize;
6884            bytes[pos..pos + 4].copy_from_slice(&offset.to_le_bytes());
6885        }
6886        let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(bytes)));
6887
6888        let scheduler = FullZipScheduler {
6889            data_buf_position: data_start,
6890            data_buf_size: data_size,
6891            rep_index: Some(FullZipRepIndexDetails {
6892                buf_position: rep_start,
6893                bytes_per_value,
6894            }),
6895            priority: 0,
6896            rows_in_page,
6897            bits_per_offset: 32,
6898            details: Arc::new(FullZipDecodeDetails {
6899                value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
6900                def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
6901                ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
6902                max_rep: 0,
6903                max_visible_def: 0,
6904            }),
6905            cached_state: None,
6906            enable_cache: false,
6907        };
6908
6909        let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
6910        let tasks = scheduler
6911            .schedule_ranges_rep(
6912                &[0..rows_in_page],
6913                &io_dyn,
6914                FullZipRepIndexDetails {
6915                    buf_position: rep_start,
6916                    bytes_per_value,
6917                },
6918            )
6919            .unwrap();
6920
6921        let requests = io.requests();
6922        assert_eq!(requests.len(), 1);
6923        assert_eq!(requests[0], vec![data_start..(data_start + data_size)]);
6924
6925        let _ = tasks.into_iter().next().unwrap().decoder_fut.await.unwrap();
6926        let requests_after_await = io.requests();
6927        assert_eq!(
6928            requests_after_await.len(),
6929            1,
6930            "full page path should not issue rep-index I/O"
6931        );
6932    }
6933
6934    /// This test is used to reproduce fuzz test https://github.com/lancedb/lance/issues/4492
6935    #[tokio::test]
6936    async fn test_fuzz_issue_4492_empty_rep_values() {
6937        use lance_datagen::{RowCount, Seed, array, gen_batch};
6938
6939        let seed = 1823859942947654717u64;
6940        let num_rows = 2741usize;
6941
6942        // Generate the exact same data that caused the failure
6943        let batch_gen = gen_batch().with_seed(Seed::from(seed));
6944        let base_generator = array::rand_type(&DataType::FixedSizeBinary(32));
6945        let list_generator = array::rand_list_any(base_generator, false);
6946
6947        let batch = batch_gen
6948            .anon_col(list_generator)
6949            .into_batch_rows(RowCount::from(num_rows as u64))
6950            .unwrap();
6951
6952        let list_array = batch.column(0).clone();
6953
6954        // Force miniblock encoding
6955        let mut metadata = HashMap::new();
6956        metadata.insert(
6957            STRUCTURAL_ENCODING_META_KEY.to_string(),
6958            STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
6959        );
6960
6961        let test_cases = TestCases::default()
6962            .with_min_file_version(LanceFileVersion::V2_1)
6963            .with_batch_size(100)
6964            .with_range(0..num_rows.min(500) as u64)
6965            .with_indices(vec![0, num_rows as u64 / 2, (num_rows - 1) as u64]);
6966
6967        check_round_trip_encoding_of_data(vec![list_array], &test_cases, metadata).await
6968    }
6969
6970    async fn test_minichunk_size_helper(
6971        string_data: Vec<Option<String>>,
6972        minichunk_size: u64,
6973        file_version: LanceFileVersion,
6974    ) {
6975        use crate::constants::MINICHUNK_SIZE_META_KEY;
6976        use crate::testing::{TestCases, check_round_trip_encoding_of_data};
6977        use arrow_array::{ArrayRef, StringArray};
6978        use std::sync::Arc;
6979
6980        let string_array: ArrayRef = Arc::new(StringArray::from(string_data));
6981
6982        let mut metadata = HashMap::new();
6983        metadata.insert(
6984            MINICHUNK_SIZE_META_KEY.to_string(),
6985            minichunk_size.to_string(),
6986        );
6987        metadata.insert(
6988            STRUCTURAL_ENCODING_META_KEY.to_string(),
6989            STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
6990        );
6991
6992        let test_cases = TestCases::default()
6993            .with_min_file_version(file_version)
6994            .with_batch_size(1000);
6995
6996        check_round_trip_encoding_of_data(vec![string_array], &test_cases, metadata).await;
6997    }
6998
6999    #[tokio::test]
7000    async fn test_minichunk_size_roundtrip() {
7001        // Test that minichunk size can be configured and works correctly in round-trip encoding
7002        let mut string_data = Vec::new();
7003        for i in 0..100 {
7004            string_data.push(Some(format!("test_string_{}", i).repeat(50)));
7005        }
7006        // configure minichunk size to 64 bytes (smaller than the default 4kb) for Lance 2.1
7007        test_minichunk_size_helper(string_data, 64, LanceFileVersion::V2_1).await;
7008    }
7009
7010    #[tokio::test]
7011    async fn test_minichunk_size_128kb_v2_2() {
7012        // Test that minichunk size can be configured to 128KB and works correctly with Lance 2.2
7013        let mut string_data = Vec::new();
7014        // create a 500kb string array
7015        for i in 0..10000 {
7016            string_data.push(Some(format!("test_string_{}", i).repeat(50)));
7017        }
7018        test_minichunk_size_helper(string_data, 128 * 1024, LanceFileVersion::V2_2).await;
7019    }
7020
7021    #[tokio::test]
7022    async fn test_binary_large_minichunk_size_over_max_miniblock_values() {
7023        let mut string_data = Vec::new();
7024        // 128kb/chunk / 6 bytes (t_9999) = 21845 items per chunk
7025        for i in 0..10000 {
7026            string_data.push(Some(format!("t_{}", i)));
7027        }
7028        test_minichunk_size_helper(string_data, 128 * 1024, LanceFileVersion::V2_2).await;
7029    }
7030
7031    #[tokio::test]
7032    async fn test_large_dictionary_general_compression() {
7033        use arrow_array::{ArrayRef, StringArray};
7034        use std::collections::HashMap;
7035        use std::sync::Arc;
7036
7037        // Create large string dictionary data (>32KiB) with low cardinality
7038        // Use 100 unique strings, each 500 bytes long = 50KB dictionary
7039        let unique_values: Vec<String> = (0..100)
7040            .map(|i| format!("value_{:04}_{}", i, "x".repeat(500)))
7041            .collect();
7042
7043        // Repeat these strings many times to create a large array
7044        let repeated_strings: Vec<_> = unique_values
7045            .iter()
7046            .cycle()
7047            .take(100_000)
7048            .map(|s| Some(s.as_str()))
7049            .collect();
7050
7051        let string_array = Arc::new(StringArray::from(repeated_strings)) as ArrayRef;
7052
7053        // Configure test to use V2_2 and verify encoding
7054        let test_cases = TestCases::default()
7055            .with_min_file_version(LanceFileVersion::V2_2)
7056            .with_verify_encoding(Arc::new(|cols: &[crate::encoder::EncodedColumn], _| {
7057                assert_eq!(cols.len(), 1);
7058                let col = &cols[0];
7059
7060                // Navigate to the dictionary encoding in the page layout
7061                if let Some(PageEncoding::Structural(page_layout)) =
7062                    &col.final_pages.first().map(|p| &p.description)
7063                    && let Some(pb21::page_layout::Layout::MiniBlockLayout(mini_block)) =
7064                        &page_layout.layout
7065                    && let Some(dictionary_encoding) = &mini_block.dictionary
7066                {
7067                    match dictionary_encoding.compression.as_ref() {
7068                        Some(Compression::General(general)) => {
7069                            // Verify it's using LZ4 or Zstd
7070                            let compression = general.compression.as_ref().unwrap();
7071                            assert!(
7072                                compression.scheme()
7073                                    == pb21::CompressionScheme::CompressionAlgorithmLz4
7074                                    || compression.scheme()
7075                                        == pb21::CompressionScheme::CompressionAlgorithmZstd,
7076                                "Expected LZ4 or Zstd compression for large dictionary"
7077                            );
7078                        }
7079                        _ => panic!("Expected General compression for large dictionary"),
7080                    }
7081                }
7082            }));
7083
7084        check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await;
7085    }
7086
7087    fn dictionary_encoding_from_page(
7088        page: &crate::encoder::EncodedPage,
7089    ) -> &crate::format::pb21::CompressiveEncoding {
7090        let PageEncoding::Structural(layout) = &page.description else {
7091            panic!("Expected structural page encoding");
7092        };
7093        let pb21::page_layout::Layout::MiniBlockLayout(layout) = layout.layout.as_ref().unwrap()
7094        else {
7095            panic!("Expected mini-block layout");
7096        };
7097        layout
7098            .dictionary
7099            .as_ref()
7100            .unwrap_or_else(|| panic!("Expected dictionary encoding"))
7101    }
7102
7103    async fn encode_variable_dict_page(
7104        metadata: HashMap<String, String>,
7105    ) -> crate::encoder::EncodedPage {
7106        use arrow_array::types::Int32Type;
7107        use arrow_array::{ArrayRef, DictionaryArray, Int32Array, StringArray};
7108
7109        let values = Arc::new(StringArray::from(
7110            (0..128)
7111                .map(|i| format!("value_{i:04}_{}", "x".repeat(256)))
7112                .collect::<Vec<_>>(),
7113        )) as ArrayRef;
7114        let keys = Int32Array::from_iter_values((0..20_000).map(|i| i % 128));
7115        let dict_array =
7116            Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap()) as ArrayRef;
7117
7118        let field = arrow_schema::Field::new(
7119            "dict_col",
7120            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
7121            false,
7122        )
7123        .with_metadata(metadata);
7124
7125        encode_first_page(field, dict_array, LanceFileVersion::V2_2).await
7126    }
7127
7128    async fn encode_auto_fixed_dict_page(
7129        metadata: HashMap<String, String>,
7130    ) -> crate::encoder::EncodedPage {
7131        use arrow_array::{ArrayRef, Decimal128Array};
7132
7133        // 128-bit fixed-width values with low cardinality to trigger dictionary encoding.
7134        let values = (0..20_000)
7135            .map(|i| match i % 3 {
7136                0 => 10_i128,
7137                1 => 20_i128,
7138                _ => 30_i128,
7139            })
7140            .collect::<Vec<_>>();
7141        let decimal = Decimal128Array::from_iter_values(values)
7142            .with_precision_and_scale(38, 0)
7143            .unwrap();
7144        let decimal = Arc::new(decimal) as ArrayRef;
7145
7146        let mut field_metadata = metadata;
7147        // Strongly encourage dictionary encoding for this synthetic test data.
7148        field_metadata.insert(
7149            "lance-encoding:dict-size-ratio".to_string(),
7150            "0.99".to_string(),
7151        );
7152        let field = arrow_schema::Field::new("fixed_col", DataType::Decimal128(38, 0), false)
7153            .with_metadata(field_metadata);
7154
7155        encode_first_page(field, decimal, LanceFileVersion::V2_2).await
7156    }
7157
7158    #[tokio::test]
7159    async fn test_dict_values_general_compression_default_lz4_for_variable_dict_values() {
7160        let page = encode_variable_dict_page(HashMap::new()).await;
7161        let dictionary_encoding = dictionary_encoding_from_page(&page);
7162        let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
7163            panic!("Expected General compression for dictionary values");
7164        };
7165        let compression = general.compression.as_ref().unwrap();
7166        assert_eq!(
7167            compression.scheme(),
7168            pb21::CompressionScheme::CompressionAlgorithmLz4
7169        );
7170    }
7171
7172    #[tokio::test]
7173    async fn test_dict_values_general_compression_default_lz4_for_fixed_dict_values() {
7174        let page = encode_auto_fixed_dict_page(HashMap::new()).await;
7175        let dictionary_encoding = dictionary_encoding_from_page(&page);
7176        let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
7177            panic!("Expected General compression for dictionary values");
7178        };
7179        let compression = general.compression.as_ref().unwrap();
7180        assert_eq!(
7181            compression.scheme(),
7182            pb21::CompressionScheme::CompressionAlgorithmLz4
7183        );
7184    }
7185
7186    #[tokio::test]
7187    async fn test_dict_values_general_compression_zstd() {
7188        let mut metadata = HashMap::new();
7189        metadata.insert(
7190            DICT_VALUES_COMPRESSION_META_KEY.to_string(),
7191            "zstd".to_string(),
7192        );
7193        let page = encode_variable_dict_page(metadata).await;
7194        let dictionary_encoding = dictionary_encoding_from_page(&page);
7195        let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
7196            panic!("Expected General compression for dictionary values");
7197        };
7198        let compression = general.compression.as_ref().unwrap();
7199        assert_eq!(
7200            compression.scheme(),
7201            pb21::CompressionScheme::CompressionAlgorithmZstd
7202        );
7203    }
7204
7205    #[tokio::test]
7206    async fn test_dict_values_general_compression_none() {
7207        let mut metadata = HashMap::new();
7208        metadata.insert(
7209            DICT_VALUES_COMPRESSION_META_KEY.to_string(),
7210            "none".to_string(),
7211        );
7212        let page = encode_variable_dict_page(metadata).await;
7213        let dictionary_encoding = dictionary_encoding_from_page(&page);
7214        assert!(
7215            !matches!(
7216                dictionary_encoding.compression.as_ref(),
7217                Some(Compression::General(_))
7218            ),
7219            "Expected dictionary values to avoid General compression"
7220        );
7221    }
7222
7223    #[test]
7224    fn test_resolve_dict_values_compression_metadata_defaults_to_lz4() {
7225        let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
7226            &HashMap::new(),
7227            None,
7228            None,
7229        );
7230        assert_eq!(metadata.get(COMPRESSION_META_KEY), Some(&"lz4".to_string()),);
7231        assert!(!metadata.contains_key(COMPRESSION_LEVEL_META_KEY));
7232    }
7233
7234    #[test]
7235    fn test_resolve_dict_values_compression_metadata_metadata_overrides_env() {
7236        let field_metadata = HashMap::from([
7237            (
7238                DICT_VALUES_COMPRESSION_META_KEY.to_string(),
7239                "none".to_string(),
7240            ),
7241            (
7242                DICT_VALUES_COMPRESSION_LEVEL_META_KEY.to_string(),
7243                "7".to_string(),
7244            ),
7245        ]);
7246        let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
7247            &field_metadata,
7248            Some("zstd".to_string()),
7249            Some("3".to_string()),
7250        );
7251        assert_eq!(
7252            metadata.get(COMPRESSION_META_KEY),
7253            Some(&"none".to_string()),
7254        );
7255        assert_eq!(
7256            metadata.get(COMPRESSION_LEVEL_META_KEY),
7257            Some(&"7".to_string()),
7258        );
7259    }
7260
7261    #[test]
7262    fn test_resolve_dict_values_compression_metadata_env_fallback() {
7263        let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
7264            &HashMap::new(),
7265            Some("zstd".to_string()),
7266            Some("9".to_string()),
7267        );
7268        assert_eq!(
7269            metadata.get(COMPRESSION_META_KEY),
7270            Some(&"zstd".to_string()),
7271        );
7272        assert_eq!(
7273            metadata.get(COMPRESSION_LEVEL_META_KEY),
7274            Some(&"9".to_string()),
7275        );
7276    }
7277
7278    #[tokio::test]
7279    async fn test_dictionary_encode_int64() {
7280        use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
7281        use crate::testing::{TestCases, check_round_trip_encoding_of_data};
7282        use crate::version::LanceFileVersion;
7283        use arrow_array::{ArrayRef, Int64Array};
7284        use std::collections::HashMap;
7285        use std::sync::Arc;
7286
7287        // Low cardinality with poor RLE opportunity.
7288        let values = (0..1000)
7289            .map(|i| match i % 3 {
7290                0 => 10i64,
7291                1 => 20i64,
7292                _ => 30i64,
7293            })
7294            .collect::<Vec<_>>();
7295        let array = Arc::new(Int64Array::from(values)) as ArrayRef;
7296
7297        let mut metadata = HashMap::new();
7298        metadata.insert(
7299            STRUCTURAL_ENCODING_META_KEY.to_string(),
7300            STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
7301        );
7302        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
7303
7304        let test_cases = TestCases::default()
7305            .with_min_file_version(LanceFileVersion::V2_2)
7306            .with_batch_size(1000)
7307            .with_range(0..1000)
7308            .with_indices(vec![0, 1, 10, 999])
7309            .with_expected_encoding("dictionary");
7310
7311        check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
7312    }
7313
7314    #[tokio::test]
7315    async fn test_dictionary_encode_float64() {
7316        use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
7317        use crate::testing::{TestCases, check_round_trip_encoding_of_data};
7318        use crate::version::LanceFileVersion;
7319        use arrow_array::{ArrayRef, Float64Array};
7320        use std::collections::HashMap;
7321        use std::sync::Arc;
7322
7323        // Low cardinality with poor RLE opportunity.
7324        let values = (0..1000)
7325            .map(|i| match i % 3 {
7326                0 => 0.1f64,
7327                1 => 0.2f64,
7328                _ => 0.3f64,
7329            })
7330            .collect::<Vec<_>>();
7331        let array = Arc::new(Float64Array::from(values)) as ArrayRef;
7332
7333        let mut metadata = HashMap::new();
7334        metadata.insert(
7335            STRUCTURAL_ENCODING_META_KEY.to_string(),
7336            STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
7337        );
7338        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
7339
7340        let test_cases = TestCases::default()
7341            .with_min_file_version(LanceFileVersion::V2_2)
7342            .with_batch_size(1000)
7343            .with_range(0..1000)
7344            .with_indices(vec![0, 1, 10, 999])
7345            .with_expected_encoding("dictionary");
7346
7347        check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
7348    }
7349
7350    #[test]
7351    fn test_miniblock_dictionary_out_of_line_bitpacking_decode() {
7352        let rows = 10_000;
7353        let unique_values = 2_000;
7354
7355        let dictionary_encoding =
7356            ProtobufUtils21::out_of_line_bitpacking(64, ProtobufUtils21::flat(11, None));
7357        let layout = pb21::MiniBlockLayout {
7358            rep_compression: None,
7359            def_compression: None,
7360            value_compression: Some(ProtobufUtils21::flat(64, None)),
7361            dictionary: Some(dictionary_encoding),
7362            num_dictionary_items: unique_values,
7363            layers: vec![pb21::RepDefLayer::RepdefAllValidItem as i32],
7364            num_buffers: 1,
7365            repetition_index_depth: 0,
7366            num_items: rows,
7367            has_large_chunk: false,
7368        };
7369
7370        let buffer_offsets_and_sizes = vec![(0, 0), (0, 0), (0, 0)];
7371        let scheduler = super::MiniBlockScheduler::try_new(
7372            &buffer_offsets_and_sizes,
7373            /*priority=*/ 0,
7374            /*items_in_page=*/ rows,
7375            &layout,
7376            &DefaultDecompressionStrategy::default(),
7377        )
7378        .unwrap();
7379
7380        let dictionary = scheduler.dictionary.unwrap();
7381        assert_eq!(dictionary.num_dictionary_items, unique_values);
7382        assert_eq!(
7383            dictionary.dictionary_data_alignment,
7384            crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT
7385        );
7386    }
7387
7388    // Dictionary encoding decision tests
7389    fn create_test_fixed_data_block(
7390        num_values: u64,
7391        cardinality: u64,
7392        bits_per_value: u64,
7393    ) -> DataBlock {
7394        assert!(cardinality > 0);
7395        assert!(cardinality <= num_values);
7396        let block_info = BlockInfo::default();
7397
7398        assert_eq!(bits_per_value % 8, 0);
7399        let data = match bits_per_value {
7400            32 => {
7401                let values = (0..num_values)
7402                    .map(|i| (i % cardinality) as u32)
7403                    .collect::<Vec<_>>();
7404                crate::buffer::LanceBuffer::reinterpret_vec(values)
7405            }
7406            64 => {
7407                let values = (0..num_values).map(|i| i % cardinality).collect::<Vec<_>>();
7408                crate::buffer::LanceBuffer::reinterpret_vec(values)
7409            }
7410            128 => {
7411                let values = (0..num_values)
7412                    .map(|i| (i % cardinality) as u128)
7413                    .collect::<Vec<_>>();
7414                crate::buffer::LanceBuffer::reinterpret_vec(values)
7415            }
7416            _ => unreachable!(),
7417        };
7418        DataBlock::FixedWidth(FixedWidthDataBlock {
7419            bits_per_value,
7420            data,
7421            num_values,
7422            block_info,
7423        })
7424    }
7425
7426    /// Helper to create VariableWidth (string) test data block with exact cardinality
7427    fn create_test_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock {
7428        use arrow_array::StringArray;
7429
7430        assert!(cardinality <= num_values && cardinality > 0);
7431
7432        let mut values = Vec::with_capacity(num_values as usize);
7433        for i in 0..num_values {
7434            values.push(format!("value_{:016}", i % cardinality));
7435        }
7436
7437        let array = StringArray::from(values);
7438        DataBlock::from_array(Arc::new(array) as ArrayRef)
7439    }
7440
7441    fn create_sorted_string_array(num_values: u64, cardinality: u64) -> ArrayRef {
7442        use arrow_array::StringArray;
7443
7444        assert!(cardinality <= num_values && cardinality > 0);
7445
7446        let mut values = Vec::with_capacity(num_values as usize);
7447        for i in 0..num_values {
7448            let value_idx = i * cardinality / num_values;
7449            values.push(format!("value_{:016}", value_idx));
7450        }
7451
7452        Arc::new(StringArray::from(values)) as ArrayRef
7453    }
7454
7455    fn create_sorted_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock {
7456        DataBlock::from_array(create_sorted_string_array(num_values, cardinality))
7457    }
7458
7459    #[test]
7460    fn test_should_dictionary_encode() {
7461        use crate::constants::DICT_SIZE_RATIO_META_KEY;
7462        use lance_core::datatypes::Field as LanceField;
7463
7464        // Create data where dict encoding saves space
7465        let block = create_test_variable_width_block(1000, 10);
7466
7467        let mut metadata = HashMap::new();
7468        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7469        let arrow_field =
7470            arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7471        let field = LanceField::try_from(&arrow_field).unwrap();
7472
7473        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7474            &block,
7475            &field,
7476            LanceFileVersion::V2_1,
7477        );
7478
7479        assert!(
7480            result.is_some(),
7481            "Should use dictionary encode based on size"
7482        );
7483    }
7484
7485    #[test]
7486    fn test_block_sampling_detects_low_cardinality_in_short_sorted_runs() {
7487        let sample_count: usize = 4096;
7488        let num_values: u64 = 200_000;
7489        let cardinality: u64 = 8_000;
7490        let run_length = num_values / cardinality;
7491        let stride = num_values as usize / sample_count;
7492        assert!(
7493            stride > run_length as usize,
7494            "test must construct the stride > run_length case"
7495        );
7496
7497        let block = create_sorted_variable_width_block(num_values, cardinality);
7498        let sample_unique_ratio =
7499            PrimitiveStructuralEncoder::sample_unique_ratio(&block, sample_count).unwrap();
7500
7501        assert!(
7502            sample_unique_ratio.is_some_and(|ratio| ratio < 0.98),
7503            "sorted low-cardinality data must not be classified as near-unique"
7504        );
7505    }
7506
7507    #[test]
7508    fn test_should_dictionary_encode_sorted_low_cardinality() {
7509        use crate::constants::DICT_SIZE_RATIO_META_KEY;
7510        use lance_core::datatypes::Field as LanceField;
7511
7512        let block = create_sorted_variable_width_block(200_000, 8_000);
7513
7514        let mut metadata = HashMap::new();
7515        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7516        let arrow_field =
7517            arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7518        let field = LanceField::try_from(&arrow_field).unwrap();
7519
7520        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7521            &block,
7522            &field,
7523            LanceFileVersion::V2_2,
7524        );
7525
7526        assert!(
7527            result.is_some(),
7528            "sorted low-cardinality data should reach dictionary encoding"
7529        );
7530    }
7531
7532    #[test]
7533    fn test_should_not_dictionary_encode_sorted_high_cardinality_short_runs() {
7534        use crate::constants::DICT_SIZE_RATIO_META_KEY;
7535        use lance_core::datatypes::Field as LanceField;
7536
7537        let num_values = 200_002;
7538        let cardinality = 100_001;
7539        let block = create_sorted_variable_width_block(num_values, cardinality);
7540
7541        let mut metadata = HashMap::new();
7542        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7543        let arrow_field =
7544            arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7545        let field = LanceField::try_from(&arrow_field).unwrap();
7546
7547        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7548            &block,
7549            &field,
7550            LanceFileVersion::V2_2,
7551        );
7552
7553        assert!(
7554            result.is_none(),
7555            "sorted high-cardinality short runs should not trigger a full dictionary probe"
7556        );
7557    }
7558
7559    #[tokio::test]
7560    async fn test_encode_sorted_low_cardinality_uses_dictionary_layout() {
7561        use crate::constants::DICT_SIZE_RATIO_META_KEY;
7562
7563        let mut metadata = HashMap::new();
7564        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7565        let field = arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7566        let array = create_sorted_string_array(200_000, 8_000);
7567
7568        let page = encode_first_page(field, array, LanceFileVersion::V2_2).await;
7569        let _ = dictionary_encoding_from_page(&page);
7570    }
7571
7572    #[test]
7573    fn test_should_not_dictionary_encode_unsupported_bits() {
7574        use crate::constants::DICT_SIZE_RATIO_META_KEY;
7575        use lance_core::datatypes::Field as LanceField;
7576
7577        let block = create_test_fixed_data_block(1000, 1000, 32);
7578
7579        let mut metadata = HashMap::new();
7580        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7581        let arrow_field =
7582            arrow_schema::Field::new("test", DataType::Int32, false).with_metadata(metadata);
7583        let field = LanceField::try_from(&arrow_field).unwrap();
7584
7585        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7586            &block,
7587            &field,
7588            LanceFileVersion::V2_1,
7589        );
7590
7591        assert!(
7592            result.is_none(),
7593            "Should not use dictionary encode for unsupported bit width"
7594        );
7595    }
7596
7597    #[test]
7598    fn test_should_not_dictionary_encode_near_unique_sample() {
7599        use crate::constants::DICT_SIZE_RATIO_META_KEY;
7600        use lance_core::datatypes::Field as LanceField;
7601
7602        let num_values = 5000;
7603        let block = create_test_variable_width_block(num_values, num_values);
7604
7605        let mut metadata = HashMap::new();
7606        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "1.0".to_string());
7607        let arrow_field =
7608            arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7609        let field = LanceField::try_from(&arrow_field).unwrap();
7610
7611        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7612            &block,
7613            &field,
7614            LanceFileVersion::V2_1,
7615        );
7616
7617        assert!(
7618            result.is_none(),
7619            "Should not probe dictionary encoding for near-unique data"
7620        );
7621    }
7622
7623    #[test]
7624    fn test_v2_1_miniblock_serializes_log_num_values_15() {
7625        let miniblocks = MiniBlockCompressed {
7626            data: vec![LanceBuffer::from(vec![1_u8; 16])],
7627            chunks: vec![
7628                MiniBlockChunk {
7629                    buffer_sizes: vec![8],
7630                    log_num_values: 15,
7631                },
7632                MiniBlockChunk {
7633                    buffer_sizes: vec![8],
7634                    log_num_values: 0,
7635                },
7636            ],
7637            num_values: 32_769,
7638        };
7639
7640        let serialized =
7641            PrimitiveStructuralEncoder::serialize_miniblocks(miniblocks, None, None, false)
7642                .unwrap();
7643
7644        let chunk_metadata = serialized.metadata.borrow_to_typed_slice::<u16>();
7645        assert_eq!(chunk_metadata.len(), 2);
7646        assert_eq!(
7647            chunk_metadata[0] & 0x0F,
7648            15,
7649            "V2.1 metadata should use all 4 bits for log_num_values"
7650        );
7651    }
7652
7653    async fn encode_first_page(
7654        field: arrow_schema::Field,
7655        array: ArrayRef,
7656        version: LanceFileVersion,
7657    ) -> crate::encoder::EncodedPage {
7658        use crate::encoder::{
7659            ColumnIndexSequence, EncodingOptions, MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers,
7660            default_encoding_strategy,
7661        };
7662        use crate::repdef::RepDefBuilder;
7663
7664        let lance_field = lance_core::datatypes::Field::try_from(&field).unwrap();
7665        let encoding_strategy = default_encoding_strategy(version);
7666        let mut column_index_seq = ColumnIndexSequence::default();
7667        let encoding_options = EncodingOptions {
7668            cache_bytes_per_column: 1,
7669            max_page_bytes: 32 * 1024 * 1024,
7670            keep_original_array: true,
7671            buffer_alignment: MIN_PAGE_BUFFER_ALIGNMENT,
7672            version,
7673        };
7674
7675        let mut encoder = encoding_strategy
7676            .create_field_encoder(
7677                encoding_strategy.as_ref(),
7678                &lance_field,
7679                &mut column_index_seq,
7680                &encoding_options,
7681            )
7682            .unwrap();
7683
7684        let mut external_buffers = OutOfLineBuffers::new(0, MIN_PAGE_BUFFER_ALIGNMENT);
7685        let repdef = RepDefBuilder::default();
7686        let num_rows = array.len() as u64;
7687        let mut pages = Vec::new();
7688        for task in encoder
7689            .maybe_encode(array, &mut external_buffers, repdef, 0, num_rows)
7690            .unwrap()
7691        {
7692            pages.push(task.await.unwrap());
7693        }
7694        for task in encoder.flush(&mut external_buffers).unwrap() {
7695            pages.push(task.await.unwrap());
7696        }
7697        pages.into_iter().next().unwrap()
7698    }
7699
7700    #[tokio::test]
7701    async fn test_constant_layout_out_of_line_fixed_size_binary_v2_2() {
7702        use crate::format::pb21::page_layout::Layout;
7703
7704        let val = vec![0xABu8; 33];
7705        let arr: ArrayRef = Arc::new(
7706            arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
7707                std::iter::repeat_n(Some(val.as_slice()), 256),
7708                33,
7709            )
7710            .unwrap(),
7711        );
7712        let field = arrow_schema::Field::new("c", DataType::FixedSizeBinary(33), true);
7713        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7714
7715        let PageEncoding::Structural(layout) = &page.description else {
7716            panic!("Expected structural encoding");
7717        };
7718        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7719            panic!("Expected constant layout in slot 2");
7720        };
7721        assert!(layout.inline_value.is_none());
7722        assert_eq!(page.data.len(), 1);
7723
7724        let test_cases = TestCases::default()
7725            .with_min_file_version(LanceFileVersion::V2_2)
7726            .with_max_file_version(LanceFileVersion::V2_2)
7727            .with_page_sizes(vec![4096]);
7728        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7729    }
7730
7731    #[tokio::test]
7732    async fn test_constant_layout_out_of_line_utf8_v2_2() {
7733        use crate::format::pb21::page_layout::Layout;
7734
7735        let arr: ArrayRef = Arc::new(arrow_array::StringArray::from_iter_values(
7736            std::iter::repeat_n("hello", 512),
7737        ));
7738        let field = arrow_schema::Field::new("c", DataType::Utf8, true);
7739        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7740
7741        let PageEncoding::Structural(layout) = &page.description else {
7742            panic!("Expected structural encoding");
7743        };
7744        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7745            panic!("Expected constant layout in slot 2");
7746        };
7747        assert!(layout.inline_value.is_none());
7748        assert_eq!(page.data.len(), 1);
7749
7750        let test_cases = TestCases::default()
7751            .with_min_file_version(LanceFileVersion::V2_2)
7752            .with_max_file_version(LanceFileVersion::V2_2)
7753            .with_page_sizes(vec![4096]);
7754        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7755    }
7756
7757    #[tokio::test]
7758    async fn test_constant_layout_nullable_item_v2_2() {
7759        use crate::format::pb21::page_layout::Layout;
7760
7761        let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![
7762            Some(7),
7763            None,
7764            Some(7),
7765            None,
7766            Some(7),
7767        ]));
7768        let field = arrow_schema::Field::new("c", DataType::Int32, true);
7769        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7770
7771        let PageEncoding::Structural(layout) = &page.description else {
7772            panic!("Expected structural encoding");
7773        };
7774        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7775            panic!("Expected constant layout in slot 2");
7776        };
7777        assert!(layout.inline_value.is_some());
7778        assert_eq!(page.data.len(), 2);
7779
7780        let test_cases = TestCases::default()
7781            .with_min_file_version(LanceFileVersion::V2_2)
7782            .with_max_file_version(LanceFileVersion::V2_2)
7783            .with_page_sizes(vec![4096]);
7784        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7785    }
7786
7787    #[tokio::test]
7788    async fn test_constant_layout_list_repdef_v2_2() {
7789        use crate::format::pb21::page_layout::Layout;
7790        use arrow_array::builder::{Int32Builder, ListBuilder};
7791
7792        let mut builder = ListBuilder::new(Int32Builder::new());
7793        builder.values().append_value(7);
7794        builder.values().append_null();
7795        builder.values().append_value(7);
7796        builder.append(true);
7797
7798        builder.append(true);
7799
7800        builder.values().append_value(7);
7801        builder.append(true);
7802
7803        builder.append_null();
7804
7805        let arr: ArrayRef = Arc::new(builder.finish());
7806        let field = arrow_schema::Field::new(
7807            "c",
7808            DataType::List(Arc::new(arrow_schema::Field::new(
7809                "item",
7810                DataType::Int32,
7811                true,
7812            ))),
7813            true,
7814        );
7815        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7816
7817        let PageEncoding::Structural(layout) = &page.description else {
7818            panic!("Expected structural encoding");
7819        };
7820        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7821            panic!("Expected constant layout in slot 2");
7822        };
7823        assert!(layout.inline_value.is_some());
7824        assert_eq!(page.data.len(), 2);
7825
7826        let test_cases = TestCases::default()
7827            .with_min_file_version(LanceFileVersion::V2_2)
7828            .with_max_file_version(LanceFileVersion::V2_2)
7829            .with_page_sizes(vec![4096]);
7830        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7831    }
7832
7833    #[tokio::test]
7834    async fn test_constant_layout_fixed_size_list_not_used_v2_2() {
7835        use crate::format::pb21::page_layout::Layout;
7836        use arrow_array::builder::{FixedSizeListBuilder, Int32Builder};
7837
7838        let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 3);
7839        for _ in 0..64 {
7840            builder.values().append_value(1);
7841            builder.values().append_null();
7842            builder.values().append_value(3);
7843            builder.append(true);
7844        }
7845        let arr: ArrayRef = Arc::new(builder.finish());
7846        let field = arrow_schema::Field::new(
7847            "c",
7848            DataType::FixedSizeList(
7849                Arc::new(arrow_schema::Field::new("item", DataType::Int32, true)),
7850                3,
7851            ),
7852            true,
7853        );
7854        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7855
7856        if let PageEncoding::Structural(layout) = &page.description {
7857            assert!(
7858                !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
7859                "FixedSizeList should not use constant layout yet"
7860            );
7861        }
7862
7863        let test_cases = TestCases::default()
7864            .with_min_file_version(LanceFileVersion::V2_2)
7865            .with_max_file_version(LanceFileVersion::V2_2)
7866            .with_page_sizes(vec![4096]);
7867        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7868    }
7869
7870    #[tokio::test]
7871    async fn test_constant_layout_not_written_before_v2_2() {
7872        use crate::format::pb21::page_layout::Layout;
7873
7874        let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![7; 1024]));
7875        let field = arrow_schema::Field::new("c", DataType::Int32, true);
7876        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_1).await;
7877
7878        let PageEncoding::Structural(layout) = &page.description else {
7879            return;
7880        };
7881        assert!(
7882            !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
7883            "Should not emit constant layout before v2.2"
7884        );
7885
7886        let test_cases = TestCases::default()
7887            .with_min_file_version(LanceFileVersion::V2_1)
7888            .with_max_file_version(LanceFileVersion::V2_1)
7889            .with_page_sizes(vec![4096]);
7890        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7891    }
7892
7893    #[tokio::test]
7894    async fn test_all_null_constant_layout_still_works_v2_2() {
7895        use crate::format::pb21::page_layout::Layout;
7896
7897        let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![None, None, None]));
7898        let field = arrow_schema::Field::new("c", DataType::Int32, true);
7899        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7900
7901        let PageEncoding::Structural(layout) = &page.description else {
7902            panic!("Expected structural encoding");
7903        };
7904        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7905            panic!("Expected layout in slot 2");
7906        };
7907        assert!(layout.inline_value.is_none());
7908        assert_eq!(page.data.len(), 0);
7909
7910        let test_cases = TestCases::default()
7911            .with_min_file_version(LanceFileVersion::V2_2)
7912            .with_max_file_version(LanceFileVersion::V2_2)
7913            .with_page_sizes(vec![4096]);
7914        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7915    }
7916
7917    #[test]
7918    fn test_encode_decode_complex_all_null_vals_roundtrip() {
7919        use crate::compression::{
7920            DecompressionStrategy, DefaultCompressionStrategy, DefaultDecompressionStrategy,
7921        };
7922
7923        let values: Arc<[u16]> = Arc::from((0..2048).map(|i| (i % 5) as u16).collect::<Vec<u16>>());
7924
7925        let compression_strategy = DefaultCompressionStrategy::default();
7926        let decompression_strategy = DefaultDecompressionStrategy::default();
7927
7928        let (compressed_buf, encoding) = PrimitiveStructuralEncoder::encode_complex_all_null_vals(
7929            &values,
7930            &compression_strategy,
7931        )
7932        .unwrap();
7933
7934        let decompressor = decompression_strategy
7935            .create_block_decompressor(&encoding)
7936            .unwrap();
7937        let decompressed = decompressor
7938            .decompress(compressed_buf, values.len() as u64)
7939            .unwrap();
7940        let decompressed_fixed_width = decompressed.as_fixed_width().unwrap();
7941        assert_eq!(decompressed_fixed_width.num_values, values.len() as u64);
7942        assert_eq!(decompressed_fixed_width.bits_per_value, 16);
7943        let rep_result = decompressed_fixed_width.data.borrow_to_typed_slice::<u16>();
7944        assert_eq!(rep_result.as_ref(), values.as_ref());
7945    }
7946
7947    #[tokio::test]
7948    async fn test_complex_all_null_compression_gated_by_version() {
7949        use crate::format::pb21::page_layout::Layout;
7950        use arrow_array::ListArray;
7951
7952        let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
7953            (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
7954        );
7955        let arr: ArrayRef = Arc::new(list_array);
7956        let field = arrow_schema::Field::new(
7957            "c",
7958            DataType::List(Arc::new(arrow_schema::Field::new(
7959                "item",
7960                DataType::Int32,
7961                true,
7962            ))),
7963            true,
7964        );
7965
7966        let page_v21 = encode_first_page(field.clone(), arr.clone(), LanceFileVersion::V2_1).await;
7967        let PageEncoding::Structural(layout_v21) = &page_v21.description else {
7968            panic!("Expected structural encoding");
7969        };
7970        let Layout::ConstantLayout(layout_v21) = layout_v21.layout.as_ref().unwrap() else {
7971            panic!("Expected constant layout");
7972        };
7973        assert!(layout_v21.rep_compression.is_none());
7974        assert!(layout_v21.def_compression.is_none());
7975        assert_eq!(layout_v21.num_rep_values, 0);
7976        assert_eq!(layout_v21.num_def_values, 0);
7977
7978        let page_v22 = encode_first_page(field, arr, LanceFileVersion::V2_2).await;
7979        let PageEncoding::Structural(layout_v22) = &page_v22.description else {
7980            panic!("Expected structural encoding");
7981        };
7982        let Layout::ConstantLayout(layout_v22) = layout_v22.layout.as_ref().unwrap() else {
7983            panic!("Expected constant layout");
7984        };
7985        assert!(layout_v22.def_compression.is_some());
7986        assert!(layout_v22.num_def_values > 0);
7987    }
7988
7989    #[tokio::test]
7990    async fn test_complex_all_null_round_trip() {
7991        use arrow_array::ListArray;
7992
7993        let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
7994            (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
7995        );
7996
7997        let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_2);
7998        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
7999            .await;
8000    }
8001
8002    // https://github.com/lance-format/lance/issues/6681
8003    #[tokio::test]
8004    async fn test_sparse_boolean_list_roundtrip() {
8005        use arrow_array::builder::{BooleanBuilder, ListBuilder};
8006
8007        let mut list_builder = ListBuilder::new(BooleanBuilder::new());
8008        for i in 0..1000i32 {
8009            if i % 64 == 0 {
8010                // Alternate true/false so the array is not constant (constant path avoids the bug).
8011                list_builder.values().append_value(i % 128 == 0);
8012                list_builder.append(true);
8013            } else {
8014                list_builder.append(false);
8015            }
8016        }
8017        let list_array = Arc::new(list_builder.finish());
8018
8019        let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1);
8020        check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
8021    }
8022}