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        RepDefSlicer, SerializedRepDefs, StructuralPagePlan, 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
2705fn corrupt_file_named(name: &str, message: impl Into<String>) -> Error {
2706    Error::corrupt_file(name.into(), message)
2707}
2708
2709impl VariableFullZipDecoder {
2710    fn new(
2711        details: Arc<FullZipDecodeDetails>,
2712        data: VecDeque<LanceBuffer>,
2713        num_rows: u64,
2714        in_bits_per_length: u8,
2715        out_bits_per_offset: u8,
2716    ) -> Result<Self> {
2717        let decompressor = match details.value_decompressor {
2718            PerValueDecompressor::Variable(ref d) => d.clone(),
2719            _ => unreachable!(),
2720        };
2721
2722        assert_eq!(in_bits_per_length % 8, 0);
2723        assert!(out_bits_per_offset == 32 || out_bits_per_offset == 64);
2724
2725        let mut decoder = Self {
2726            details,
2727            decompressor,
2728            data: LanceBuffer::empty(),
2729            offsets: LanceBuffer::empty(),
2730            rep: LanceBuffer::empty().borrow_to_typed_slice(),
2731            def: LanceBuffer::empty().borrow_to_typed_slice(),
2732            bits_per_offset: out_bits_per_offset,
2733            repdef_starts: Vec::with_capacity(num_rows as usize + 1),
2734            data_starts: Vec::with_capacity(num_rows as usize + 1),
2735            offset_starts: Vec::with_capacity(num_rows as usize + 1),
2736            visible_item_counts: Vec::with_capacity(num_rows as usize + 1),
2737            current_idx: 0,
2738            num_rows,
2739        };
2740
2741        // There's no great time to do this and this is the least worst time.  If we don't unzip then
2742        // we can't slice the data during the decode phase.  This is because we need the offsets to be
2743        // unpacked to know where the values start and end.
2744        //
2745        // We don't want to unzip on the decode thread because that is a single-threaded path
2746        // We don't want to unzip on the scheduling thread because that is a single-threaded path
2747        //
2748        // Fortunately, we know variable length data will always be read indirectly and so we can do it
2749        // here, which should be on the indirect thread.  The primary disadvantage to doing it here is that
2750        // we load all the data into memory and then throw it away only to load it all into memory again during
2751        // the decode.
2752        //
2753        // There are some alternatives to investigate:
2754        //   - Instead of just reading the beginning and end of the rep index we could read the entire
2755        //     range in between.  This will give us the break points that we need for slicing and won't increase
2756        //     the number of IOPs but it will mean we are doing more total I/O and we need to load the rep index
2757        //     even when doing a full scan.
2758        //   - We could force each decode task to do a full unzip of all the data.  Each decode task now
2759        //     has to do more work but the work is all fused.
2760        //   - We could just try doing this work on the decode thread and see if it is a problem.
2761        decoder.unzip(data, in_bits_per_length, out_bits_per_offset, num_rows)?;
2762
2763        Ok(decoder)
2764    }
2765
2766    fn slice_batch_data_and_rebase_offsets_typed<T>(
2767        data: &LanceBuffer,
2768        offsets: &LanceBuffer,
2769    ) -> Result<(LanceBuffer, LanceBuffer)>
2770    where
2771        T: arrow_buffer::ArrowNativeType
2772            + Copy
2773            + PartialOrd
2774            + std::ops::Sub<Output = T>
2775            + std::fmt::Display
2776            + TryInto<usize>,
2777    {
2778        let offsets_slice = offsets.borrow_to_typed_slice::<T>();
2779        let offsets_slice = offsets_slice.as_ref();
2780        if offsets_slice.is_empty() {
2781            return Err(Error::internal(
2782                "Variable offsets cannot be empty".to_string(),
2783            ));
2784        }
2785
2786        let base = offsets_slice[0];
2787        let end = *offsets_slice.last().unwrap();
2788        if end < base {
2789            return Err(Error::internal(format!(
2790                "Invalid variable offsets: end ({end}) is less than base ({base})"
2791            )));
2792        }
2793
2794        let data_start = base.try_into().map_err(|_| {
2795            Error::internal(format!("Variable offset ({base}) does not fit into usize"))
2796        })?;
2797        let data_end = end.try_into().map_err(|_| {
2798            Error::internal(format!("Variable offset ({end}) does not fit into usize"))
2799        })?;
2800        if data_end > data.len() {
2801            return Err(Error::internal(format!(
2802                "Invalid variable offsets: end ({data_end}) exceeds data len ({})",
2803                data.len()
2804            )));
2805        }
2806
2807        let mut rebased_offsets = Vec::with_capacity(offsets_slice.len());
2808        for &offset in offsets_slice {
2809            if offset < base {
2810                return Err(Error::internal(format!(
2811                    "Invalid variable offsets: offset ({offset}) is less than base ({base})"
2812                )));
2813            }
2814            rebased_offsets.push(offset - base);
2815        }
2816
2817        let sliced_data = data.slice_with_length(data_start, data_end - data_start);
2818        // Copy into a compact buffer so each output batch owns only what it references.
2819        let sliced_data = LanceBuffer::copy_slice(&sliced_data);
2820        let rebased_offsets = LanceBuffer::reinterpret_vec(rebased_offsets);
2821        Ok((sliced_data, rebased_offsets))
2822    }
2823
2824    fn slice_batch_data_and_rebase_offsets(
2825        data: &LanceBuffer,
2826        offsets: &LanceBuffer,
2827        bits_per_offset: u8,
2828    ) -> Result<(LanceBuffer, LanceBuffer)> {
2829        match bits_per_offset {
2830            32 => Self::slice_batch_data_and_rebase_offsets_typed::<u32>(data, offsets),
2831            64 => Self::slice_batch_data_and_rebase_offsets_typed::<u64>(data, offsets),
2832            _ => Err(Error::internal(format!(
2833                "Unsupported bits_per_offset={bits_per_offset}"
2834            ))),
2835        }
2836    }
2837
2838    /// Reads a single length prefix from the front of `data`.
2839    ///
2840    /// The bytes come from the file. A page whose item walk ends with a partial
2841    /// trailing item leaves fewer than `bits_per_offset / 8` bytes here, so this
2842    /// is bounds checked and reports a corrupt file rather than reading past the
2843    /// end of the buffer.
2844    fn parse_length(data: &[u8], bits_per_offset: u8) -> Result<u64> {
2845        let width = bits_per_offset as usize / 8;
2846        if data.len() < width {
2847            return Err(corrupt_file_named(
2848                "variable_full_zip",
2849                format!(
2850                    "truncated length prefix: {} byte(s) remain in the page buffer but a \
2851                     {}-bit length prefix requires {}",
2852                    data.len(),
2853                    bits_per_offset,
2854                    width
2855                ),
2856            ));
2857        }
2858        Ok(match bits_per_offset {
2859            8 => data[0] as u64,
2860            16 => u16::from_le_bytes(data[..2].try_into().unwrap()) as u64,
2861            32 => u32::from_le_bytes(data[..4].try_into().unwrap()) as u64,
2862            64 => u64::from_le_bytes(data[..8].try_into().unwrap()),
2863            _ => unreachable!(),
2864        })
2865    }
2866
2867    fn unzip(
2868        &mut self,
2869        data: VecDeque<LanceBuffer>,
2870        in_bits_per_length: u8,
2871        out_bits_per_offset: u8,
2872        num_rows: u64,
2873    ) -> Result<()> {
2874        // This undercounts if there are lists but, at this point, we don't really know how many items we have
2875        let mut rep = Vec::with_capacity(num_rows as usize);
2876        let mut def = Vec::with_capacity(num_rows as usize);
2877        let bytes_cw = self.details.ctrl_word_parser.bytes_per_word() * num_rows as usize;
2878
2879        // This undercounts if there are lists
2880        // It can also overcount if there are invisible items
2881        let bytes_per_offset = out_bits_per_offset as usize / 8;
2882        let bytes_offsets = bytes_per_offset * (num_rows as usize + 1);
2883        let mut offsets_data = Vec::with_capacity(bytes_offsets);
2884
2885        let bytes_per_length = in_bits_per_length as usize / 8;
2886        let bytes_lengths = bytes_per_length * num_rows as usize;
2887
2888        let bytes_data = data.iter().map(|d| d.len()).sum::<usize>();
2889        // This overcounts since bytes_lengths and bytes_cw are undercounts
2890        // It can also undercount if there are invisible items (hence the saturating_sub)
2891        let mut unzipped_data =
2892            Vec::with_capacity((bytes_data - bytes_cw).saturating_sub(bytes_lengths));
2893
2894        let mut current_offset = 0_u64;
2895        let mut visible_item_count = 0_u64;
2896        for databuf in data.into_iter() {
2897            let mut databuf = databuf.as_ref();
2898            while !databuf.is_empty() {
2899                let data_start = unzipped_data.len();
2900                let offset_start = offsets_data.len();
2901                // We might have only-rep or only-def, neither, or both.  They move at the same
2902                // speed though so we only need one index into it
2903                let repdef_start = rep.len().max(def.len());
2904                // TODO: Kind of inefficient we parse the control word twice here
2905                let ctrl_desc = self.details.ctrl_word_parser.parse_desc(
2906                    databuf,
2907                    self.details.max_rep,
2908                    self.details.max_visible_def,
2909                );
2910                self.details
2911                    .ctrl_word_parser
2912                    .parse(databuf, &mut rep, &mut def);
2913                databuf = &databuf[self.details.ctrl_word_parser.bytes_per_word()..];
2914
2915                if ctrl_desc.is_new_row {
2916                    self.repdef_starts.push(repdef_start);
2917                    self.data_starts.push(data_start);
2918                    self.offset_starts.push(offset_start);
2919                    self.visible_item_counts.push(visible_item_count);
2920                }
2921                if ctrl_desc.is_visible {
2922                    visible_item_count += 1;
2923                    if ctrl_desc.is_valid_item {
2924                        let length = Self::parse_length(databuf, in_bits_per_length)?;
2925                        match out_bits_per_offset {
2926                            32 => offsets_data
2927                                .extend_from_slice(&(current_offset as u32).to_le_bytes()),
2928                            64 => offsets_data.extend_from_slice(&current_offset.to_le_bytes()),
2929                            _ => unreachable!(),
2930                        };
2931                        databuf = &databuf[bytes_per_offset..];
2932                        unzipped_data.extend_from_slice(&databuf[..length as usize]);
2933                        databuf = &databuf[length as usize..];
2934                        current_offset += length;
2935                    } else {
2936                        // Null items still get an offset
2937                        match out_bits_per_offset {
2938                            32 => offsets_data
2939                                .extend_from_slice(&(current_offset as u32).to_le_bytes()),
2940                            64 => offsets_data.extend_from_slice(&current_offset.to_le_bytes()),
2941                            _ => unreachable!(),
2942                        }
2943                    }
2944                }
2945            }
2946        }
2947        self.repdef_starts.push(rep.len().max(def.len()));
2948        self.data_starts.push(unzipped_data.len());
2949        self.offset_starts.push(offsets_data.len());
2950        self.visible_item_counts.push(visible_item_count);
2951        match out_bits_per_offset {
2952            32 => offsets_data.extend_from_slice(&(current_offset as u32).to_le_bytes()),
2953            64 => offsets_data.extend_from_slice(&current_offset.to_le_bytes()),
2954            _ => unreachable!(),
2955        };
2956        self.rep = ScalarBuffer::from(rep);
2957        self.def = ScalarBuffer::from(def);
2958        self.data = LanceBuffer::from(unzipped_data);
2959        self.offsets = LanceBuffer::from(offsets_data);
2960        Ok(())
2961    }
2962}
2963
2964impl StructuralPageDecoder for VariableFullZipDecoder {
2965    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
2966        let start = self.current_idx;
2967        let end = start + num_rows as usize;
2968
2969        let offset_start = self.offset_starts[start];
2970        let offset_end = self.offset_starts[end] + (self.bits_per_offset as usize / 8);
2971        let offsets = self
2972            .offsets
2973            .slice_with_length(offset_start, offset_end - offset_start);
2974        // Keep each batch's variable data buffer bounded to the selected rows.
2975        let (data, offsets) =
2976            Self::slice_batch_data_and_rebase_offsets(&self.data, &offsets, self.bits_per_offset)?;
2977
2978        let repdef_start = self.repdef_starts[start];
2979        let repdef_end = self.repdef_starts[end];
2980        let rep = if self.rep.is_empty() {
2981            self.rep.clone()
2982        } else {
2983            self.rep.slice(repdef_start, repdef_end - repdef_start)
2984        };
2985        let def = if self.def.is_empty() {
2986            self.def.clone()
2987        } else {
2988            self.def.slice(repdef_start, repdef_end - repdef_start)
2989        };
2990
2991        let visible_item_counts_start = self.visible_item_counts[start];
2992        let visible_item_counts_end = self.visible_item_counts[end];
2993        let num_visible_items = visible_item_counts_end - visible_item_counts_start;
2994
2995        self.current_idx += num_rows as usize;
2996
2997        Ok(Box::new(VariableFullZipDecodeTask {
2998            details: self.details.clone(),
2999            decompressor: self.decompressor.clone(),
3000            data,
3001            offsets,
3002            bits_per_offset: self.bits_per_offset,
3003            num_visible_items,
3004            rep,
3005            def,
3006        }))
3007    }
3008
3009    fn num_rows(&self) -> u64 {
3010        self.num_rows
3011    }
3012}
3013
3014#[derive(Debug)]
3015struct VariableFullZipDecodeTask {
3016    details: Arc<FullZipDecodeDetails>,
3017    decompressor: Arc<dyn VariablePerValueDecompressor>,
3018    data: LanceBuffer,
3019    offsets: LanceBuffer,
3020    bits_per_offset: u8,
3021    num_visible_items: u64,
3022    rep: ScalarBuffer<u16>,
3023    def: ScalarBuffer<u16>,
3024}
3025
3026impl DecodePageTask for VariableFullZipDecodeTask {
3027    fn decode(self: Box<Self>) -> Result<DecodedPage> {
3028        let block = VariableWidthBlock {
3029            data: self.data,
3030            offsets: self.offsets,
3031            bits_per_offset: self.bits_per_offset,
3032            num_values: self.num_visible_items,
3033            block_info: BlockInfo::new(),
3034        };
3035        let decomopressed = self.decompressor.decompress(block)?;
3036        let rep = if self.rep.is_empty() {
3037            None
3038        } else {
3039            Some(self.rep.to_vec())
3040        };
3041        let def = if self.def.is_empty() {
3042            None
3043        } else {
3044            Some(self.def.to_vec())
3045        };
3046        let unraveler = RepDefUnraveler::new(
3047            rep,
3048            def,
3049            self.details.def_meaning.clone(),
3050            self.num_visible_items,
3051        );
3052        Ok(DecodedPage {
3053            data: decomopressed,
3054            repdef: unraveler,
3055        })
3056    }
3057}
3058
3059#[derive(Debug)]
3060struct FullZipDecodeTaskItem {
3061    data: PerValueDataBlock,
3062    rows_in_buf: u64,
3063}
3064
3065/// A task to unzip and decompress full-zip encoded data when that data
3066/// has a fixed-width.
3067#[derive(Debug)]
3068struct FixedFullZipDecodeTask {
3069    details: Arc<FullZipDecodeDetails>,
3070    data: Vec<FullZipDecodeTaskItem>,
3071    num_rows: usize,
3072    bytes_per_value: usize,
3073}
3074
3075impl DecodePageTask for FixedFullZipDecodeTask {
3076    fn decode(self: Box<Self>) -> Result<DecodedPage> {
3077        // Multiply by 2 to make a stab at the size of the output buffer (which will be decompressed and thus bigger)
3078        let estimated_size_bytes = self
3079            .data
3080            .iter()
3081            .map(|task_item| task_item.data.data_size() as usize)
3082            .sum::<usize>()
3083            * 2;
3084        let mut data_builder =
3085            DataBlockBuilder::with_capacity_estimate(estimated_size_bytes as u64);
3086
3087        if self.details.ctrl_word_parser.bytes_per_word() == 0 {
3088            // Fast path, no need to unzip because there is no rep/def
3089            //
3090            // We decompress each buffer and add it to our output buffer
3091            for task_item in self.data.into_iter() {
3092                let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
3093                    unreachable!()
3094                };
3095                let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
3096                else {
3097                    unreachable!()
3098                };
3099                debug_assert_eq!(fixed_data.num_values, task_item.rows_in_buf);
3100                let decompressed = decompressor.decompress(fixed_data, task_item.rows_in_buf)?;
3101                data_builder.append(&decompressed, 0..task_item.rows_in_buf);
3102            }
3103
3104            let unraveler = RepDefUnraveler::new(
3105                None,
3106                None,
3107                self.details.def_meaning.clone(),
3108                self.num_rows as u64,
3109            );
3110
3111            Ok(DecodedPage {
3112                data: data_builder.finish(),
3113                repdef: unraveler,
3114            })
3115        } else {
3116            // Slow path, unzipping needed
3117            let mut rep = Vec::with_capacity(self.num_rows);
3118            let mut def = Vec::with_capacity(self.num_rows);
3119
3120            for task_item in self.data.into_iter() {
3121                let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
3122                    unreachable!()
3123                };
3124                let mut buf_slice = fixed_data.data.as_ref();
3125                let num_values = fixed_data.num_values as usize;
3126                // We will be unzipping repdef in to `rep` and `def` and the
3127                // values into `values` (which contains the compressed values)
3128                let mut values = Vec::with_capacity(
3129                    fixed_data.data.len()
3130                        - (self.details.ctrl_word_parser.bytes_per_word() * num_values),
3131                );
3132                let mut visible_items = 0;
3133                for _ in 0..num_values {
3134                    // Extract rep/def
3135                    self.details
3136                        .ctrl_word_parser
3137                        .parse(buf_slice, &mut rep, &mut def);
3138                    buf_slice = &buf_slice[self.details.ctrl_word_parser.bytes_per_word()..];
3139
3140                    let is_visible = def
3141                        .last()
3142                        .map(|d| *d <= self.details.max_visible_def)
3143                        .unwrap_or(true);
3144                    if is_visible {
3145                        // Extract value
3146                        values.extend_from_slice(buf_slice[..self.bytes_per_value].as_ref());
3147                        buf_slice = &buf_slice[self.bytes_per_value..];
3148                        visible_items += 1;
3149                    }
3150                }
3151
3152                // Finally, we decompress the values and add them to our output buffer
3153                let values_buf = LanceBuffer::from(values);
3154                let fixed_data = FixedWidthDataBlock {
3155                    bits_per_value: self.bytes_per_value as u64 * 8,
3156                    block_info: BlockInfo::new(),
3157                    data: values_buf,
3158                    num_values: visible_items,
3159                };
3160                let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
3161                else {
3162                    unreachable!()
3163                };
3164                let decompressed = decompressor.decompress(fixed_data, visible_items)?;
3165                data_builder.append(&decompressed, 0..visible_items);
3166            }
3167
3168            let repetition = if rep.is_empty() { None } else { Some(rep) };
3169            let definition = if def.is_empty() { None } else { Some(def) };
3170
3171            let unraveler = RepDefUnraveler::new(
3172                repetition,
3173                definition,
3174                self.details.def_meaning.clone(),
3175                self.num_rows as u64,
3176            );
3177            let data = data_builder.finish();
3178
3179            Ok(DecodedPage {
3180                data,
3181                repdef: unraveler,
3182            })
3183        }
3184    }
3185}
3186
3187#[derive(Debug)]
3188struct StructuralPrimitiveFieldSchedulingJob<'a> {
3189    scheduler: &'a StructuralPrimitiveFieldScheduler,
3190    ranges: Vec<Range<u64>>,
3191    page_idx: usize,
3192    range_idx: usize,
3193    global_row_offset: u64,
3194}
3195
3196impl<'a> StructuralPrimitiveFieldSchedulingJob<'a> {
3197    pub fn new(scheduler: &'a StructuralPrimitiveFieldScheduler, ranges: Vec<Range<u64>>) -> Self {
3198        Self {
3199            scheduler,
3200            ranges,
3201            page_idx: 0,
3202            range_idx: 0,
3203            global_row_offset: 0,
3204        }
3205    }
3206}
3207
3208impl StructuralSchedulingJob for StructuralPrimitiveFieldSchedulingJob<'_> {
3209    fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result<Vec<ScheduledScanLine>> {
3210        if self.range_idx >= self.ranges.len() {
3211            return Ok(Vec::new());
3212        }
3213        // Get our current range
3214        let mut range = self.ranges[self.range_idx].clone();
3215        let priority = range.start;
3216
3217        let mut cur_page = &self.scheduler.page_schedulers[self.page_idx];
3218        trace!(
3219            "Current range is {:?} and current page has {} rows",
3220            range, cur_page.num_rows
3221        );
3222        // Skip entire pages until we have some overlap with our next range
3223        while cur_page.num_rows + self.global_row_offset <= range.start {
3224            self.global_row_offset += cur_page.num_rows;
3225            self.page_idx += 1;
3226            trace!("Skipping entire page of {} rows", cur_page.num_rows);
3227            cur_page = &self.scheduler.page_schedulers[self.page_idx];
3228        }
3229
3230        // Now the cur_page has overlap with range.  Continue looping through ranges
3231        // until we find a range that exceeds the current page
3232
3233        let mut ranges_in_page = Vec::new();
3234        while cur_page.num_rows + self.global_row_offset > range.start {
3235            range.start = range.start.max(self.global_row_offset);
3236            let start_in_page = range.start - self.global_row_offset;
3237            let end_in_page = start_in_page + (range.end - range.start);
3238            let end_in_page = end_in_page.min(cur_page.num_rows);
3239            let last_in_range = (end_in_page + self.global_row_offset) >= range.end;
3240
3241            ranges_in_page.push(start_in_page..end_in_page);
3242            if last_in_range {
3243                self.range_idx += 1;
3244                if self.range_idx == self.ranges.len() {
3245                    break;
3246                }
3247                range = self.ranges[self.range_idx].clone();
3248            } else {
3249                break;
3250            }
3251        }
3252
3253        trace!(
3254            "Scheduling {} rows across {} ranges from page with {} rows (priority={}, column_index={}, page_index={})",
3255            ranges_in_page.iter().map(|r| r.end - r.start).sum::<u64>(),
3256            ranges_in_page.len(),
3257            cur_page.num_rows,
3258            priority,
3259            self.scheduler.column_index,
3260            cur_page.page_index,
3261        );
3262
3263        self.global_row_offset += cur_page.num_rows;
3264        self.page_idx += 1;
3265
3266        let page_decoders = cur_page
3267            .scheduler
3268            .schedule_ranges(&ranges_in_page, context.io())?;
3269
3270        let cur_path = context.current_path();
3271        page_decoders
3272            .into_iter()
3273            .map(|page_load_task| {
3274                let cur_path = cur_path.clone();
3275                let page_decoder = page_load_task.decoder_fut;
3276                let unloaded_page = async move {
3277                    let page_decoder = page_decoder.await?;
3278                    Ok(LoadedPageShard {
3279                        decoder: page_decoder,
3280                        path: cur_path,
3281                    })
3282                }
3283                .boxed();
3284                Ok(ScheduledScanLine {
3285                    decoders: vec![MessageType::UnloadedPage(UnloadedPageShard(unloaded_page))],
3286                    rows_scheduled: page_load_task.num_rows,
3287                })
3288            })
3289            .collect::<Result<Vec<_>>>()
3290    }
3291}
3292
3293#[derive(Debug)]
3294struct PageInfoAndScheduler {
3295    page_index: usize,
3296    num_rows: u64,
3297    scheduler: Box<dyn StructuralPageScheduler>,
3298}
3299
3300/// A scheduler for a leaf node
3301///
3302/// Here we look at the layout of the various pages and delegate scheduling to a scheduler
3303/// appropriate for the layout of the page.
3304#[derive(Debug)]
3305pub struct StructuralPrimitiveFieldScheduler {
3306    page_schedulers: Vec<PageInfoAndScheduler>,
3307    column_index: u32,
3308    // Identifies the requested decode shape (e.g. blob descriptor struct vs
3309    // raw bytes). Blob columns can produce multiple page scheduler variants
3310    // for the same physical column depending on the target field's data type,
3311    // and the cached page state types differ per variant. The view tag is
3312    // mixed into the cache key so different variants do not collide.
3313    view_tag: String,
3314}
3315
3316impl StructuralPrimitiveFieldScheduler {
3317    pub fn try_new(
3318        column_info: &ColumnInfo,
3319        decompressors: &dyn DecompressionStrategy,
3320        cache_repetition_index: bool,
3321        target_field: &Field,
3322    ) -> Result<Self> {
3323        let page_schedulers = column_info
3324            .page_infos
3325            .iter()
3326            .enumerate()
3327            .map(|(page_index, page_info)| {
3328                Self::page_info_to_scheduler(
3329                    page_info,
3330                    page_index,
3331                    decompressors,
3332                    cache_repetition_index,
3333                    target_field,
3334                )
3335            })
3336            .collect::<Result<Vec<_>>>()?;
3337        Ok(Self {
3338            page_schedulers,
3339            column_index: column_info.index,
3340            view_tag: format!("{:?}", target_field.data_type()),
3341        })
3342    }
3343
3344    fn page_layout_to_scheduler(
3345        page_info: &PageInfo,
3346        page_layout: &PageLayout,
3347        decompressors: &dyn DecompressionStrategy,
3348        cache_repetition_index: bool,
3349        target_field: &Field,
3350    ) -> Result<Box<dyn StructuralPageScheduler>> {
3351        use pb21::page_layout::Layout;
3352        Ok(match page_layout.layout.as_ref().expect_ok()? {
3353            Layout::MiniBlockLayout(mini_block) => Box::new(MiniBlockScheduler::try_new(
3354                &page_info.buffer_offsets_and_sizes,
3355                page_info.priority,
3356                mini_block.num_items,
3357                mini_block,
3358                decompressors,
3359            )?),
3360            Layout::FullZipLayout(full_zip) => {
3361                let mut scheduler = FullZipScheduler::try_new(
3362                    &page_info.buffer_offsets_and_sizes,
3363                    page_info.priority,
3364                    page_info.num_rows,
3365                    full_zip,
3366                    decompressors,
3367                )?;
3368                scheduler.enable_cache = cache_repetition_index;
3369                Box::new(scheduler)
3370            }
3371            Layout::ConstantLayout(constant_layout) => {
3372                let def_meaning = constant_layout
3373                    .layers
3374                    .iter()
3375                    .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
3376                    .collect::<Vec<_>>();
3377                let has_scalar_value = constant_layout.inline_value.is_some()
3378                    || page_info.buffer_offsets_and_sizes.len() == 1
3379                    || page_info.buffer_offsets_and_sizes.len() == 3;
3380                if has_scalar_value {
3381                    Box::new(constant::ConstantPageScheduler::try_new(
3382                        page_info.buffer_offsets_and_sizes.clone(),
3383                        constant_layout.inline_value.clone(),
3384                        target_field.data_type(),
3385                        def_meaning.into(),
3386                    )?) as Box<dyn StructuralPageScheduler>
3387                } else if def_meaning.len() == 1
3388                    && def_meaning[0] == DefinitionInterpretation::NullableItem
3389                {
3390                    Box::new(SimpleAllNullScheduler::default()) as Box<dyn StructuralPageScheduler>
3391                } else {
3392                    let rep_decompressor = constant_layout
3393                        .rep_compression
3394                        .as_ref()
3395                        .map(|encoding| decompressors.create_block_decompressor(encoding))
3396                        .transpose()?
3397                        .map(Arc::from);
3398
3399                    let def_decompressor = constant_layout
3400                        .def_compression
3401                        .as_ref()
3402                        .map(|encoding| decompressors.create_block_decompressor(encoding))
3403                        .transpose()?
3404                        .map(Arc::from);
3405
3406                    Box::new(ComplexAllNullScheduler::new(
3407                        page_info.buffer_offsets_and_sizes.clone(),
3408                        def_meaning.into(),
3409                        rep_decompressor,
3410                        def_decompressor,
3411                        constant_layout.num_rep_values,
3412                        constant_layout.num_def_values,
3413                    )) as Box<dyn StructuralPageScheduler>
3414                }
3415            }
3416            Layout::BlobLayout(blob) => {
3417                let inner_scheduler = Self::page_layout_to_scheduler(
3418                    page_info,
3419                    blob.inner_layout.as_ref().expect_ok()?.as_ref(),
3420                    decompressors,
3421                    cache_repetition_index,
3422                    target_field,
3423                )?;
3424                let def_meaning = blob
3425                    .layers
3426                    .iter()
3427                    .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
3428                    .collect::<Vec<_>>();
3429                if matches!(target_field.data_type(), DataType::Struct(_)) {
3430                    // User wants to decode blob into struct
3431                    Box::new(BlobDescriptionPageScheduler::new(
3432                        inner_scheduler,
3433                        def_meaning.into(),
3434                    ))
3435                } else {
3436                    // User wants to decode blob into binary data
3437                    Box::new(BlobPageScheduler::new(
3438                        inner_scheduler,
3439                        page_info.priority,
3440                        page_info.num_rows,
3441                        def_meaning.into(),
3442                    ))
3443                }
3444            }
3445        })
3446    }
3447
3448    fn page_info_to_scheduler(
3449        page_info: &PageInfo,
3450        page_index: usize,
3451        decompressors: &dyn DecompressionStrategy,
3452        cache_repetition_index: bool,
3453        target_field: &Field,
3454    ) -> Result<PageInfoAndScheduler> {
3455        let page_layout = page_info.encoding.as_structural();
3456        let scheduler = Self::page_layout_to_scheduler(
3457            page_info,
3458            page_layout,
3459            decompressors,
3460            cache_repetition_index,
3461            target_field,
3462        )?;
3463        Ok(PageInfoAndScheduler {
3464            page_index,
3465            num_rows: page_info.num_rows,
3466            scheduler,
3467        })
3468    }
3469}
3470
3471pub trait CachedPageData: Any + Send + Sync + DeepSizeOf + 'static {
3472    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static>;
3473}
3474
3475pub struct NoCachedPageData;
3476
3477impl DeepSizeOf for NoCachedPageData {
3478    fn deep_size_of_children(&self, _ctx: &mut Context) -> usize {
3479        0
3480    }
3481}
3482impl CachedPageData for NoCachedPageData {
3483    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
3484        self
3485    }
3486}
3487
3488pub struct CachedFieldData {
3489    pages: Vec<Arc<dyn CachedPageData>>,
3490}
3491
3492impl DeepSizeOf for CachedFieldData {
3493    fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
3494        self.pages.deep_size_of_children(ctx)
3495    }
3496}
3497
3498// Cache key for field data
3499//
3500// Both `column_index` and `view_tag` are part of the key because a single
3501// physical column can be decoded under more than one shape — a blob column,
3502// for instance, materializes as a `Struct<position, size>` descriptor in one
3503// scheduler variant and as the raw `LargeBinary` bytes in another. Each
3504// variant builds different `CachedPageData` types per page, so two readers
3505// that hit the same `column_index` with different shapes used to collide and
3506// crash with a downcast failure when loading cached state.
3507#[derive(Debug, Clone)]
3508pub struct FieldDataCacheKey {
3509    pub column_index: u32,
3510    pub view_tag: String,
3511}
3512
3513impl CacheKey for FieldDataCacheKey {
3514    type ValueType = CachedFieldData;
3515
3516    fn key(&self) -> std::borrow::Cow<'_, str> {
3517        format!("{}:{}", self.column_index, self.view_tag).into()
3518    }
3519
3520    fn type_name() -> &'static str {
3521        "FieldData"
3522    }
3523}
3524
3525impl StructuralFieldScheduler for StructuralPrimitiveFieldScheduler {
3526    fn initialize<'a>(
3527        &'a mut self,
3528        _filter: &'a FilterExpression,
3529        context: &'a SchedulerContext,
3530    ) -> BoxFuture<'a, Result<()>> {
3531        let cache_key = FieldDataCacheKey {
3532            column_index: self.column_index,
3533            view_tag: self.view_tag.clone(),
3534        };
3535        let cache = context.cache().clone();
3536
3537        async move {
3538            if let Some(cached_data) = cache.get_with_key(&cache_key).await {
3539                self.page_schedulers
3540                    .iter_mut()
3541                    .zip(cached_data.pages.iter())
3542                    .for_each(|(page_scheduler, cached_data)| {
3543                        page_scheduler.scheduler.load(cached_data);
3544                    });
3545                return Ok(());
3546            }
3547
3548            let page_data = self
3549                .page_schedulers
3550                .iter_mut()
3551                .map(|s| s.scheduler.initialize(context.io()))
3552                .collect::<FuturesOrdered<_>>();
3553
3554            let page_data = page_data.try_collect::<Vec<_>>().await?;
3555            let cached_data = Arc::new(CachedFieldData { pages: page_data });
3556            cache.insert_with_key(&cache_key, cached_data).await;
3557            Ok(())
3558        }
3559        .boxed()
3560    }
3561
3562    fn schedule_ranges<'a>(
3563        &'a self,
3564        ranges: &[Range<u64>],
3565        _filter: &FilterExpression,
3566    ) -> Result<Box<dyn StructuralSchedulingJob + 'a>> {
3567        let ranges = ranges.to_vec();
3568        Ok(Box::new(StructuralPrimitiveFieldSchedulingJob::new(
3569            self, ranges,
3570        )))
3571    }
3572}
3573
3574/// Takes the output from several pages decoders and
3575/// concatenates them.
3576#[derive(Debug)]
3577pub struct StructuralCompositeDecodeArrayTask {
3578    tasks: Vec<Box<dyn DecodePageTask>>,
3579    should_validate: bool,
3580    data_type: DataType,
3581}
3582
3583impl StructuralCompositeDecodeArrayTask {
3584    fn restore_validity(
3585        array: Arc<dyn Array>,
3586        unraveler: &mut CompositeRepDefUnraveler,
3587    ) -> Arc<dyn Array> {
3588        let validity = unraveler.unravel_validity(array.len());
3589        let Some(validity) = validity else {
3590            return array;
3591        };
3592        if array.data_type() == &DataType::Null {
3593            // We unravel from a null array but we don't add the null buffer because arrow-rs doesn't like it
3594            return array;
3595        }
3596        assert_eq!(validity.len(), array.len());
3597        // SAFETY: We've should have already asserted the buffers are all valid, we are just
3598        // adding null buffers to the array here
3599        make_array(unsafe {
3600            array
3601                .to_data()
3602                .into_builder()
3603                .nulls(Some(validity))
3604                .build_unchecked()
3605        })
3606    }
3607}
3608
3609impl StructuralDecodeArrayTask for StructuralCompositeDecodeArrayTask {
3610    fn decode(self: Box<Self>) -> Result<DecodedArray> {
3611        let mut arrays = Vec::with_capacity(self.tasks.len());
3612        let mut unravelers = Vec::with_capacity(self.tasks.len());
3613        let mut data_size = 0u64;
3614        for task in self.tasks {
3615            let decoded = task.decode()?;
3616            data_size += decoded.data.data_size();
3617            unravelers.push(decoded.repdef);
3618
3619            let array = make_array(
3620                decoded
3621                    .data
3622                    .into_arrow(self.data_type.clone(), self.should_validate)?,
3623            );
3624
3625            arrays.push(array);
3626        }
3627        let array_refs = arrays.iter().map(|arr| arr.as_ref()).collect::<Vec<_>>();
3628        let array = arrow_select::concat::concat(&array_refs)?;
3629        let mut repdef = CompositeRepDefUnraveler::new(unravelers);
3630
3631        let array = Self::restore_validity(array, &mut repdef);
3632
3633        Ok(DecodedArray {
3634            array,
3635            repdef,
3636            data_size,
3637        })
3638    }
3639}
3640
3641#[derive(Debug)]
3642pub struct StructuralPrimitiveFieldDecoder {
3643    field: Arc<ArrowField>,
3644    page_decoders: VecDeque<Box<dyn StructuralPageDecoder>>,
3645    should_validate: bool,
3646    rows_drained_in_current: u64,
3647}
3648
3649impl StructuralPrimitiveFieldDecoder {
3650    pub fn new(field: &Arc<ArrowField>, should_validate: bool) -> Self {
3651        Self {
3652            field: field.clone(),
3653            page_decoders: VecDeque::new(),
3654            should_validate,
3655            rows_drained_in_current: 0,
3656        }
3657    }
3658}
3659
3660impl StructuralFieldDecoder for StructuralPrimitiveFieldDecoder {
3661    fn accept_page(&mut self, child: LoadedPageShard) -> Result<()> {
3662        assert!(child.path.is_empty());
3663        self.page_decoders.push_back(child.decoder);
3664        Ok(())
3665    }
3666
3667    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn StructuralDecodeArrayTask>> {
3668        let mut remaining = num_rows;
3669        let mut tasks = Vec::new();
3670        while remaining > 0 {
3671            let cur_page = self.page_decoders.front_mut().unwrap();
3672            let num_in_page = cur_page.num_rows() - self.rows_drained_in_current;
3673            let to_take = num_in_page.min(remaining);
3674
3675            let task = cur_page.drain(to_take)?;
3676            tasks.push(task);
3677
3678            if to_take == num_in_page {
3679                self.page_decoders.pop_front();
3680                self.rows_drained_in_current = 0;
3681            } else {
3682                self.rows_drained_in_current += to_take;
3683            }
3684
3685            remaining -= to_take;
3686        }
3687        Ok(Box::new(StructuralCompositeDecodeArrayTask {
3688            tasks,
3689            should_validate: self.should_validate,
3690            data_type: self.field.data_type().clone(),
3691        }))
3692    }
3693
3694    fn data_type(&self) -> &DataType {
3695        self.field.data_type()
3696    }
3697}
3698
3699/// The serialized representation of full-zip data
3700struct SerializedFullZip {
3701    /// The zipped values buffer
3702    values: LanceBuffer,
3703    /// The repetition index (only present if there is repetition)
3704    repetition_index: Option<LanceBuffer>,
3705}
3706
3707// We align and pad mini-blocks to 8 byte boundaries for two reasons.  First,
3708// to allow us to store a chunk size in 12 bits.
3709//
3710// If we directly record the size in bytes with 12 bits we would be limited to
3711// 4KiB which is too small.  Since we know each mini-block consists of 8 byte
3712// words we can store the # of words instead which gives us 32KiB.
3713//
3714// Second, each chunk in a mini-block is aligned to 8 bytes.  This allows multi-byte
3715// values like offsets to be stored in a mini-block and safely read back out.  It also
3716// helps ensure zero-copy reads in cases where zero-copy is possible (e.g. no decoding
3717// needed).
3718//
3719// Note: by "aligned to 8 bytes" we mean BOTH "aligned to 8 bytes from the start of
3720// the page" and "aligned to 8 bytes from the start of the file."
3721const MINIBLOCK_ALIGNMENT: usize = 8;
3722
3723/// An encoder for primitive (leaf) arrays
3724///
3725/// This encoder is fairly complicated and follows a number of paths depending
3726/// on the data.
3727///
3728/// First, we convert the validity & offsets information into repetition and
3729/// definition levels.  Then we compress the data itself into a single buffer.
3730///
3731/// If the data is narrow then we encode the data in small chunks (each chunk
3732/// should be a few disk sectors and contains a buffer of repetition, a buffer
3733/// of definition, and a buffer of value data).  This approach is called
3734/// "mini-block".  These mini-blocks are stored into a single data buffer.
3735///
3736/// If the data is wide then we zip together the repetition and definition value
3737/// with the value data into a single buffer.  This approach is called "zipped".
3738///
3739/// If there is any repetition information then we create a repetition index
3740///
3741/// In addition, the compression process may create zero or more metadata buffers.
3742/// For example, a dictionary compression will create dictionary metadata.  Any
3743/// mini-block approach has a metadata buffer of block sizes.  This metadata is
3744/// stored in a separate buffer on disk and read at initialization time.
3745///
3746/// TODO: We should concatenate metadata buffers from all pages into a single buffer
3747/// at (roughly) the end of the file so there is, at most, one read per column of
3748/// metadata per file.
3749pub struct PrimitiveStructuralEncoder {
3750    // Accumulates arrays until we have enough data to justify a disk page
3751    accumulation_queue: AccumulationQueue,
3752
3753    keep_original_array: bool,
3754    support_large_chunk: bool,
3755    accumulated_repdefs: Vec<RepDefBuilder>,
3756    // The compression strategy we will use to compress the data
3757    compression_strategy: Arc<dyn CompressionStrategy>,
3758    column_index: u32,
3759    field: Field,
3760    encoding_metadata: Arc<HashMap<String, String>>,
3761    version: LanceFileVersion,
3762}
3763
3764struct CompressedLevelsChunk {
3765    data: LanceBuffer,
3766    num_levels: u16,
3767}
3768
3769struct CompressedLevels {
3770    data: Vec<CompressedLevelsChunk>,
3771    compression: CompressiveEncoding,
3772    rep_index: Option<LanceBuffer>,
3773}
3774
3775struct SerializedMiniBlockPage {
3776    num_buffers: u64,
3777    data: LanceBuffer,
3778    metadata: LanceBuffer,
3779}
3780
3781#[derive(Debug, Clone, Copy)]
3782struct DictEncodingBudget {
3783    max_dict_entries: u32,
3784    max_encoded_size: usize,
3785}
3786
3787// A primitive page after optional structural splitting.
3788struct PrimitivePageData {
3789    // Arrow leaf arrays that contain this page's visible values.
3790    arrays: Vec<ArrayRef>,
3791    // Repetition / definition levels aligned to this page.
3792    repdef: SerializedRepDefs,
3793    // Top-level row number of the first row in this page.
3794    row_number: u64,
3795    // Number of top-level rows in this page.
3796    num_rows: u64,
3797    // Present when one top-level row is too large for one miniblock rep/def chunk.
3798    unsplittable_miniblock_levels: Option<u64>,
3799}
3800
3801// Immutable encoder state shared by per-page encode tasks.
3802//
3803// Cloning this only clones Arc-backed configuration and field metadata.  Page data
3804// stays in PrimitivePageData and is moved into exactly one task.
3805#[derive(Clone)]
3806struct PrimitiveEncodeContext {
3807    // Column being encoded.
3808    column_idx: u32,
3809    // Logical field metadata for compression/layout selection.
3810    field: Field,
3811    // Compression strategy shared across pages.
3812    compression_strategy: Arc<dyn CompressionStrategy>,
3813    // Field-level encoding metadata such as structural encoding overrides.
3814    encoding_metadata: Arc<HashMap<String, String>>,
3815    // Whether miniblock chunks may use the v2.2 large-chunk metadata.
3816    support_large_chunk: bool,
3817    // Lance file version selected by the writer.
3818    version: LanceFileVersion,
3819    // True when the only rep/def information is simple nullable validity.
3820    is_simple_validity: bool,
3821    // True when the field has any non-empty rep/def information.
3822    has_repdef_info: bool,
3823}
3824
3825impl PrimitiveStructuralEncoder {
3826    pub fn try_new(
3827        options: &EncodingOptions,
3828        compression_strategy: Arc<dyn CompressionStrategy>,
3829        column_index: u32,
3830        field: Field,
3831        encoding_metadata: Arc<HashMap<String, String>>,
3832    ) -> Result<Self> {
3833        Ok(Self {
3834            accumulation_queue: AccumulationQueue::new(
3835                options.cache_bytes_per_column,
3836                column_index,
3837                options.keep_original_array,
3838            ),
3839            support_large_chunk: options.support_large_chunk(),
3840            keep_original_array: options.keep_original_array,
3841            accumulated_repdefs: Vec::new(),
3842            column_index,
3843            compression_strategy,
3844            field,
3845            encoding_metadata,
3846            version: options.version,
3847        })
3848    }
3849
3850    // TODO: This is a heuristic we may need to tune at some point
3851    //
3852    // As data gets narrow then the "zipping" process gets too expensive
3853    //   and we prefer mini-block
3854    // As data gets wide then the # of values per block shrinks (very wide)
3855    //   data doesn't even fit in a mini-block and the block overhead gets
3856    //   too large and we prefer zipped.
3857    fn is_narrow(data_block: &DataBlock) -> bool {
3858        const MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE: u64 = 256;
3859
3860        if let Some(max_len_array) = data_block.get_stat(Stat::MaxLength) {
3861            let max_len_array = max_len_array
3862                .as_any()
3863                .downcast_ref::<PrimitiveArray<UInt64Type>>()
3864                .unwrap();
3865            if max_len_array.value(0) < MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE {
3866                return true;
3867            }
3868        }
3869        false
3870    }
3871
3872    fn prefers_miniblock(
3873        data_block: &DataBlock,
3874        encoding_metadata: &HashMap<String, String>,
3875    ) -> bool {
3876        // If the user specifically requested miniblock then use it
3877        if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
3878            return user_requested.to_lowercase() == STRUCTURAL_ENCODING_MINIBLOCK;
3879        }
3880        // Otherwise only use miniblock if it is narrow
3881        Self::is_narrow(data_block)
3882    }
3883
3884    fn prefers_fullzip(encoding_metadata: &HashMap<String, String>) -> bool {
3885        // Fullzip is the backup option so the only reason we wouldn't use it is if the
3886        // user specifically requested not to use it (in which case we're probably going
3887        // to emit an error)
3888        if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
3889            return user_requested.to_lowercase() == STRUCTURAL_ENCODING_FULLZIP;
3890        }
3891        true
3892    }
3893
3894    // Converts value data, repetition levels, and definition levels into a single
3895    // buffer of mini-blocks.  In addition, creates a buffer of mini-block metadata
3896    // which tells us the size of each block.  Finally, if repetition is present then
3897    // we also create a buffer for the repetition index.
3898    //
3899    // Each chunk is serialized as:
3900    // | num_bufs (1 byte) | buf_lens (2 bytes per buffer) | P | buf0 | P | buf1 | ... | bufN | P |
3901    //
3902    // P - Padding inserted to ensure each buffer is 8-byte aligned and the buffer size is a multiple
3903    //     of 8 bytes (so that the next chunk is 8-byte aligned).
3904    //
3905    // Each block has a u16 word of metadata.  The upper 12 bits contain the
3906    // # of 8-byte words in the block (if the block does not fill the final word
3907    // then up to 7 bytes of padding are added).  The lower 4 bits describe the log_2
3908    // number of values (e.g. if there are 1024 then the lower 4 bits will be
3909    // 0xA)  All blocks except the last must have power-of-two number of values.
3910    // This not only makes metadata smaller but it makes decoding easier since
3911    // batch sizes are typically a power of 2.  4 bits would allow us to express
3912    // up to 32Ki values.
3913    //
3914    // This means blocks can have 1 to 32Ki values and 8 - 32Ki bytes.
3915    //
3916    // All metadata words are serialized (as little endian) into a single buffer
3917    // of metadata values.
3918    //
3919    // If there is repetition then we also create a repetition index.  This is a
3920    // single buffer of integer vectors (stored in row major order).  There is one
3921    // entry for each chunk.  The size of the vector is based on the depth of random
3922    // access we want to support.
3923    //
3924    // A vector of size 2 is the minimum and will support row-based random access (e.g.
3925    // "take the 57th row").  A vector of size 3 will support 1 level of nested access
3926    // (e.g. "take the 3rd item in the 57th row").  A vector of size 4 will support 2
3927    // levels of nested access and so on.
3928    //
3929    // The first number in the vector is the number of top-level rows that complete in
3930    // the chunk.  The second number is the number of second-level rows that complete
3931    // after the final top-level row completed (or beginning of the chunk if no top-level
3932    // row completes in the chunk).  And so on.  The final number in the vector is always
3933    // the number of leftover items not covered by earlier entries in the vector.
3934    //
3935    // Currently we are limited to 0 levels of nested access but that will change in the
3936    // future.
3937    //
3938    // The repetition index and the chunk metadata are read at initialization time and
3939    // cached in memory.
3940    fn serialize_miniblocks(
3941        miniblocks: MiniBlockCompressed,
3942        rep: Option<Vec<CompressedLevelsChunk>>,
3943        def: Option<Vec<CompressedLevelsChunk>>,
3944        support_large_chunk: bool,
3945    ) -> Result<SerializedMiniBlockPage> {
3946        let bytes_rep = rep
3947            .as_ref()
3948            .map(|rep| rep.iter().map(|r| r.data.len()).sum::<usize>())
3949            .unwrap_or(0);
3950        let bytes_def = def
3951            .as_ref()
3952            .map(|def| def.iter().map(|d| d.data.len()).sum::<usize>())
3953            .unwrap_or(0);
3954        let bytes_data = miniblocks.data.iter().map(|d| d.len()).sum::<usize>();
3955        let mut num_buffers = miniblocks.data.len();
3956        if rep.is_some() {
3957            num_buffers += 1;
3958        }
3959        if def.is_some() {
3960            num_buffers += 1;
3961        }
3962        // 2 bytes for the length of each buffer and up to 7 bytes of padding per buffer
3963        let max_extra = 9 * num_buffers;
3964        let mut data_buffer = Vec::with_capacity(bytes_rep + bytes_def + bytes_data + max_extra);
3965        let chunk_size_bytes = if support_large_chunk { 4 } else { 2 };
3966        let mut meta_buffer = Vec::with_capacity(miniblocks.chunks.len() * chunk_size_bytes);
3967
3968        let mut rep_iter = rep.map(|r| r.into_iter());
3969        let mut def_iter = def.map(|d| d.into_iter());
3970
3971        let mut buffer_offsets = vec![0; miniblocks.data.len()];
3972        for chunk in miniblocks.chunks {
3973            let start_pos = data_buffer.len();
3974            // Start of chunk should be aligned
3975            debug_assert_eq!(start_pos % MINIBLOCK_ALIGNMENT, 0);
3976
3977            let rep = rep_iter.as_mut().map(|r| r.next().unwrap());
3978            let def = def_iter.as_mut().map(|d| d.next().unwrap());
3979
3980            // Write the number of levels, or 0 if there is no rep/def
3981            let num_levels = rep
3982                .as_ref()
3983                .map(|r| r.num_levels)
3984                .unwrap_or(def.as_ref().map(|d| d.num_levels).unwrap_or(0));
3985            data_buffer.extend_from_slice(&num_levels.to_le_bytes());
3986
3987            // Write the buffer lengths
3988            if let Some(rep) = rep.as_ref() {
3989                let bytes_rep = u16::try_from(rep.data.len()).map_err(|_| {
3990                    Error::internal(format!(
3991                        "Repetition buffer size ({} bytes) too large",
3992                        rep.data.len()
3993                    ))
3994                })?;
3995                data_buffer.extend_from_slice(&bytes_rep.to_le_bytes());
3996            }
3997            if let Some(def) = def.as_ref() {
3998                let bytes_def = u16::try_from(def.data.len()).map_err(|_| {
3999                    Error::internal(format!(
4000                        "Definition buffer size ({} bytes) too large",
4001                        def.data.len()
4002                    ))
4003                })?;
4004                data_buffer.extend_from_slice(&bytes_def.to_le_bytes());
4005            }
4006
4007            if support_large_chunk {
4008                for &buffer_size in &chunk.buffer_sizes {
4009                    data_buffer.extend_from_slice(&buffer_size.to_le_bytes());
4010                }
4011            } else {
4012                for &buffer_size in &chunk.buffer_sizes {
4013                    let buffer_size = u16::try_from(buffer_size).map_err(|_| {
4014                        Error::internal(format!(
4015                            "Mini-block buffer size ({} bytes) too large for 16-bit metadata",
4016                            buffer_size
4017                        ))
4018                    })?;
4019                    data_buffer.extend_from_slice(&buffer_size.to_le_bytes());
4020                }
4021            }
4022
4023            // Pad
4024            let add_padding = |data_buffer: &mut Vec<u8>| {
4025                let pad = pad_bytes::<MINIBLOCK_ALIGNMENT>(data_buffer.len());
4026                data_buffer.extend(iter::repeat_n(FILL_BYTE, pad));
4027            };
4028            add_padding(&mut data_buffer);
4029
4030            // Write the buffers themselves
4031            if let Some(rep) = rep.as_ref() {
4032                data_buffer.extend_from_slice(&rep.data);
4033                add_padding(&mut data_buffer);
4034            }
4035            if let Some(def) = def.as_ref() {
4036                data_buffer.extend_from_slice(&def.data);
4037                add_padding(&mut data_buffer);
4038            }
4039            for (buffer_size, (buffer, buffer_offset)) in chunk
4040                .buffer_sizes
4041                .iter()
4042                .zip(miniblocks.data.iter().zip(buffer_offsets.iter_mut()))
4043            {
4044                let start = *buffer_offset;
4045                let end = start + *buffer_size as usize;
4046                *buffer_offset += *buffer_size as usize;
4047                data_buffer.extend_from_slice(&buffer[start..end]);
4048                add_padding(&mut data_buffer);
4049            }
4050
4051            let chunk_bytes = data_buffer.len() - start_pos;
4052            let max_chunk_size = if support_large_chunk {
4053                1_u64 << 31 // 28 bits of 8-byte words in u32 metadata
4054            } else {
4055                32 * 1024 // 32KiB limit with u16 metadata
4056            };
4057            if chunk_bytes == 0 || chunk_bytes as u64 > max_chunk_size {
4058                return Err(Error::internal(format!(
4059                    "Mini-block chunk size {} bytes exceeds the {} byte metadata limit",
4060                    chunk_bytes, max_chunk_size
4061                )));
4062            }
4063            if chunk_bytes % MINIBLOCK_ALIGNMENT != 0 {
4064                return Err(Error::internal(format!(
4065                    "Mini-block chunk size {} bytes is not aligned to {} bytes",
4066                    chunk_bytes, MINIBLOCK_ALIGNMENT
4067                )));
4068            }
4069            if chunk.log_num_values > 15 {
4070                return Err(Error::internal(format!(
4071                    "Mini-block log_num_values {} exceeds the 4-bit metadata limit",
4072                    chunk.log_num_values
4073                )));
4074            }
4075            // We subtract 1 here from chunk_bytes because we want to be able to express
4076            // a size of 32KiB and not (32Ki - 8)B which is what we'd get otherwise with
4077            // 0xFFF
4078            let divided_bytes = chunk_bytes / MINIBLOCK_ALIGNMENT;
4079            let divided_bytes_minus_one = (divided_bytes - 1) as u64;
4080
4081            let metadata = (divided_bytes_minus_one << 4) | chunk.log_num_values as u64;
4082            if support_large_chunk {
4083                meta_buffer.extend_from_slice(&(metadata as u32).to_le_bytes());
4084            } else {
4085                meta_buffer.extend_from_slice(&(metadata as u16).to_le_bytes());
4086            }
4087        }
4088
4089        let data_buffer = LanceBuffer::from(data_buffer);
4090        let metadata_buffer = LanceBuffer::from(meta_buffer);
4091
4092        Ok(SerializedMiniBlockPage {
4093            num_buffers: miniblocks.data.len() as u64,
4094            data: data_buffer,
4095            metadata: metadata_buffer,
4096        })
4097    }
4098
4099    /// Compresses a buffer of levels into chunks
4100    ///
4101    /// If these are repetition levels then we also calculate the repetition index here (that
4102    /// is the third return value)
4103    fn compress_levels(
4104        mut levels: RepDefSlicer<'_>,
4105        num_elements: u64,
4106        compression_strategy: &dyn CompressionStrategy,
4107        chunks: &[MiniBlockChunk],
4108        // This will be 0 if we are compressing def levels
4109        max_rep: u16,
4110    ) -> Result<CompressedLevels> {
4111        let mut rep_index = if max_rep > 0 {
4112            Vec::with_capacity(chunks.len())
4113        } else {
4114            vec![]
4115        };
4116        // Make the levels into a FixedWidth data block
4117        let num_levels = levels.num_levels() as u64;
4118        let levels_buf = levels.all_levels().clone();
4119
4120        let mut fixed_width_block = FixedWidthDataBlock {
4121            data: levels_buf,
4122            bits_per_value: 16,
4123            num_values: num_levels,
4124            block_info: BlockInfo::new(),
4125        };
4126        // Compute statistics to enable optimal compression for rep/def levels
4127        fixed_width_block.compute_stat();
4128
4129        let levels_block = DataBlock::FixedWidth(fixed_width_block);
4130        let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
4131        // Pick a block compressor
4132        let (compressor, compressor_desc) =
4133            compression_strategy.create_block_compressor(&levels_field, &levels_block)?;
4134        // Compress blocks of levels (sized according to the chunks)
4135        let mut level_chunks = Vec::with_capacity(chunks.len());
4136        let mut values_counter = 0;
4137        for (chunk_idx, chunk) in chunks.iter().enumerate() {
4138            let chunk_num_values = chunk.num_values(values_counter, num_elements);
4139            debug_assert!(chunk_num_values > 0);
4140            values_counter += chunk_num_values;
4141            let chunk_levels = if chunk_idx < chunks.len() - 1 {
4142                levels.slice_next(chunk_num_values as usize)
4143            } else {
4144                levels.slice_rest()
4145            };
4146            let num_chunk_levels = (chunk_levels.len() / 2) as u64;
4147            if max_rep > 0 {
4148                // If max_rep > 0 then we are working with rep levels and we need
4149                // to calculate the repetition index.  The repetition index for a
4150                // chunk is currently 2 values (in the future it may be more).
4151                //
4152                // The first value is the number of rows that _finish_ in the
4153                // chunk.
4154                //
4155                // The second value is the number of "leftovers" after the last
4156                // finished row in the chunk.
4157                let rep_values = chunk_levels.borrow_to_typed_slice::<u16>();
4158                let rep_values = rep_values.as_ref();
4159
4160                // We skip 1 here because a max_rep at spot 0 doesn't count as a finished list (we
4161                // will count it in the previous chunk)
4162                let mut num_rows = rep_values.iter().skip(1).filter(|v| **v == max_rep).count();
4163                let num_leftovers = if chunk_idx < chunks.len() - 1 {
4164                    rep_values
4165                        .iter()
4166                        .rev()
4167                        .position(|v| *v == max_rep)
4168                        // # of leftovers includes the max_rep spot
4169                        .map(|pos| pos + 1)
4170                        .unwrap_or(rep_values.len())
4171                } else {
4172                    // Last chunk can't have leftovers
4173                    0
4174                };
4175
4176                if chunk_idx != 0 && rep_values.first() == Some(&max_rep) {
4177                    // This chunk starts with a new row and so, if we thought we had leftovers
4178                    // in the previous chunk, we were mistaken
4179                    // TODO: Can use unchecked here
4180                    let rep_len = rep_index.len();
4181                    if rep_index[rep_len - 1] != 0 {
4182                        // We thought we had leftovers but that was actually a full row
4183                        rep_index[rep_len - 2] += 1;
4184                        rep_index[rep_len - 1] = 0;
4185                    }
4186                }
4187
4188                if chunk_idx == chunks.len() - 1 {
4189                    // The final list
4190                    num_rows += 1;
4191                }
4192                rep_index.push(num_rows as u64);
4193                rep_index.push(num_leftovers as u64);
4194            }
4195            let mut chunk_fixed_width = FixedWidthDataBlock {
4196                data: chunk_levels,
4197                bits_per_value: 16,
4198                num_values: num_chunk_levels,
4199                block_info: BlockInfo::new(),
4200            };
4201            chunk_fixed_width.compute_stat();
4202            let chunk_levels_block = DataBlock::FixedWidth(chunk_fixed_width);
4203            let compressed_levels = compressor.compress(chunk_levels_block)?;
4204            let num_levels = u16::try_from(num_chunk_levels).map_err(|_| {
4205                Error::invalid_input_source(
4206                    format!(
4207                        "Mini-block cannot encode {} rep/def levels in one chunk. \
4208                         This usually means a top-level row contains too much nested structure \
4209                         for the current layout.",
4210                        num_chunk_levels
4211                    )
4212                    .into(),
4213                )
4214            })?;
4215            level_chunks.push(CompressedLevelsChunk {
4216                data: compressed_levels,
4217                num_levels,
4218            });
4219        }
4220        debug_assert_eq!(levels.num_levels_remaining(), 0);
4221        let rep_index = if rep_index.is_empty() {
4222            None
4223        } else {
4224            Some(LanceBuffer::reinterpret_vec(rep_index))
4225        };
4226        Ok(CompressedLevels {
4227            data: level_chunks,
4228            compression: compressor_desc,
4229            rep_index,
4230        })
4231    }
4232
4233    fn encode_simple_all_null(
4234        column_idx: u32,
4235        num_rows: u64,
4236        row_number: u64,
4237    ) -> Result<EncodedPage> {
4238        let description =
4239            ProtobufUtils21::constant_layout(&[DefinitionInterpretation::NullableItem], None);
4240        Ok(EncodedPage {
4241            column_idx,
4242            data: vec![],
4243            description: PageEncoding::Structural(description),
4244            num_rows,
4245            row_number,
4246        })
4247    }
4248
4249    fn encode_complex_all_null_vals(
4250        data: &Arc<[u16]>,
4251        compression_strategy: &dyn CompressionStrategy,
4252    ) -> Result<(LanceBuffer, pb21::CompressiveEncoding)> {
4253        let buffer = LanceBuffer::reinterpret_slice(data.clone());
4254        let mut fixed_width_block = FixedWidthDataBlock {
4255            data: buffer,
4256            bits_per_value: 16,
4257            num_values: data.len() as u64,
4258            block_info: BlockInfo::new(),
4259        };
4260        fixed_width_block.compute_stat();
4261
4262        let levels_block = DataBlock::FixedWidth(fixed_width_block);
4263        let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
4264        let (compressor, encoding) =
4265            compression_strategy.create_block_compressor(&levels_field, &levels_block)?;
4266        let compressed_buffer = compressor.compress(levels_block)?;
4267        Ok((compressed_buffer, encoding))
4268    }
4269
4270    // Encodes a page where all values are null but we have rep/def
4271    // information that we need to store (e.g. to distinguish between
4272    // different kinds of null)
4273    fn encode_complex_all_null(
4274        column_idx: u32,
4275        repdef: crate::repdef::SerializedRepDefs,
4276        row_number: u64,
4277        num_rows: u64,
4278        version: LanceFileVersion,
4279        compression_strategy: &dyn CompressionStrategy,
4280    ) -> Result<EncodedPage> {
4281        if version.resolve() < LanceFileVersion::V2_2 {
4282            let rep_bytes = if let Some(rep) = repdef.repetition_levels.as_ref() {
4283                LanceBuffer::reinterpret_slice(rep.clone())
4284            } else {
4285                LanceBuffer::empty()
4286            };
4287
4288            let def_bytes = if let Some(def) = repdef.definition_levels.as_ref() {
4289                LanceBuffer::reinterpret_slice(def.clone())
4290            } else {
4291                LanceBuffer::empty()
4292            };
4293
4294            let description = ProtobufUtils21::constant_layout(&repdef.def_meaning, None);
4295            return Ok(EncodedPage {
4296                column_idx,
4297                data: vec![rep_bytes, def_bytes],
4298                description: PageEncoding::Structural(description),
4299                num_rows,
4300                row_number,
4301            });
4302        }
4303
4304        let (rep_bytes, rep_encoding, num_rep_values) = if let Some(rep) =
4305            repdef.repetition_levels.as_ref()
4306        {
4307            let num_values = rep.len() as u64;
4308            let (buffer, encoding) = Self::encode_complex_all_null_vals(rep, compression_strategy)?;
4309            (buffer, Some(encoding), num_values)
4310        } else {
4311            (LanceBuffer::empty(), None, 0)
4312        };
4313
4314        let (def_bytes, def_encoding, num_def_values) = if let Some(def) =
4315            repdef.definition_levels.as_ref()
4316        {
4317            let num_values = def.len() as u64;
4318            let (buffer, encoding) = Self::encode_complex_all_null_vals(def, compression_strategy)?;
4319            (buffer, Some(encoding), num_values)
4320        } else {
4321            (LanceBuffer::empty(), None, 0)
4322        };
4323
4324        let description = ProtobufUtils21::compressed_all_null_constant_layout(
4325            &repdef.def_meaning,
4326            rep_encoding,
4327            def_encoding,
4328            num_rep_values,
4329            num_def_values,
4330        );
4331        Ok(EncodedPage {
4332            column_idx,
4333            data: vec![rep_bytes, def_bytes],
4334            description: PageEncoding::Structural(description),
4335            num_rows,
4336            row_number,
4337        })
4338    }
4339
4340    fn leaf_validity(
4341        repdef: &crate::repdef::SerializedRepDefs,
4342        num_values: usize,
4343    ) -> Result<Option<BooleanBuffer>> {
4344        let rep = repdef
4345            .repetition_levels
4346            .as_ref()
4347            .map(|rep| rep.as_ref().to_vec());
4348        let def = repdef
4349            .definition_levels
4350            .as_ref()
4351            .map(|def| def.as_ref().to_vec());
4352        let mut unraveler = RepDefUnraveler::new(
4353            rep,
4354            def,
4355            repdef.def_meaning.clone().into(),
4356            num_values as u64,
4357        );
4358        if unraveler.is_all_valid() {
4359            return Ok(None);
4360        }
4361        let mut validity = BooleanBufferBuilder::new(num_values);
4362        unraveler.unravel_validity(&mut validity);
4363        Ok(Some(validity.finish()))
4364    }
4365
4366    fn is_constant_values(
4367        arrays: &[ArrayRef],
4368        scalar: &ArrayRef,
4369        validity: Option<&BooleanBuffer>,
4370    ) -> Result<bool> {
4371        debug_assert_eq!(scalar.len(), 1);
4372        debug_assert_eq!(scalar.null_count(), 0);
4373
4374        match scalar.data_type() {
4375            DataType::Boolean => {
4376                let mut global_idx = 0usize;
4377                let scalar_val = scalar.as_boolean().value(0);
4378                for arr in arrays {
4379                    let bool_arr = arr.as_boolean();
4380                    for i in 0..arr.len() {
4381                        let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4382                        global_idx += 1;
4383                        if !is_valid {
4384                            continue;
4385                        }
4386                        if bool_arr.value(i) != scalar_val {
4387                            return Ok(false);
4388                        }
4389                    }
4390                }
4391                Ok(true)
4392            }
4393            DataType::Utf8 => Self::is_constant_utf8::<i32>(arrays, scalar, validity),
4394            DataType::LargeUtf8 => Self::is_constant_utf8::<i64>(arrays, scalar, validity),
4395            DataType::Binary => Self::is_constant_binary::<i32>(arrays, scalar, validity),
4396            DataType::LargeBinary => Self::is_constant_binary::<i64>(arrays, scalar, validity),
4397            data_type => {
4398                let mut global_idx = 0usize;
4399                let Some(byte_width) = data_type.byte_width_opt() else {
4400                    return Ok(false);
4401                };
4402                let scalar_data = scalar.to_data();
4403                if scalar_data.buffers().len() != 1 || !scalar_data.child_data().is_empty() {
4404                    return Ok(false);
4405                }
4406                let scalar_bytes = scalar_data.buffers()[0].as_slice();
4407                if scalar_bytes.len() != byte_width {
4408                    return Ok(false);
4409                }
4410
4411                for arr in arrays {
4412                    let data = arr.to_data();
4413                    if data.buffers().is_empty() {
4414                        return Ok(false);
4415                    }
4416                    let buf = data.buffers()[0].as_slice();
4417                    let base = data.offset();
4418                    for i in 0..arr.len() {
4419                        let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4420                        global_idx += 1;
4421                        if !is_valid {
4422                            continue;
4423                        }
4424                        let start = (base + i) * byte_width;
4425                        if buf[start..start + byte_width] != scalar_bytes[..] {
4426                            return Ok(false);
4427                        }
4428                    }
4429                }
4430                Ok(true)
4431            }
4432        }
4433    }
4434
4435    fn is_constant_utf8<O: arrow_array::OffsetSizeTrait>(
4436        arrays: &[ArrayRef],
4437        scalar: &ArrayRef,
4438        validity: Option<&BooleanBuffer>,
4439    ) -> Result<bool> {
4440        debug_assert_eq!(scalar.len(), 1);
4441        let scalar_val = scalar.as_string::<O>().value(0).as_bytes();
4442        let mut global_idx = 0usize;
4443        for arr in arrays {
4444            let str_arr = arr.as_string::<O>();
4445            for i in 0..arr.len() {
4446                let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4447                global_idx += 1;
4448                if !is_valid {
4449                    continue;
4450                }
4451                if str_arr.value(i).as_bytes() != scalar_val {
4452                    return Ok(false);
4453                }
4454            }
4455        }
4456        Ok(true)
4457    }
4458
4459    fn is_constant_binary<O: arrow_array::OffsetSizeTrait>(
4460        arrays: &[ArrayRef],
4461        scalar: &ArrayRef,
4462        validity: Option<&BooleanBuffer>,
4463    ) -> Result<bool> {
4464        debug_assert_eq!(scalar.len(), 1);
4465        let scalar_val = scalar.as_binary::<O>().value(0);
4466        let mut global_idx = 0usize;
4467        for arr in arrays {
4468            let bin_arr = arr.as_binary::<O>();
4469            for i in 0..arr.len() {
4470                let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4471                global_idx += 1;
4472                if !is_valid {
4473                    continue;
4474                }
4475                if bin_arr.value(i) != scalar_val {
4476                    return Ok(false);
4477                }
4478            }
4479        }
4480        Ok(true)
4481    }
4482
4483    fn find_constant_scalar(
4484        arrays: &[ArrayRef],
4485        validity: Option<&BooleanBuffer>,
4486    ) -> Result<Option<ArrayRef>> {
4487        if arrays.is_empty() {
4488            return Ok(None);
4489        }
4490
4491        let global_scalar_idx = if let Some(validity) = validity {
4492            let Some(idx) = (0..validity.len()).find(|&i| validity.value(i)) else {
4493                return Ok(None);
4494            };
4495            idx
4496        } else {
4497            0
4498        };
4499
4500        let mut idx_remaining = global_scalar_idx;
4501        let mut scalar_arr_idx = 0usize;
4502        while scalar_arr_idx < arrays.len() {
4503            let len = arrays[scalar_arr_idx].len();
4504            if idx_remaining < len {
4505                break;
4506            }
4507            idx_remaining -= len;
4508            scalar_arr_idx += 1;
4509        }
4510
4511        if scalar_arr_idx >= arrays.len() {
4512            return Ok(None);
4513        }
4514
4515        let scalar =
4516            lance_arrow::scalar::extract_scalar_value(&arrays[scalar_arr_idx], idx_remaining)?;
4517        if scalar.null_count() != 0 {
4518            return Ok(None);
4519        }
4520        if !Self::is_constant_values(arrays, &scalar, validity)? {
4521            return Ok(None);
4522        }
4523        Ok(Some(scalar))
4524    }
4525
4526    fn resolve_dict_values_compression_metadata(
4527        field_metadata: &HashMap<String, String>,
4528        env_compression: Option<String>,
4529        env_compression_level: Option<String>,
4530    ) -> HashMap<String, String> {
4531        let mut metadata = HashMap::new();
4532
4533        let compression = field_metadata
4534            .get(DICT_VALUES_COMPRESSION_META_KEY)
4535            .cloned()
4536            .or(env_compression)
4537            .unwrap_or_else(|| DEFAULT_DICT_VALUES_COMPRESSION.to_string());
4538        metadata.insert(COMPRESSION_META_KEY.to_string(), compression);
4539
4540        if let Some(compression_level) = field_metadata
4541            .get(DICT_VALUES_COMPRESSION_LEVEL_META_KEY)
4542            .cloned()
4543            .or(env_compression_level)
4544        {
4545            metadata.insert(COMPRESSION_LEVEL_META_KEY.to_string(), compression_level);
4546        }
4547
4548        metadata
4549    }
4550
4551    fn build_dict_values_compressor_field(field: &Field) -> Result<Field> {
4552        // This is an internal synthetic field used only to feed metadata into
4553        // `create_block_compressor` for dictionary values. The concrete type/name here
4554        // are not semantically meaningful; we rely on explicit metadata below to control
4555        // general compression selection for dictionary values.
4556        let mut dict_values_field = Field::new_arrow("", DataType::UInt16, false)?;
4557        dict_values_field.metadata = Self::resolve_dict_values_compression_metadata(
4558            &field.metadata,
4559            env::var(DICT_VALUES_COMPRESSION_ENV_VAR).ok(),
4560            env::var(DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR).ok(),
4561        );
4562        Ok(dict_values_field)
4563    }
4564
4565    #[allow(clippy::too_many_arguments)]
4566    fn encode_miniblock(
4567        column_idx: u32,
4568        field: &Field,
4569        compression_strategy: &dyn CompressionStrategy,
4570        data: DataBlock,
4571        repdef: crate::repdef::SerializedRepDefs,
4572        row_number: u64,
4573        dictionary_data: Option<DataBlock>,
4574        num_rows: u64,
4575        support_large_chunk: bool,
4576    ) -> Result<EncodedPage> {
4577        if let DataBlock::AllNull(_null_block) = data {
4578            // We should not be using mini-block for all-null.  There are other structural
4579            // encodings for that.
4580            unreachable!()
4581        }
4582
4583        let num_items = data.num_values();
4584
4585        let compressor = compression_strategy.create_miniblock_compressor(field, &data)?;
4586        let (compressed_data, value_encoding) = compressor.compress(data)?;
4587
4588        let max_rep = repdef.def_meaning.iter().filter(|l| l.is_list()).count() as u16;
4589
4590        let mut compressed_rep = repdef
4591            .rep_slicer()
4592            .map(|rep_slicer| {
4593                Self::compress_levels(
4594                    rep_slicer,
4595                    num_items,
4596                    compression_strategy,
4597                    &compressed_data.chunks,
4598                    max_rep,
4599                )
4600            })
4601            .transpose()?;
4602
4603        let (rep_index, rep_index_depth) =
4604            match compressed_rep.as_mut().and_then(|cr| cr.rep_index.as_mut()) {
4605                Some(rep_index) => (Some(rep_index.clone()), 1),
4606                None => (None, 0),
4607            };
4608
4609        let mut compressed_def = repdef
4610            .def_slicer()
4611            .map(|def_slicer| {
4612                Self::compress_levels(
4613                    def_slicer,
4614                    num_items,
4615                    compression_strategy,
4616                    &compressed_data.chunks,
4617                    /*max_rep=*/ 0,
4618                )
4619            })
4620            .transpose()?;
4621
4622        // TODO: Parquet sparsely encodes values here.  We could do the same but
4623        // then we won't have log2 values per chunk.  This means more metadata
4624        // and potentially more decoder asymmetry.  However, it may be worth
4625        // investigating at some point
4626
4627        let rep_data = compressed_rep
4628            .as_mut()
4629            .map(|cr| std::mem::take(&mut cr.data));
4630        let def_data = compressed_def
4631            .as_mut()
4632            .map(|cd| std::mem::take(&mut cd.data));
4633
4634        let serialized =
4635            Self::serialize_miniblocks(compressed_data, rep_data, def_data, support_large_chunk)?;
4636
4637        // Metadata, Data, Dictionary, (maybe) Repetition Index
4638        let mut data = Vec::with_capacity(4);
4639        data.push(serialized.metadata);
4640        data.push(serialized.data);
4641
4642        if let Some(dictionary_data) = dictionary_data {
4643            let num_dictionary_items = dictionary_data.num_values();
4644            let dict_values_field = Self::build_dict_values_compressor_field(field)?;
4645
4646            let (compressor, dictionary_encoding) = compression_strategy
4647                .create_block_compressor(&dict_values_field, &dictionary_data)?;
4648            let dictionary_buffer = compressor.compress(dictionary_data)?;
4649
4650            data.push(dictionary_buffer);
4651            if let Some(rep_index) = rep_index {
4652                data.push(rep_index);
4653            }
4654
4655            let description = ProtobufUtils21::miniblock_layout(
4656                compressed_rep.map(|cr| cr.compression),
4657                compressed_def.map(|cd| cd.compression),
4658                value_encoding,
4659                rep_index_depth,
4660                serialized.num_buffers,
4661                Some((dictionary_encoding, num_dictionary_items)),
4662                &repdef.def_meaning,
4663                num_items,
4664                support_large_chunk,
4665            );
4666            Ok(EncodedPage {
4667                num_rows,
4668                column_idx,
4669                data,
4670                description: PageEncoding::Structural(description),
4671                row_number,
4672            })
4673        } else {
4674            let description = ProtobufUtils21::miniblock_layout(
4675                compressed_rep.map(|cr| cr.compression),
4676                compressed_def.map(|cd| cd.compression),
4677                value_encoding,
4678                rep_index_depth,
4679                serialized.num_buffers,
4680                None,
4681                &repdef.def_meaning,
4682                num_items,
4683                support_large_chunk,
4684            );
4685
4686            if let Some(rep_index) = rep_index {
4687                let view = rep_index.borrow_to_typed_slice::<u64>();
4688                let total = view.chunks_exact(2).map(|c| c[0]).sum::<u64>();
4689                debug_assert_eq!(total, num_rows);
4690
4691                data.push(rep_index);
4692            }
4693
4694            Ok(EncodedPage {
4695                num_rows,
4696                column_idx,
4697                data,
4698                description: PageEncoding::Structural(description),
4699                row_number,
4700            })
4701        }
4702    }
4703
4704    // For fixed-size data we encode < control word | data > for each value
4705    fn serialize_full_zip_fixed(
4706        fixed: FixedWidthDataBlock,
4707        mut repdef: ControlWordIterator,
4708        num_values: u64,
4709    ) -> Result<SerializedFullZip> {
4710        if !fixed.bits_per_value.is_multiple_of(8) {
4711            return Err(Error::invalid_input_source(
4712                format!(
4713                    "Full-zip fixed-width values must be byte aligned, got {} bits per value",
4714                    fixed.bits_per_value
4715                )
4716                .into(),
4717            ));
4718        }
4719
4720        let len = fixed.data.len() + repdef.bytes_per_word() * num_values as usize;
4721        let mut zipped_data = Vec::with_capacity(len);
4722
4723        let max_rep_index_val = if repdef.has_repetition() {
4724            len as u64
4725        } else {
4726            // Setting this to 0 means we won't write a repetition index
4727            0
4728        };
4729        let mut rep_index_builder =
4730            BytepackedIntegerEncoder::with_capacity(num_values as usize + 1, max_rep_index_val);
4731
4732        let bytes_per_value = fixed.bits_per_value as usize / 8;
4733        let mut offset = 0;
4734
4735        if bytes_per_value == 0 {
4736            // No data, just dump the repdef into the buffer
4737            while let Some(control) = repdef.append_next(&mut zipped_data) {
4738                if control.is_new_row {
4739                    // We have finished a row
4740                    debug_assert!(offset <= len);
4741                    // SAFETY: We know that `start <= len`
4742                    unsafe { rep_index_builder.append(offset as u64) };
4743                }
4744                offset = zipped_data.len();
4745            }
4746        } else {
4747            // We have data, zip it with the repdef
4748            let mut data_iter = fixed.data.chunks_exact(bytes_per_value);
4749            while let Some(control) = repdef.append_next(&mut zipped_data) {
4750                if control.is_new_row {
4751                    // We have finished a row
4752                    debug_assert!(offset <= len);
4753                    // SAFETY: We know that `start <= len`
4754                    unsafe { rep_index_builder.append(offset as u64) };
4755                }
4756                if control.is_visible {
4757                    let value = data_iter.next().unwrap();
4758                    zipped_data.extend_from_slice(value);
4759                }
4760                offset = zipped_data.len();
4761            }
4762        }
4763
4764        debug_assert_eq!(zipped_data.len(), len);
4765        // Put the final value in the rep index
4766        // SAFETY: `zipped_data.len() == len`
4767        unsafe {
4768            rep_index_builder.append(zipped_data.len() as u64);
4769        }
4770
4771        let zipped_data = LanceBuffer::from(zipped_data);
4772        let rep_index = rep_index_builder.into_data();
4773        let rep_index = if rep_index.is_empty() {
4774            None
4775        } else {
4776            Some(LanceBuffer::from(rep_index))
4777        };
4778        Ok(SerializedFullZip {
4779            values: zipped_data,
4780            repetition_index: rep_index,
4781        })
4782    }
4783
4784    // For variable-size data we encode < control word | length | data > for each value
4785    //
4786    // In addition, we create a second buffer, the repetition index
4787    fn serialize_full_zip_variable(
4788        variable: VariableWidthBlock,
4789        mut repdef: ControlWordIterator,
4790        num_items: u64,
4791    ) -> Result<SerializedFullZip> {
4792        let bytes_per_offset = variable.bits_per_offset as usize / 8;
4793        if !variable.bits_per_offset.is_multiple_of(8) {
4794            return Err(Error::invalid_input_source(
4795                format!(
4796                    "Full-zip variable-width offsets must be byte aligned, got {} bits per offset",
4797                    variable.bits_per_offset
4798                )
4799                .into(),
4800            ));
4801        }
4802        let len = variable.data.len()
4803            + repdef.bytes_per_word() * num_items as usize
4804            + bytes_per_offset * variable.num_values as usize;
4805        let mut buf = Vec::with_capacity(len);
4806
4807        let max_rep_index_val = len as u64;
4808        let mut rep_index_builder =
4809            BytepackedIntegerEncoder::with_capacity(num_items as usize + 1, max_rep_index_val);
4810
4811        // TODO: byte pack the item lengths with varint encoding
4812        match bytes_per_offset {
4813            4 => {
4814                let offs = variable.offsets.borrow_to_typed_slice::<u32>();
4815                let mut rep_offset = 0;
4816                let mut windows_iter = offs.as_ref().windows(2);
4817                while let Some(control) = repdef.append_next(&mut buf) {
4818                    if control.is_new_row {
4819                        // We have finished a row
4820                        debug_assert!(rep_offset <= len);
4821                        // SAFETY: We know that `buf.len() <= len`
4822                        unsafe { rep_index_builder.append(rep_offset as u64) };
4823                    }
4824                    if control.is_visible {
4825                        let window = windows_iter.next().unwrap();
4826                        if control.is_valid_item {
4827                            buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
4828                            buf.extend_from_slice(
4829                                &variable.data[window[0] as usize..window[1] as usize],
4830                            );
4831                        }
4832                    }
4833                    rep_offset = buf.len();
4834                }
4835            }
4836            8 => {
4837                let offs = variable.offsets.borrow_to_typed_slice::<u64>();
4838                let mut rep_offset = 0;
4839                let mut windows_iter = offs.as_ref().windows(2);
4840                while let Some(control) = repdef.append_next(&mut buf) {
4841                    if control.is_new_row {
4842                        // We have finished a row
4843                        debug_assert!(rep_offset <= len);
4844                        // SAFETY: We know that `buf.len() <= len`
4845                        unsafe { rep_index_builder.append(rep_offset as u64) };
4846                    }
4847                    if control.is_visible {
4848                        let window = windows_iter.next().unwrap();
4849                        if control.is_valid_item {
4850                            buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
4851                            buf.extend_from_slice(
4852                                &variable.data[window[0] as usize..window[1] as usize],
4853                            );
4854                        }
4855                    }
4856                    rep_offset = buf.len();
4857                }
4858            }
4859            _ => {
4860                return Err(Error::invalid_input_source(
4861                    format!(
4862                        "Full-zip variable-width offsets must be 32 or 64 bits, got {} bits",
4863                        variable.bits_per_offset
4864                    )
4865                    .into(),
4866                ));
4867            }
4868        }
4869
4870        // We might have saved a few bytes by not copying lengths when the length was zero.  However,
4871        // if we are over `len` then we have a bug.
4872        debug_assert!(buf.len() <= len);
4873        // Put the final value in the rep index
4874        // SAFETY: `zipped_data.len() == len`
4875        unsafe {
4876            rep_index_builder.append(buf.len() as u64);
4877        }
4878
4879        let zipped_data = LanceBuffer::from(buf);
4880        let rep_index = rep_index_builder.into_data();
4881        debug_assert!(!rep_index.is_empty());
4882        let rep_index = Some(LanceBuffer::from(rep_index));
4883        Ok(SerializedFullZip {
4884            values: zipped_data,
4885            repetition_index: rep_index,
4886        })
4887    }
4888
4889    /// Serializes data into a single buffer according to the full-zip format which zips
4890    /// together the repetition, definition, and value data into a single buffer.
4891    fn serialize_full_zip(
4892        compressed_data: PerValueDataBlock,
4893        repdef: ControlWordIterator,
4894        num_items: u64,
4895    ) -> Result<SerializedFullZip> {
4896        match compressed_data {
4897            PerValueDataBlock::Fixed(fixed) => {
4898                Self::serialize_full_zip_fixed(fixed, repdef, num_items)
4899            }
4900            PerValueDataBlock::Variable(var) => {
4901                Self::serialize_full_zip_variable(var, repdef, num_items)
4902            }
4903        }
4904    }
4905
4906    fn expand_boolean_to_bytes(fixed: FixedWidthDataBlock) -> FixedWidthDataBlock {
4907        debug_assert_eq!(fixed.bits_per_value, 1);
4908        let num_values = fixed.num_values as usize;
4909        let bool_buf = BooleanBuffer::new(fixed.data.into_buffer(), 0, num_values);
4910        let expanded: Vec<u8> = (0..num_values).map(|i| bool_buf.value(i) as u8).collect();
4911        FixedWidthDataBlock {
4912            data: LanceBuffer::from(expanded),
4913            bits_per_value: 8,
4914            num_values: fixed.num_values,
4915            block_info: BlockInfo::new(),
4916        }
4917    }
4918
4919    fn encode_full_zip(
4920        column_idx: u32,
4921        field: &Field,
4922        compression_strategy: &dyn CompressionStrategy,
4923        data: DataBlock,
4924        repdef: crate::repdef::SerializedRepDefs,
4925        row_number: u64,
4926        num_lists: u64,
4927    ) -> Result<EncodedPage> {
4928        let max_rep = repdef
4929            .repetition_levels
4930            .as_ref()
4931            .map_or(0, |r| r.iter().max().copied().unwrap_or(0));
4932        let max_def = repdef
4933            .definition_levels
4934            .as_ref()
4935            .map_or(0, |d| d.iter().max().copied().unwrap_or(0));
4936
4937        // To handle FSL we just flatten
4938        // let data = data.flatten();
4939
4940        let (num_items, num_visible_items) =
4941            if let Some(rep_levels) = repdef.repetition_levels.as_ref() {
4942                // If there are rep levels there may be "invisible" items and we need to encode
4943                // rep_levels.len() things which might be larger than data.num_values()
4944                (rep_levels.len() as u64, data.num_values())
4945            } else {
4946                // If there are no rep levels then we encode data.num_values() things
4947                (data.num_values(), data.num_values())
4948            };
4949
4950        let max_visible_def = repdef.max_visible_level.unwrap_or(u16::MAX);
4951
4952        let repdef_iter = build_control_word_iterator(
4953            repdef.repetition_levels.as_deref(),
4954            max_rep,
4955            repdef.definition_levels.as_deref(),
4956            max_def,
4957            max_visible_def,
4958            num_items as usize,
4959        );
4960        let bits_rep = repdef_iter.bits_rep();
4961        let bits_def = repdef_iter.bits_def();
4962
4963        // Full-zip requires byte-aligned values; expand 1-bit booleans to 1 byte each.
4964        let data = match data {
4965            DataBlock::FixedWidth(fixed) if fixed.bits_per_value == 1 => {
4966                DataBlock::FixedWidth(Self::expand_boolean_to_bytes(fixed))
4967            }
4968            other => other,
4969        };
4970
4971        let compressor = compression_strategy.create_per_value(field, &data)?;
4972        let (compressed_data, value_encoding) = compressor.compress(data)?;
4973
4974        let description = match &compressed_data {
4975            PerValueDataBlock::Fixed(fixed) => ProtobufUtils21::fixed_full_zip_layout(
4976                bits_rep,
4977                bits_def,
4978                fixed.bits_per_value as u32,
4979                value_encoding,
4980                &repdef.def_meaning,
4981                num_items as u32,
4982                num_visible_items as u32,
4983            ),
4984            PerValueDataBlock::Variable(variable) => ProtobufUtils21::variable_full_zip_layout(
4985                bits_rep,
4986                bits_def,
4987                variable.bits_per_offset as u32,
4988                value_encoding,
4989                &repdef.def_meaning,
4990                num_items as u32,
4991                num_visible_items as u32,
4992            ),
4993        };
4994
4995        let zipped = Self::serialize_full_zip(compressed_data, repdef_iter, num_items)?;
4996
4997        let data = if let Some(repindex) = zipped.repetition_index {
4998            vec![zipped.values, repindex]
4999        } else {
5000            vec![zipped.values]
5001        };
5002
5003        Ok(EncodedPage {
5004            num_rows: num_lists,
5005            column_idx,
5006            data,
5007            description: PageEncoding::Structural(description),
5008            row_number,
5009        })
5010    }
5011
5012    fn should_dictionary_encode(
5013        data_block: &DataBlock,
5014        field: &Field,
5015        version: LanceFileVersion,
5016    ) -> Option<DictEncodingBudget> {
5017        const DEFAULT_SAMPLE_SIZE: usize = 4096;
5018        const DEFAULT_SAMPLE_UNIQUE_RATIO: f64 = 0.98;
5019
5020        // Since we only dictionary encode FixedWidth and VariableWidth blocks for now, we skip
5021        // estimating the size for other types.
5022        match data_block {
5023            DataBlock::FixedWidth(fixed) => {
5024                if fixed.bits_per_value == 64 && version < LanceFileVersion::V2_2 {
5025                    return None;
5026                }
5027                if fixed.bits_per_value != 64 && fixed.bits_per_value != 128 {
5028                    return None;
5029                }
5030                if fixed.bits_per_value % 8 != 0 {
5031                    return None;
5032                }
5033            }
5034            DataBlock::VariableWidth(var) => {
5035                if var.bits_per_offset != 32 && var.bits_per_offset != 64 {
5036                    return None;
5037                }
5038            }
5039            _ => return None,
5040        }
5041
5042        // Don't dictionary encode tiny arrays.
5043        let too_small = env::var("LANCE_ENCODING_DICT_TOO_SMALL")
5044            .ok()
5045            .and_then(|val| val.parse().ok())
5046            .unwrap_or(100);
5047        if data_block.num_values() < too_small {
5048            return None;
5049        }
5050
5051        let num_values = data_block.num_values();
5052
5053        // Apply divisor threshold and cap. This is intentionally conservative: the goal is to
5054        // avoid spending too much CPU trying to estimate very high cardinalities.
5055        let divisor: u64 = field
5056            .metadata
5057            .get(DICT_DIVISOR_META_KEY)
5058            .and_then(|val| val.parse().ok())
5059            .or_else(|| {
5060                env::var("LANCE_ENCODING_DICT_DIVISOR")
5061                    .ok()
5062                    .and_then(|val| val.parse().ok())
5063            })
5064            .unwrap_or(DEFAULT_DICT_DIVISOR);
5065
5066        let max_cardinality: u64 = env::var("LANCE_ENCODING_DICT_MAX_CARDINALITY")
5067            .ok()
5068            .and_then(|val| val.parse().ok())
5069            .unwrap_or(DEFAULT_DICT_MAX_CARDINALITY);
5070
5071        let threshold_cardinality = num_values
5072            .checked_div(divisor.max(1))
5073            .unwrap_or(0)
5074            .min(max_cardinality);
5075        if threshold_cardinality == 0 {
5076            return None;
5077        }
5078
5079        // Get size ratio from metadata or env var.
5080        let threshold_ratio = field
5081            .metadata
5082            .get(DICT_SIZE_RATIO_META_KEY)
5083            .and_then(|val| val.parse::<f64>().ok())
5084            .or_else(|| {
5085                env::var("LANCE_ENCODING_DICT_SIZE_RATIO")
5086                    .ok()
5087                    .and_then(|val| val.parse().ok())
5088            })
5089            .unwrap_or(DEFAULT_DICT_SIZE_RATIO);
5090
5091        if threshold_ratio <= 0.0 || threshold_ratio > 1.0 {
5092            panic!(
5093                "Invalid parameter: dict-size-ratio is {} which is not in the range (0, 1].",
5094                threshold_ratio
5095            );
5096        }
5097
5098        let data_size = data_block.data_size();
5099        if data_size == 0 {
5100            return None;
5101        }
5102
5103        let max_encoded_size = (data_size as f64 * threshold_ratio) as u64;
5104        let max_encoded_size = usize::try_from(max_encoded_size).ok()?;
5105
5106        // Avoid probing dictionary encoding on data that appears to be near-unique
5107        // or likely to exceed the dictionary budget.
5108        if let Some(sample_unique_ratio) =
5109            Self::sample_unique_ratio(data_block, DEFAULT_SAMPLE_SIZE)?
5110        {
5111            if sample_unique_ratio >= DEFAULT_SAMPLE_UNIQUE_RATIO {
5112                return None;
5113            }
5114
5115            let projected_cardinality = (sample_unique_ratio * num_values as f64).ceil() as u64;
5116            if projected_cardinality > threshold_cardinality {
5117                return None;
5118            }
5119        }
5120
5121        let max_dict_entries = u32::try_from(threshold_cardinality.min(i32::MAX as u64)).ok()?;
5122        Some(DictEncodingBudget {
5123            max_dict_entries,
5124            max_encoded_size,
5125        })
5126    }
5127
5128    /// Samples whether a page looks near-unique before attempting dictionary encoding.
5129    ///
5130    /// The probe uses deterministic block sampling (not RNG sampling), which keeps
5131    /// the check cheap and reproducible across runs. The result is only a gate for
5132    /// whether we try dictionary encoding, not a cardinality statistic.
5133    /// Returns `Some(None)` when there are too few reliable samples or the block type does not
5134    /// support dictionary encoding. Returns `None` for malformed data.
5135    fn sample_unique_ratio(data_block: &DataBlock, max_samples: usize) -> Option<Option<f64>> {
5136        use std::collections::HashSet;
5137
5138        const NUM_SAMPLE_BLOCKS: usize = 32;
5139        const MIN_RELIABLE_SAMPLES: usize = 1024;
5140
5141        let num_values = usize::try_from(data_block.num_values()).ok()?;
5142        if num_values == 0 {
5143            return Some(None);
5144        }
5145
5146        let sample_count = num_values.min(max_samples).max(1);
5147        if sample_count < MIN_RELIABLE_SAMPLES {
5148            return Some(None);
5149        }
5150
5151        let block_count = NUM_SAMPLE_BLOCKS.min(sample_count).min(num_values).max(1);
5152        let samples_per_block = (sample_count / block_count).max(1);
5153        let mut indices = Vec::with_capacity(sample_count);
5154        for block_idx in 0..block_count {
5155            let block_start = block_idx * num_values / block_count;
5156            let next_block_start = ((block_idx + 1) * num_values / block_count).min(num_values);
5157            let block_len = next_block_start.saturating_sub(block_start);
5158            let samples_in_block = samples_per_block.min(block_len);
5159            indices.extend((0..samples_in_block).map(|offset| block_start + offset));
5160        }
5161
5162        if indices.len() < MIN_RELIABLE_SAMPLES {
5163            return Some(None);
5164        }
5165
5166        let ratio = match data_block {
5167            DataBlock::FixedWidth(fixed) => match fixed.bits_per_value {
5168                64 => {
5169                    let values = fixed.data.borrow_to_typed_slice::<u64>();
5170                    let values = values.as_ref();
5171                    let mut unique: HashSet<u64> =
5172                        HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
5173                    for idx in indices.iter().copied() {
5174                        unique.insert(values.get(idx).copied()?);
5175                    }
5176                    unique.len() as f64 / indices.len() as f64
5177                }
5178                128 => {
5179                    let values = fixed.data.borrow_to_typed_slice::<u128>();
5180                    let values = values.as_ref();
5181                    let mut unique: HashSet<u128> =
5182                        HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
5183                    for idx in indices.iter().copied() {
5184                        unique.insert(values.get(idx).copied()?);
5185                    }
5186                    unique.len() as f64 / indices.len() as f64
5187                }
5188                _ => return Some(None),
5189            },
5190            DataBlock::VariableWidth(var) => {
5191                use xxhash_rust::xxh3::xxh3_64;
5192
5193                // Hash variable-width slices instead of storing borrowed slice keys.
5194                let mut unique: HashSet<u64> =
5195                    HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
5196                match var.bits_per_offset {
5197                    32 => {
5198                        let offsets_ref = var.offsets.borrow_to_typed_slice::<u32>();
5199                        let offsets: &[u32] = offsets_ref.as_ref();
5200                        for i in indices.iter().copied() {
5201                            let start = usize::try_from(*offsets.get(i)?).ok()?;
5202                            let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
5203                            if start > end || end > var.data.len() {
5204                                return None;
5205                            }
5206                            unique.insert(xxh3_64(&var.data[start..end]));
5207                        }
5208                    }
5209                    64 => {
5210                        let offsets_ref = var.offsets.borrow_to_typed_slice::<u64>();
5211                        let offsets: &[u64] = offsets_ref.as_ref();
5212                        for i in indices.iter().copied() {
5213                            let start = usize::try_from(*offsets.get(i)?).ok()?;
5214                            let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
5215                            if start > end || end > var.data.len() {
5216                                return None;
5217                            }
5218                            unique.insert(xxh3_64(&var.data[start..end]));
5219                        }
5220                    }
5221                    _ => return Some(None),
5222                }
5223                unique.len() as f64 / indices.len() as f64
5224            }
5225            _ => return Some(None),
5226        };
5227
5228        Some(Some(ratio))
5229    }
5230
5231    fn slice_repdef(repdef: &SerializedRepDefs, range: Range<usize>) -> SerializedRepDefs {
5232        let repetition_levels = repdef
5233            .repetition_levels
5234            .as_ref()
5235            .map(|levels| levels[range.clone()].to_vec());
5236        let definition_levels = repdef
5237            .definition_levels
5238            .as_ref()
5239            .map(|levels| levels[range].to_vec());
5240        SerializedRepDefs::new_with_fixed_size_list_levels(
5241            repetition_levels,
5242            definition_levels,
5243            repdef.def_meaning.clone(),
5244            repdef.has_fixed_size_list_levels(),
5245        )
5246    }
5247
5248    fn slice_arrays(
5249        arrays: &[ArrayRef],
5250        value_start: u64,
5251        num_values: u64,
5252    ) -> Result<Vec<ArrayRef>> {
5253        if num_values == 0 {
5254            return Ok(Vec::new());
5255        }
5256
5257        let mut values_to_skip = usize::try_from(value_start).map_err(|_| {
5258            Error::invalid_input(format!("Value start {} is too large", value_start))
5259        })?;
5260        let mut values_remaining = usize::try_from(num_values).map_err(|_| {
5261            Error::invalid_input(format!("Value count {} is too large", num_values))
5262        })?;
5263        let mut sliced = Vec::new();
5264
5265        for array in arrays {
5266            if values_to_skip >= array.len() {
5267                values_to_skip -= array.len();
5268                continue;
5269            }
5270
5271            let offset = values_to_skip;
5272            let len = (array.len() - offset).min(values_remaining);
5273            sliced.push(array.slice(offset, len));
5274            values_remaining -= len;
5275            values_to_skip = 0;
5276
5277            if values_remaining == 0 {
5278                break;
5279            }
5280        }
5281
5282        if values_remaining != 0 {
5283            return Err(Error::internal(format!(
5284                "Page split requested {} values starting at {}, but the page did not contain enough values",
5285                num_values, value_start
5286            )));
5287        }
5288
5289        Ok(sliced)
5290    }
5291
5292    fn split_structural_pages_for_miniblock_budget(
5293        arrays: Vec<ArrayRef>,
5294        repdef: SerializedRepDefs,
5295        plan: StructuralPagePlan,
5296        row_number: u64,
5297        num_rows: u64,
5298    ) -> Result<Vec<PrimitivePageData>> {
5299        if plan == StructuralPagePlan::Fits {
5300            return Ok(vec![PrimitivePageData {
5301                arrays,
5302                repdef,
5303                row_number,
5304                num_rows,
5305                unsplittable_miniblock_levels: None,
5306            }]);
5307        }
5308        if let StructuralPagePlan::UnsplittableOverBudget(num_levels) = plan {
5309            return Ok(vec![PrimitivePageData {
5310                arrays,
5311                repdef,
5312                row_number,
5313                num_rows,
5314                unsplittable_miniblock_levels: Some(num_levels),
5315            }]);
5316        }
5317
5318        let StructuralPagePlan::Split(splits) = plan else {
5319            unreachable!();
5320        };
5321
5322        let mut pages = Vec::with_capacity(splits.len());
5323        for split in splits {
5324            let arrays = Self::slice_arrays(&arrays, split.value_start, split.num_values)?;
5325            let repdef = Self::slice_repdef(&repdef, split.level_range);
5326            pages.push(PrimitivePageData {
5327                arrays,
5328                repdef,
5329                row_number: row_number + split.row_start,
5330                num_rows: split.num_rows,
5331                unsplittable_miniblock_levels: None,
5332            });
5333        }
5334        Ok(pages)
5335    }
5336
5337    fn encode_page(ctx: PrimitiveEncodeContext, page: PrimitivePageData) -> Result<EncodedPage> {
5338        let PrimitiveEncodeContext {
5339            column_idx,
5340            field,
5341            compression_strategy,
5342            encoding_metadata,
5343            support_large_chunk,
5344            version,
5345            is_simple_validity,
5346            has_repdef_info,
5347        } = ctx;
5348        let PrimitivePageData {
5349            arrays,
5350            repdef,
5351            row_number,
5352            num_rows,
5353            unsplittable_miniblock_levels,
5354        } = page;
5355        let num_values = arrays.iter().map(|arr| arr.len() as u64).sum();
5356
5357        if num_values == 0 {
5358            // This page contains only structural events, such as empty/null list rows.
5359            // The existing complex-null layout stores the rep/def stream without value buffers.
5360            log::debug!(
5361                "Encoding column {} with {} items ({} rows) using complex-null layout",
5362                column_idx,
5363                num_values,
5364                num_rows
5365            );
5366            return Self::encode_complex_all_null(
5367                column_idx,
5368                repdef,
5369                row_number,
5370                num_rows,
5371                version,
5372                compression_strategy.as_ref(),
5373            );
5374        }
5375
5376        let leaf_validity = Self::leaf_validity(&repdef, num_values as usize)?;
5377        let all_null = leaf_validity
5378            .as_ref()
5379            .map(|validity| validity.count_set_bits() == 0)
5380            .unwrap_or(false);
5381
5382        if all_null {
5383            return if is_simple_validity {
5384                log::debug!(
5385                    "Encoding column {} with {} items ({} rows) using simple-null layout",
5386                    column_idx,
5387                    num_values,
5388                    num_rows
5389                );
5390                Self::encode_simple_all_null(column_idx, num_values, row_number)
5391            } else {
5392                log::debug!(
5393                    "Encoding column {} with {} items ({} rows) using complex-null layout",
5394                    column_idx,
5395                    num_values,
5396                    num_rows
5397                );
5398                Self::encode_complex_all_null(
5399                    column_idx,
5400                    repdef,
5401                    row_number,
5402                    num_rows,
5403                    version,
5404                    compression_strategy.as_ref(),
5405                )
5406            };
5407        }
5408
5409        if let DataType::Struct(fields) = &field.data_type()
5410            && fields.is_empty()
5411        {
5412            if has_repdef_info {
5413                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()));
5414            }
5415            // This is maybe a little confusing but the reader should never look at this anyways and it
5416            // seems like overkill to invent a new layout just for "empty structs".
5417            return Self::encode_simple_all_null(column_idx, num_values, row_number);
5418        }
5419
5420        let data_block = DataBlock::from_arrays(&arrays, num_values);
5421
5422        if version.resolve() >= LanceFileVersion::V2_2
5423            && let Some(scalar) = Self::find_constant_scalar(&arrays, leaf_validity.as_ref())?
5424        {
5425            log::debug!(
5426                "Encoding column {} with {} items ({} rows) using constant layout",
5427                column_idx,
5428                num_values,
5429                num_rows
5430            );
5431            return constant::encode_constant_page(
5432                column_idx, scalar, repdef, row_number, num_rows,
5433            );
5434        }
5435
5436        if let Some(num_levels) = unsplittable_miniblock_levels {
5437            let requested_encoding = encoding_metadata
5438                .get(STRUCTURAL_ENCODING_META_KEY)
5439                .map(|requested| requested.to_lowercase());
5440            let fullzip_error = match &data_block {
5441                DataBlock::FixedWidth(fixed) if !fixed.bits_per_value.is_multiple_of(8) => {
5442                    Some(format!(
5443                        "Full-zip fixed-width values must be byte aligned, got {} bits per value",
5444                        fixed.bits_per_value
5445                    ))
5446                }
5447                DataBlock::VariableWidth(variable)
5448                    if !variable.bits_per_offset.is_multiple_of(8) =>
5449                {
5450                    Some(format!(
5451                        "Full-zip variable-width offsets must be byte aligned, got {} bits per offset",
5452                        variable.bits_per_offset
5453                    ))
5454                }
5455                DataBlock::VariableWidth(variable)
5456                    if variable.bits_per_offset != 32 && variable.bits_per_offset != 64 =>
5457                {
5458                    Some(format!(
5459                        "Full-zip variable-width offsets must be 32 or 64 bits, got {} bits",
5460                        variable.bits_per_offset
5461                    ))
5462                }
5463                DataBlock::Struct(struct_data_block)
5464                    if !struct_data_block.has_variable_width_child() =>
5465                {
5466                    Some(
5467                        "Full-zip packed struct requires at least one variable-width child"
5468                            .to_string(),
5469                    )
5470                }
5471                DataBlock::Dictionary(_) => {
5472                    Some("Full-zip does not encode dictionary data blocks directly".to_string())
5473                }
5474                DataBlock::FixedSizeList(fsl) => match fsl.clone().try_into_flat() {
5475                    Some(flat) if flat.bits_per_value.is_multiple_of(8) => None,
5476                    Some(flat) => Some(format!(
5477                        "Full-zip fixed-size-list values must be byte aligned after flattening, got {} bits per value",
5478                        flat.bits_per_value
5479                    )),
5480                    None => Some(
5481                        "Full-zip fixed-size-list capability requires a flat fixed-width child"
5482                            .to_string(),
5483                    ),
5484                },
5485                DataBlock::FixedWidth(_) | DataBlock::VariableWidth(_) | DataBlock::Struct(_) => {
5486                    None
5487                }
5488                other => Some(format!(
5489                    "Full-zip does not support value block type {}",
5490                    other.name()
5491                )),
5492            };
5493            match requested_encoding.as_deref() {
5494                Some(STRUCTURAL_ENCODING_FULLZIP) => {
5495                    if let Some(reason) = fullzip_error {
5496                        return Err(Error::invalid_input_source(reason.into()));
5497                    }
5498                    return Self::encode_full_zip(
5499                        column_idx,
5500                        &field,
5501                        compression_strategy.as_ref(),
5502                        data_block,
5503                        repdef,
5504                        row_number,
5505                        num_rows,
5506                    );
5507                }
5508                Some(STRUCTURAL_ENCODING_MINIBLOCK) | None => {
5509                    if requested_encoding.is_none() && fullzip_error.is_none() {
5510                        log::debug!(
5511                            "Encoding column {} with {} items using full-zip layout because mini-block cannot split the structural page",
5512                            column_idx,
5513                            num_values
5514                        );
5515                        return Self::encode_full_zip(
5516                            column_idx,
5517                            &field,
5518                            compression_strategy.as_ref(),
5519                            data_block,
5520                            repdef,
5521                            row_number,
5522                            num_rows,
5523                        );
5524                    }
5525                    return Err(Error::invalid_input_source(
5526                        format!(
5527                            "Mini-block cannot encode {} rep/def levels in one top-level row. \
5528                             This usually means the row contains too much nested structure \
5529                             for the current layout.",
5530                            num_levels
5531                        )
5532                        .into(),
5533                    ));
5534                }
5535                _ => {}
5536            }
5537        }
5538
5539        let requires_full_zip_packed_struct =
5540            if let DataBlock::Struct(ref struct_data_block) = data_block {
5541                struct_data_block.has_variable_width_child()
5542            } else {
5543                false
5544            };
5545
5546        if requires_full_zip_packed_struct {
5547            log::debug!(
5548                "Encoding column {} with {} items using full-zip packed struct layout",
5549                column_idx,
5550                num_values
5551            );
5552            return Self::encode_full_zip(
5553                column_idx,
5554                &field,
5555                compression_strategy.as_ref(),
5556                data_block,
5557                repdef,
5558                row_number,
5559                num_rows,
5560            );
5561        }
5562
5563        if let DataBlock::Dictionary(dict) = data_block {
5564            log::debug!(
5565                "Encoding column {} with {} items using dictionary encoding (already dictionary encoded)",
5566                column_idx,
5567                num_values
5568            );
5569            let (mut indices_data_block, dictionary_data_block) = dict.into_parts();
5570            // TODO: https://github.com/lancedb/lance/issues/4809
5571            // If we compute stats on dictionary_data_block => panic.
5572            // If we don't compute stats on indices_data_block => panic.
5573            // This is messy.  Don't make me call compute_stat ever.
5574            indices_data_block.compute_stat();
5575            return Self::encode_miniblock(
5576                column_idx,
5577                &field,
5578                compression_strategy.as_ref(),
5579                indices_data_block,
5580                repdef,
5581                row_number,
5582                Some(dictionary_data_block),
5583                num_rows,
5584                support_large_chunk,
5585            );
5586        }
5587
5588        // Try dictionary encoding first if applicable. If encoding aborts, fall back to the
5589        // preferred structural encoding.
5590        let dict_result = Self::should_dictionary_encode(&data_block, &field, version).and_then(|budget| {
5591                log::debug!(
5592                    "Encoding column {} with {} items using dictionary encoding (mini-block layout)",
5593                    column_idx,
5594                    num_values
5595                );
5596                dict::dictionary_encode(
5597                    &data_block,
5598                    budget.max_dict_entries,
5599                    budget.max_encoded_size,
5600                )
5601            });
5602
5603        if let Some((indices_data_block, dictionary_data_block)) = dict_result {
5604            Self::encode_miniblock(
5605                column_idx,
5606                &field,
5607                compression_strategy.as_ref(),
5608                indices_data_block,
5609                repdef,
5610                row_number,
5611                Some(dictionary_data_block),
5612                num_rows,
5613                support_large_chunk,
5614            )
5615        } else if Self::prefers_miniblock(&data_block, encoding_metadata.as_ref()) {
5616            log::debug!(
5617                "Encoding column {} with {} items using mini-block layout",
5618                column_idx,
5619                num_values
5620            );
5621            Self::encode_miniblock(
5622                column_idx,
5623                &field,
5624                compression_strategy.as_ref(),
5625                data_block,
5626                repdef,
5627                row_number,
5628                None,
5629                num_rows,
5630                support_large_chunk,
5631            )
5632        } else if Self::prefers_fullzip(encoding_metadata.as_ref()) {
5633            log::debug!(
5634                "Encoding column {} with {} items using full-zip layout",
5635                column_idx,
5636                num_values
5637            );
5638            Self::encode_full_zip(
5639                column_idx,
5640                &field,
5641                compression_strategy.as_ref(),
5642                data_block,
5643                repdef,
5644                row_number,
5645                num_rows,
5646            )
5647        } else {
5648            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()))
5649        }
5650    }
5651
5652    // Creates encode tasks, consuming all buffered data
5653    fn do_flush(
5654        &mut self,
5655        arrays: Vec<ArrayRef>,
5656        repdefs: Vec<RepDefBuilder>,
5657        row_number: u64,
5658        num_rows: u64,
5659    ) -> Result<Vec<EncodeTask>> {
5660        let num_values = arrays.iter().map(|arr| arr.len() as u64).sum();
5661        let is_simple_validity = repdefs.iter().all(|rd| rd.is_simple_validity());
5662        let has_repdef_info = repdefs.iter().any(|rd| !rd.is_empty());
5663        let (repdef, structural_plan) = RepDefBuilder::serialize_with_structural_plan(
5664            repdefs,
5665            miniblock::max_repdef_levels_per_chunk,
5666            num_rows,
5667            num_values,
5668        )?;
5669        let pages = Self::split_structural_pages_for_miniblock_budget(
5670            arrays,
5671            repdef,
5672            structural_plan,
5673            row_number,
5674            num_rows,
5675        )?;
5676
5677        let mut tasks = Vec::with_capacity(pages.len());
5678        let ctx = PrimitiveEncodeContext {
5679            column_idx: self.column_index,
5680            field: self.field.clone(),
5681            compression_strategy: self.compression_strategy.clone(),
5682            encoding_metadata: self.encoding_metadata.clone(),
5683            support_large_chunk: self.support_large_chunk,
5684            version: self.version,
5685            is_simple_validity,
5686            has_repdef_info,
5687        };
5688        for page in pages {
5689            let ctx = ctx.clone();
5690            let task = spawn_cpu(move || Self::encode_page(ctx, page)).boxed();
5691            tasks.push(task);
5692        }
5693        Ok(tasks)
5694    }
5695
5696    fn extract_validity_buf(
5697        array: Arc<dyn Array>,
5698        repdef: &mut RepDefBuilder,
5699        keep_original_array: bool,
5700    ) -> Result<Arc<dyn Array>> {
5701        if let Some(validity) = array.nulls() {
5702            if keep_original_array {
5703                repdef.add_validity_bitmap(validity.clone());
5704            } else {
5705                repdef.add_validity_bitmap(deep_copy_nulls(Some(validity)).unwrap());
5706            }
5707            let data_no_nulls = array.to_data().into_builder().nulls(None).build()?;
5708            Ok(make_array(data_no_nulls))
5709        } else {
5710            repdef.add_no_null(array.len());
5711            Ok(array)
5712        }
5713    }
5714
5715    fn extract_validity(
5716        mut array: Arc<dyn Array>,
5717        repdef: &mut RepDefBuilder,
5718        keep_original_array: bool,
5719    ) -> Result<Arc<dyn Array>> {
5720        match array.data_type() {
5721            DataType::Null => {
5722                repdef.add_validity_bitmap(NullBuffer::new(BooleanBuffer::new_unset(array.len())));
5723                Ok(array)
5724            }
5725            DataType::Dictionary(_, _) => {
5726                array = dict::normalize_dict_nulls(array)?;
5727                Self::extract_validity_buf(array, repdef, keep_original_array)
5728            }
5729            // Extract our validity buf but NOT any child validity bufs. (they will be encoded in
5730            // as part of the values).  Note: for FSL we do not use repdef.add_fsl because we do
5731            // NOT want to increase the repdef depth.
5732            //
5733            // This would be quite catasrophic for something like vector embeddings.  Imagine we
5734            // had thousands of vectors and some were null but no vector contained null items.  If
5735            // we treated the vectors (primitive FSL) like we treat structural FSL we would end up
5736            // with a rep/def value for every single item in the vector.
5737            _ => Self::extract_validity_buf(array, repdef, keep_original_array),
5738        }
5739    }
5740}
5741
5742impl FieldEncoder for PrimitiveStructuralEncoder {
5743    // Buffers data, if there is enough to write a page then we create an encode task
5744    fn maybe_encode(
5745        &mut self,
5746        array: ArrayRef,
5747        _external_buffers: &mut OutOfLineBuffers,
5748        mut repdef: RepDefBuilder,
5749        row_number: u64,
5750        num_rows: u64,
5751    ) -> Result<Vec<EncodeTask>> {
5752        let array = Self::extract_validity(array, &mut repdef, self.keep_original_array)?;
5753        self.accumulated_repdefs.push(repdef);
5754
5755        if let Some((arrays, row_number, num_rows)) =
5756            self.accumulation_queue.insert(array, row_number, num_rows)
5757        {
5758            let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
5759            Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
5760        } else {
5761            Ok(vec![])
5762        }
5763    }
5764
5765    // If there is any data left in the buffer then create an encode task from it
5766    fn flush(&mut self, _external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
5767        if let Some((arrays, row_number, num_rows)) = self.accumulation_queue.flush() {
5768            let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
5769            Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
5770        } else {
5771            Ok(vec![])
5772        }
5773    }
5774
5775    fn num_columns(&self) -> u32 {
5776        1
5777    }
5778
5779    fn finish(
5780        &mut self,
5781        _external_buffers: &mut OutOfLineBuffers,
5782    ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
5783        std::future::ready(Ok(vec![EncodedColumn::default()])).boxed()
5784    }
5785}
5786
5787#[cfg(test)]
5788#[allow(clippy::single_range_in_vec_init)]
5789mod tests {
5790    use super::{
5791        ChunkInstructions, DataBlock, DecodeMiniBlockTask, FixedPerValueDecompressor,
5792        FixedWidthDataBlock, FullZipCacheableState, FullZipDecodeDetails, FullZipReadSource,
5793        FullZipRepIndexDetails, FullZipScheduler, MiniBlockChunk, MiniBlockCompressed,
5794        MiniBlockRepIndex, PerValueDecompressor, PreambleAction, StructuralPageScheduler,
5795        VariableFullZipDecoder,
5796    };
5797    use crate::buffer::LanceBuffer;
5798    use crate::compression::DefaultDecompressionStrategy;
5799    use crate::constants::{
5800        COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_VALUES_COMPRESSION_LEVEL_META_KEY,
5801        DICT_VALUES_COMPRESSION_META_KEY, STRUCTURAL_ENCODING_META_KEY,
5802        STRUCTURAL_ENCODING_MINIBLOCK,
5803    };
5804    use crate::data::BlockInfo;
5805    use crate::decoder::PageEncoding;
5806    use crate::encodings::logical::primitive::{
5807        ChunkDrainInstructions, PrimitiveStructuralEncoder,
5808    };
5809    use crate::format::ProtobufUtils21;
5810    use crate::format::pb21;
5811    use crate::format::pb21::compressive_encoding::Compression;
5812    use crate::repdef::build_control_word_iterator;
5813    use crate::testing::{TestCases, check_round_trip_encoding_of_data};
5814    use crate::version::LanceFileVersion;
5815    use arrow_array::{ArrayRef, Int8Array, StringArray};
5816    use arrow_schema::DataType;
5817    use std::collections::HashMap;
5818    use std::{collections::VecDeque, sync::Arc};
5819
5820    #[test]
5821    fn test_is_narrow() {
5822        let int8_array = Int8Array::from(vec![1, 2, 3]);
5823        let array_ref: ArrayRef = Arc::new(int8_array);
5824        let block = DataBlock::from_array(array_ref);
5825
5826        assert!(PrimitiveStructuralEncoder::is_narrow(&block));
5827
5828        let string_array = StringArray::from(vec![Some("hello"), Some("world")]);
5829        let block = DataBlock::from_array(string_array);
5830        assert!(PrimitiveStructuralEncoder::is_narrow(&block));
5831
5832        let string_array = StringArray::from(vec![
5833            Some("hello world".repeat(100)),
5834            Some("world".to_string()),
5835        ]);
5836        let block = DataBlock::from_array(string_array);
5837        assert!((!PrimitiveStructuralEncoder::is_narrow(&block)));
5838    }
5839
5840    #[test]
5841    fn test_fullzip_fixed_rejects_non_byte_aligned_values() {
5842        let fixed = FixedWidthDataBlock {
5843            data: LanceBuffer::from(vec![0_u8]),
5844            bits_per_value: 1,
5845            num_values: 8,
5846            block_info: BlockInfo::new(),
5847        };
5848        let repdef = build_control_word_iterator(None, 0, None, 0, u16::MAX, 8);
5849
5850        let Err(err) = PrimitiveStructuralEncoder::serialize_full_zip_fixed(fixed, repdef, 8)
5851        else {
5852            panic!("expected full-zip to reject 1-bit fixed-width values");
5853        };
5854        assert!(
5855            err.to_string().contains("byte aligned"),
5856            "unexpected error: {err}"
5857        );
5858    }
5859
5860    #[test]
5861    fn test_map_range() {
5862        // Null in the middle
5863        // [[A, B, C], [D, E], NULL, [F, G, H]]
5864        let rep = Some(vec![1, 0, 0, 1, 0, 1, 1, 0, 0]);
5865        let def = Some(vec![0, 0, 0, 0, 0, 1, 0, 0, 0]);
5866        let max_visible_def = 0;
5867        let total_items = 8;
5868        let max_rep = 1;
5869
5870        let check = |range, expected_item_range, expected_level_range| {
5871            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5872                range,
5873                rep.as_ref(),
5874                def.as_ref(),
5875                max_rep,
5876                max_visible_def,
5877                total_items,
5878                PreambleAction::Absent,
5879            );
5880            assert_eq!(item_range, expected_item_range);
5881            assert_eq!(level_range, expected_level_range);
5882        };
5883
5884        check(0..1, 0..3, 0..3);
5885        check(1..2, 3..5, 3..5);
5886        check(2..3, 5..5, 5..6);
5887        check(3..4, 5..8, 6..9);
5888        check(0..2, 0..5, 0..5);
5889        check(1..3, 3..5, 3..6);
5890        check(2..4, 5..8, 5..9);
5891        check(0..3, 0..5, 0..6);
5892        check(1..4, 3..8, 3..9);
5893        check(0..4, 0..8, 0..9);
5894
5895        // Null at start
5896        // [NULL, [A, B], [C]]
5897        let rep = Some(vec![1, 1, 0, 1]);
5898        let def = Some(vec![1, 0, 0, 0]);
5899        let max_visible_def = 0;
5900        let total_items = 3;
5901
5902        let check = |range, expected_item_range, expected_level_range| {
5903            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5904                range,
5905                rep.as_ref(),
5906                def.as_ref(),
5907                max_rep,
5908                max_visible_def,
5909                total_items,
5910                PreambleAction::Absent,
5911            );
5912            assert_eq!(item_range, expected_item_range);
5913            assert_eq!(level_range, expected_level_range);
5914        };
5915
5916        check(0..1, 0..0, 0..1);
5917        check(1..2, 0..2, 1..3);
5918        check(2..3, 2..3, 3..4);
5919        check(0..2, 0..2, 0..3);
5920        check(1..3, 0..3, 1..4);
5921        check(0..3, 0..3, 0..4);
5922
5923        // Null at end
5924        // [[A], [B, C], NULL]
5925        let rep = Some(vec![1, 1, 0, 1]);
5926        let def = Some(vec![0, 0, 0, 1]);
5927        let max_visible_def = 0;
5928        let total_items = 3;
5929
5930        let check = |range, expected_item_range, expected_level_range| {
5931            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5932                range,
5933                rep.as_ref(),
5934                def.as_ref(),
5935                max_rep,
5936                max_visible_def,
5937                total_items,
5938                PreambleAction::Absent,
5939            );
5940            assert_eq!(item_range, expected_item_range);
5941            assert_eq!(level_range, expected_level_range);
5942        };
5943
5944        check(0..1, 0..1, 0..1);
5945        check(1..2, 1..3, 1..3);
5946        check(2..3, 3..3, 3..4);
5947        check(0..2, 0..3, 0..3);
5948        check(1..3, 1..3, 1..4);
5949        check(0..3, 0..3, 0..4);
5950
5951        // No nulls, with repetition
5952        // [[A, B], [C, D], [E, F]]
5953        let rep = Some(vec![1, 0, 1, 0, 1, 0]);
5954        let def: Option<&[u16]> = None;
5955        let max_visible_def = 0;
5956        let total_items = 6;
5957
5958        let check = |range, expected_item_range, expected_level_range| {
5959            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5960                range,
5961                rep.as_ref(),
5962                def.as_ref(),
5963                max_rep,
5964                max_visible_def,
5965                total_items,
5966                PreambleAction::Absent,
5967            );
5968            assert_eq!(item_range, expected_item_range);
5969            assert_eq!(level_range, expected_level_range);
5970        };
5971
5972        check(0..1, 0..2, 0..2);
5973        check(1..2, 2..4, 2..4);
5974        check(2..3, 4..6, 4..6);
5975        check(0..2, 0..4, 0..4);
5976        check(1..3, 2..6, 2..6);
5977        check(0..3, 0..6, 0..6);
5978
5979        // No repetition, with nulls (this case is trivial)
5980        // [A, B, NULL, C]
5981        let rep: Option<&[u16]> = None;
5982        let def = Some(vec![0, 0, 1, 0]);
5983        let max_visible_def = 1;
5984        let total_items = 4;
5985
5986        let check = |range, expected_item_range, expected_level_range| {
5987            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5988                range,
5989                rep.as_ref(),
5990                def.as_ref(),
5991                max_rep,
5992                max_visible_def,
5993                total_items,
5994                PreambleAction::Absent,
5995            );
5996            assert_eq!(item_range, expected_item_range);
5997            assert_eq!(level_range, expected_level_range);
5998        };
5999
6000        check(0..1, 0..1, 0..1);
6001        check(1..2, 1..2, 1..2);
6002        check(2..3, 2..3, 2..3);
6003        check(0..2, 0..2, 0..2);
6004        check(1..3, 1..3, 1..3);
6005        check(0..3, 0..3, 0..3);
6006
6007        // Tricky case, this chunk is a continuation and starts with a rep-index = 0
6008        // [[..., A] [B, C], NULL]
6009        //
6010        // What we do will depend on the preamble action
6011        let rep = Some(vec![0, 1, 0, 1]);
6012        let def = Some(vec![0, 0, 0, 1]);
6013        let max_visible_def = 0;
6014        let total_items = 3;
6015
6016        let check = |range, expected_item_range, expected_level_range| {
6017            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6018                range,
6019                rep.as_ref(),
6020                def.as_ref(),
6021                max_rep,
6022                max_visible_def,
6023                total_items,
6024                PreambleAction::Take,
6025            );
6026            assert_eq!(item_range, expected_item_range);
6027            assert_eq!(level_range, expected_level_range);
6028        };
6029
6030        // If we are taking the preamble then the range must start at 0
6031        check(0..1, 0..3, 0..3);
6032        check(0..2, 0..3, 0..4);
6033
6034        let check = |range, expected_item_range, expected_level_range| {
6035            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6036                range,
6037                rep.as_ref(),
6038                def.as_ref(),
6039                max_rep,
6040                max_visible_def,
6041                total_items,
6042                PreambleAction::Skip,
6043            );
6044            assert_eq!(item_range, expected_item_range);
6045            assert_eq!(level_range, expected_level_range);
6046        };
6047
6048        check(0..1, 1..3, 1..3);
6049        check(1..2, 3..3, 3..4);
6050        check(0..2, 1..3, 1..4);
6051
6052        // Another preamble case but now it doesn't end with a new list
6053        // [[..., A], NULL, [D, E]]
6054        //
6055        // What we do will depend on the preamble action
6056        let rep = Some(vec![0, 1, 1, 0]);
6057        let def = Some(vec![0, 1, 0, 0]);
6058        let max_visible_def = 0;
6059        let total_items = 4;
6060
6061        let check = |range, expected_item_range, expected_level_range| {
6062            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6063                range,
6064                rep.as_ref(),
6065                def.as_ref(),
6066                max_rep,
6067                max_visible_def,
6068                total_items,
6069                PreambleAction::Take,
6070            );
6071            assert_eq!(item_range, expected_item_range);
6072            assert_eq!(level_range, expected_level_range);
6073        };
6074
6075        // If we are taking the preamble then the range must start at 0
6076        check(0..1, 0..1, 0..2);
6077        check(0..2, 0..3, 0..4);
6078
6079        let check = |range, expected_item_range, expected_level_range| {
6080            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6081                range,
6082                rep.as_ref(),
6083                def.as_ref(),
6084                max_rep,
6085                max_visible_def,
6086                total_items,
6087                PreambleAction::Skip,
6088            );
6089            assert_eq!(item_range, expected_item_range);
6090            assert_eq!(level_range, expected_level_range);
6091        };
6092
6093        // If we are taking the preamble then the range must start at 0
6094        check(0..1, 1..1, 1..2);
6095        check(1..2, 1..3, 2..4);
6096        check(0..2, 1..3, 1..4);
6097
6098        // Now a preamble case without any definition levels
6099        // [[..., A] [B, C], [D]]
6100        let rep = Some(vec![0, 1, 0, 1]);
6101        let def: Option<Vec<u16>> = None;
6102        let max_visible_def = 0;
6103        let total_items = 4;
6104
6105        let check = |range, expected_item_range, expected_level_range| {
6106            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6107                range,
6108                rep.as_ref(),
6109                def.as_ref(),
6110                max_rep,
6111                max_visible_def,
6112                total_items,
6113                PreambleAction::Take,
6114            );
6115            assert_eq!(item_range, expected_item_range);
6116            assert_eq!(level_range, expected_level_range);
6117        };
6118
6119        // If we are taking the preamble then the range must start at 0
6120        check(0..1, 0..3, 0..3);
6121        check(0..2, 0..4, 0..4);
6122
6123        let check = |range, expected_item_range, expected_level_range| {
6124            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6125                range,
6126                rep.as_ref(),
6127                def.as_ref(),
6128                max_rep,
6129                max_visible_def,
6130                total_items,
6131                PreambleAction::Skip,
6132            );
6133            assert_eq!(item_range, expected_item_range);
6134            assert_eq!(level_range, expected_level_range);
6135        };
6136
6137        check(0..1, 1..3, 1..3);
6138        check(1..2, 3..4, 3..4);
6139        check(0..2, 1..4, 1..4);
6140
6141        // If we have nested lists then non-top level lists may be empty/null
6142        // and we need to make sure we still handle them as invisible items (we
6143        // failed to do this previously)
6144        let rep = Some(vec![2, 1, 2, 0, 1, 2]);
6145        let def = Some(vec![0, 1, 2, 0, 0, 0]);
6146        let max_rep = 2;
6147        let max_visible_def = 0;
6148        let total_items = 4;
6149
6150        let check = |range, expected_item_range, expected_level_range| {
6151            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6152                range,
6153                rep.as_ref(),
6154                def.as_ref(),
6155                max_rep,
6156                max_visible_def,
6157                total_items,
6158                PreambleAction::Absent,
6159            );
6160            assert_eq!(item_range, expected_item_range);
6161            assert_eq!(level_range, expected_level_range);
6162        };
6163
6164        check(0..3, 0..4, 0..6);
6165        check(0..1, 0..1, 0..2);
6166        check(1..2, 1..3, 2..5);
6167        check(2..3, 3..4, 5..6);
6168
6169        // Invisible items in a preamble that we are taking (regressing a previous failure)
6170        let rep = Some(vec![0, 0, 1, 0, 1, 1]);
6171        let def = Some(vec![0, 1, 0, 0, 0, 0]);
6172        let max_rep = 1;
6173        let max_visible_def = 0;
6174        let total_items = 5;
6175
6176        let check = |range, expected_item_range, expected_level_range| {
6177            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6178                range,
6179                rep.as_ref(),
6180                def.as_ref(),
6181                max_rep,
6182                max_visible_def,
6183                total_items,
6184                PreambleAction::Take,
6185            );
6186            assert_eq!(item_range, expected_item_range);
6187            assert_eq!(level_range, expected_level_range);
6188        };
6189
6190        check(0..0, 0..1, 0..2);
6191        check(0..1, 0..3, 0..4);
6192        check(0..2, 0..4, 0..5);
6193
6194        // Skip preamble (with invis items) and skip a few rows (with invis items)
6195        // and then take a few rows but not all the rows
6196        let rep = Some(vec![0, 1, 0, 1, 0, 1, 0, 1]);
6197        let def = Some(vec![1, 0, 1, 1, 0, 0, 0, 0]);
6198        let max_rep = 1;
6199        let max_visible_def = 0;
6200        let total_items = 5;
6201
6202        let check = |range, expected_item_range, expected_level_range| {
6203            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6204                range,
6205                rep.as_ref(),
6206                def.as_ref(),
6207                max_rep,
6208                max_visible_def,
6209                total_items,
6210                PreambleAction::Skip,
6211            );
6212            assert_eq!(item_range, expected_item_range);
6213            assert_eq!(level_range, expected_level_range);
6214        };
6215
6216        check(2..3, 2..4, 5..7);
6217    }
6218
6219    #[test]
6220    fn test_slice_batch_data_and_rebase_offsets_u32() {
6221        let data = LanceBuffer::copy_slice(b"0123456789abcdefghij");
6222        let offsets = LanceBuffer::reinterpret_vec(vec![6_u32, 8_u32, 8_u32, 12_u32]);
6223
6224        let (sliced_data, normalized_offsets) =
6225            VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
6226                .unwrap();
6227
6228        assert_eq!(sliced_data.as_ref(), b"6789ab");
6229        let normalized = normalized_offsets.borrow_to_typed_slice::<u32>();
6230        assert_eq!(normalized.as_ref(), &[0, 2, 2, 6]);
6231    }
6232
6233    #[test]
6234    fn test_slice_batch_data_and_rebase_offsets_u64() {
6235        let data = LanceBuffer::copy_slice(b"abcdefghijklmnopqrstuvwxyz");
6236        let offsets = LanceBuffer::reinterpret_vec(vec![10_u64, 12_u64, 16_u64, 20_u64]);
6237
6238        let (sliced_data, normalized_offsets) =
6239            VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 64)
6240                .unwrap();
6241
6242        assert_eq!(sliced_data.as_ref(), b"klmnopqrst");
6243        let normalized = normalized_offsets.borrow_to_typed_slice::<u64>();
6244        assert_eq!(normalized.as_ref(), &[0, 2, 6, 10]);
6245    }
6246
6247    #[test]
6248    fn test_slice_batch_data_and_rebase_offsets_rejects_invalid_offsets() {
6249        let data = LanceBuffer::copy_slice(b"abcd");
6250        let offsets = LanceBuffer::reinterpret_vec(vec![3_u32, 2_u32]);
6251
6252        let err = VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
6253            .expect_err("offset end before start should error");
6254        assert!(err.to_string().contains("less than base"));
6255    }
6256
6257    #[test]
6258    fn test_schedule_instructions() {
6259        // Convert repetition index to bytes for testing
6260        let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
6261        let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
6262        let repetition_index = MiniBlockRepIndex::decode_from_bytes(&rep_bytes, 2);
6263
6264        let check = |user_ranges, expected_instructions| {
6265            let instructions =
6266                ChunkInstructions::schedule_instructions(&repetition_index, user_ranges);
6267            assert_eq!(instructions, expected_instructions);
6268        };
6269
6270        // The instructions we expect if we're grabbing the whole range
6271        let expected_take_all = vec![
6272            ChunkInstructions {
6273                chunk_idx: 0,
6274                preamble: PreambleAction::Absent,
6275                rows_to_skip: 0,
6276                rows_to_take: 6,
6277                take_trailer: true,
6278            },
6279            ChunkInstructions {
6280                chunk_idx: 1,
6281                preamble: PreambleAction::Take,
6282                rows_to_skip: 0,
6283                rows_to_take: 2,
6284                take_trailer: false,
6285            },
6286            ChunkInstructions {
6287                chunk_idx: 2,
6288                preamble: PreambleAction::Absent,
6289                rows_to_skip: 0,
6290                rows_to_take: 5,
6291                take_trailer: true,
6292            },
6293            ChunkInstructions {
6294                chunk_idx: 3,
6295                preamble: PreambleAction::Take,
6296                rows_to_skip: 0,
6297                rows_to_take: 1,
6298                take_trailer: false,
6299            },
6300        ];
6301
6302        // Take all as 1 range
6303        check(&[0..14], expected_take_all.clone());
6304
6305        // Take all a individual rows
6306        check(
6307            &[
6308                0..1,
6309                1..2,
6310                2..3,
6311                3..4,
6312                4..5,
6313                5..6,
6314                6..7,
6315                7..8,
6316                8..9,
6317                9..10,
6318                10..11,
6319                11..12,
6320                12..13,
6321                13..14,
6322            ],
6323            expected_take_all,
6324        );
6325
6326        // Test some partial takes
6327
6328        // 2 rows in the same chunk but not contiguous
6329        check(
6330            &[0..1, 3..4],
6331            vec![
6332                ChunkInstructions {
6333                    chunk_idx: 0,
6334                    preamble: PreambleAction::Absent,
6335                    rows_to_skip: 0,
6336                    rows_to_take: 1,
6337                    take_trailer: false,
6338                },
6339                ChunkInstructions {
6340                    chunk_idx: 0,
6341                    preamble: PreambleAction::Absent,
6342                    rows_to_skip: 3,
6343                    rows_to_take: 1,
6344                    take_trailer: false,
6345                },
6346            ],
6347        );
6348
6349        // Taking just a trailer/preamble
6350        check(
6351            &[5..6],
6352            vec![
6353                ChunkInstructions {
6354                    chunk_idx: 0,
6355                    preamble: PreambleAction::Absent,
6356                    rows_to_skip: 5,
6357                    rows_to_take: 1,
6358                    take_trailer: true,
6359                },
6360                ChunkInstructions {
6361                    chunk_idx: 1,
6362                    preamble: PreambleAction::Take,
6363                    rows_to_skip: 0,
6364                    rows_to_take: 0,
6365                    take_trailer: false,
6366                },
6367            ],
6368        );
6369
6370        // Skipping an entire chunk
6371        check(
6372            &[7..10],
6373            vec![
6374                ChunkInstructions {
6375                    chunk_idx: 1,
6376                    preamble: PreambleAction::Skip,
6377                    rows_to_skip: 1,
6378                    rows_to_take: 1,
6379                    take_trailer: false,
6380                },
6381                ChunkInstructions {
6382                    chunk_idx: 2,
6383                    preamble: PreambleAction::Absent,
6384                    rows_to_skip: 0,
6385                    rows_to_take: 2,
6386                    take_trailer: false,
6387                },
6388            ],
6389        );
6390    }
6391
6392    #[test]
6393    fn test_drain_instructions() {
6394        fn drain_from_instructions(
6395            instructions: &mut VecDeque<ChunkInstructions>,
6396            mut rows_desired: u64,
6397            need_preamble: &mut bool,
6398            skip_in_chunk: &mut u64,
6399        ) -> Vec<ChunkDrainInstructions> {
6400            // Note: instructions.len() is an upper bound, we typically take much fewer
6401            let mut drain_instructions = Vec::with_capacity(instructions.len());
6402            while rows_desired > 0 || *need_preamble {
6403                let (next_instructions, consumed_chunk) = instructions
6404                    .front()
6405                    .unwrap()
6406                    .drain_from_instruction(&mut rows_desired, need_preamble, skip_in_chunk);
6407                if consumed_chunk {
6408                    instructions.pop_front();
6409                }
6410                drain_instructions.push(next_instructions);
6411            }
6412            drain_instructions
6413        }
6414
6415        // Convert repetition index to bytes for testing
6416        let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
6417        let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
6418        let repetition_index = MiniBlockRepIndex::decode_from_bytes(&rep_bytes, 2);
6419        let user_ranges = vec![1..7, 10..14];
6420
6421        // First, schedule the ranges
6422        let scheduled = ChunkInstructions::schedule_instructions(&repetition_index, &user_ranges);
6423
6424        let mut to_drain = VecDeque::from(scheduled.clone());
6425
6426        // Now we drain in batches of 4
6427
6428        let mut need_preamble = false;
6429        let mut skip_in_chunk = 0;
6430
6431        let next_batch =
6432            drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
6433
6434        assert!(!need_preamble);
6435        assert_eq!(skip_in_chunk, 4);
6436        assert_eq!(
6437            next_batch,
6438            vec![ChunkDrainInstructions {
6439                chunk_instructions: scheduled[0].clone(),
6440                rows_to_take: 4,
6441                rows_to_skip: 0,
6442                preamble_action: PreambleAction::Absent,
6443            }]
6444        );
6445
6446        let next_batch =
6447            drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
6448
6449        assert!(!need_preamble);
6450        assert_eq!(skip_in_chunk, 2);
6451
6452        assert_eq!(
6453            next_batch,
6454            vec![
6455                ChunkDrainInstructions {
6456                    chunk_instructions: scheduled[0].clone(),
6457                    rows_to_take: 1,
6458                    rows_to_skip: 4,
6459                    preamble_action: PreambleAction::Absent,
6460                },
6461                ChunkDrainInstructions {
6462                    chunk_instructions: scheduled[1].clone(),
6463                    rows_to_take: 1,
6464                    rows_to_skip: 0,
6465                    preamble_action: PreambleAction::Take,
6466                },
6467                ChunkDrainInstructions {
6468                    chunk_instructions: scheduled[2].clone(),
6469                    rows_to_take: 2,
6470                    rows_to_skip: 0,
6471                    preamble_action: PreambleAction::Absent,
6472                }
6473            ]
6474        );
6475
6476        let next_batch =
6477            drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
6478
6479        assert!(!need_preamble);
6480        assert_eq!(skip_in_chunk, 0);
6481
6482        assert_eq!(
6483            next_batch,
6484            vec![
6485                ChunkDrainInstructions {
6486                    chunk_instructions: scheduled[2].clone(),
6487                    rows_to_take: 1,
6488                    rows_to_skip: 2,
6489                    preamble_action: PreambleAction::Absent,
6490                },
6491                ChunkDrainInstructions {
6492                    chunk_instructions: scheduled[3].clone(),
6493                    rows_to_take: 1,
6494                    rows_to_skip: 0,
6495                    preamble_action: PreambleAction::Take,
6496                },
6497            ]
6498        );
6499
6500        // Regression case.  Need a chunk with preamble, rows, and trailer (the middle chunk here)
6501        let rep_data: Vec<u64> = vec![5, 2, 3, 3, 20, 0];
6502        let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
6503        let repetition_index = MiniBlockRepIndex::decode_from_bytes(&rep_bytes, 2);
6504        let user_ranges = vec![0..28];
6505
6506        // First, schedule the ranges
6507        let scheduled = ChunkInstructions::schedule_instructions(&repetition_index, &user_ranges);
6508
6509        let mut to_drain = VecDeque::from(scheduled.clone());
6510
6511        // Drain first chunk and some of second chunk
6512
6513        let mut need_preamble = false;
6514        let mut skip_in_chunk = 0;
6515
6516        let next_batch =
6517            drain_from_instructions(&mut to_drain, 7, &mut need_preamble, &mut skip_in_chunk);
6518
6519        assert_eq!(
6520            next_batch,
6521            vec![
6522                ChunkDrainInstructions {
6523                    chunk_instructions: scheduled[0].clone(),
6524                    rows_to_take: 6,
6525                    rows_to_skip: 0,
6526                    preamble_action: PreambleAction::Absent,
6527                },
6528                ChunkDrainInstructions {
6529                    chunk_instructions: scheduled[1].clone(),
6530                    rows_to_take: 1,
6531                    rows_to_skip: 0,
6532                    preamble_action: PreambleAction::Take,
6533                },
6534            ]
6535        );
6536
6537        assert!(!need_preamble);
6538        assert_eq!(skip_in_chunk, 1);
6539
6540        // Now, the tricky part.  We drain the second chunk, including the trailer, and need to make sure
6541        // we get a drain task to take the preamble of the third chunk (and nothing else)
6542        let next_batch =
6543            drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
6544
6545        assert_eq!(
6546            next_batch,
6547            vec![
6548                ChunkDrainInstructions {
6549                    chunk_instructions: scheduled[1].clone(),
6550                    rows_to_take: 2,
6551                    rows_to_skip: 1,
6552                    preamble_action: PreambleAction::Skip,
6553                },
6554                ChunkDrainInstructions {
6555                    chunk_instructions: scheduled[2].clone(),
6556                    rows_to_take: 0,
6557                    rows_to_skip: 0,
6558                    preamble_action: PreambleAction::Take,
6559                },
6560            ]
6561        );
6562
6563        assert!(!need_preamble);
6564        assert_eq!(skip_in_chunk, 0);
6565    }
6566
6567    #[tokio::test]
6568    async fn test_fullzip_initialize_is_lazy() {
6569        use futures::{FutureExt, future::BoxFuture};
6570        use std::ops::Range;
6571        use std::sync::Mutex;
6572
6573        #[derive(Debug, Clone)]
6574        struct RecordingScheduler {
6575            data: bytes::Bytes,
6576            requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
6577        }
6578
6579        impl RecordingScheduler {
6580            fn new(data: bytes::Bytes) -> Self {
6581                Self {
6582                    data,
6583                    requests: Arc::new(Mutex::new(Vec::new())),
6584                }
6585            }
6586
6587            fn requests(&self) -> Vec<Vec<Range<u64>>> {
6588                self.requests.lock().unwrap().clone()
6589            }
6590        }
6591
6592        impl crate::EncodingsIo for RecordingScheduler {
6593            fn submit_request(
6594                &self,
6595                ranges: Vec<Range<u64>>,
6596                _priority: u64,
6597            ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
6598                self.requests.lock().unwrap().push(ranges.clone());
6599                let data = ranges
6600                    .into_iter()
6601                    .map(|range| self.data.slice(range.start as usize..range.end as usize))
6602                    .collect::<Vec<_>>();
6603                std::future::ready(Ok(data)).boxed()
6604            }
6605        }
6606
6607        #[derive(Debug)]
6608        struct TestFixedDecompressor;
6609
6610        impl FixedPerValueDecompressor for TestFixedDecompressor {
6611            fn decompress(
6612                &self,
6613                _data: FixedWidthDataBlock,
6614                _num_rows: u64,
6615            ) -> crate::Result<DataBlock> {
6616                unimplemented!("Test decompressor")
6617            }
6618
6619            fn bits_per_value(&self) -> u64 {
6620                32
6621            }
6622        }
6623
6624        let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(vec![
6625            0;
6626            16 * 1024
6627        ])));
6628        let mut scheduler = FullZipScheduler {
6629            data_buf_position: 0,
6630            data_buf_size: 4096,
6631            rep_index: Some(FullZipRepIndexDetails {
6632                buf_position: 1000,
6633                bytes_per_value: 4,
6634            }),
6635            priority: 0,
6636            rows_in_page: 100,
6637            bits_per_offset: 32,
6638            details: Arc::new(FullZipDecodeDetails {
6639                value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
6640                def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
6641                ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
6642                max_rep: 0,
6643                max_visible_def: 0,
6644            }),
6645            cached_state: None,
6646            enable_cache: false,
6647        };
6648
6649        let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
6650        let cached_data = scheduler.initialize(&io_dyn).await.unwrap();
6651
6652        assert!(
6653            cached_data
6654                .as_arc_any()
6655                .downcast_ref::<super::NoCachedPageData>()
6656                .is_some(),
6657            "FullZip initialize should not eagerly load repetition index data"
6658        );
6659        assert!(scheduler.cached_state.is_none());
6660        assert!(
6661            io.requests().is_empty(),
6662            "FullZip initialize should not issue any I/O"
6663        );
6664    }
6665
6666    #[tokio::test]
6667    async fn test_fullzip_read_source_slices_prefetched_page() {
6668        let page_start = 200_u64;
6669        let page_data = LanceBuffer::copy_slice(&[0, 1, 2, 3, 4, 5, 6, 7]);
6670        let source = FullZipReadSource::PrefetchedPage {
6671            base_offset: page_start,
6672            data: page_data,
6673        };
6674        let ranges = vec![
6675            page_start..(page_start + 3),
6676            (page_start + 4)..(page_start + 8),
6677        ];
6678        let mut data = source.fetch(&ranges, 0).await.unwrap();
6679        assert_eq!(data.pop_front().unwrap().as_ref(), &[0, 1, 2]);
6680        assert_eq!(data.pop_front().unwrap().as_ref(), &[4, 5, 6, 7]);
6681    }
6682
6683    #[tokio::test]
6684    async fn test_fullzip_initialize_caches_rep_index_when_enabled() {
6685        use futures::{FutureExt, future::BoxFuture};
6686        use std::ops::Range;
6687        use std::sync::Mutex;
6688
6689        #[derive(Debug, Clone)]
6690        struct RecordingScheduler {
6691            data: bytes::Bytes,
6692            requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
6693        }
6694
6695        impl RecordingScheduler {
6696            fn new(data: bytes::Bytes) -> Self {
6697                Self {
6698                    data,
6699                    requests: Arc::new(Mutex::new(Vec::new())),
6700                }
6701            }
6702
6703            fn requests(&self) -> Vec<Vec<Range<u64>>> {
6704                self.requests.lock().unwrap().clone()
6705            }
6706        }
6707
6708        impl crate::EncodingsIo for RecordingScheduler {
6709            fn submit_request(
6710                &self,
6711                ranges: Vec<Range<u64>>,
6712                _priority: u64,
6713            ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
6714                self.requests.lock().unwrap().push(ranges.clone());
6715                let data = ranges
6716                    .into_iter()
6717                    .map(|range| self.data.slice(range.start as usize..range.end as usize))
6718                    .collect::<Vec<_>>();
6719                std::future::ready(Ok(data)).boxed()
6720            }
6721        }
6722
6723        #[derive(Debug)]
6724        struct TestFixedDecompressor;
6725
6726        impl FixedPerValueDecompressor for TestFixedDecompressor {
6727            fn decompress(
6728                &self,
6729                _data: FixedWidthDataBlock,
6730                _num_rows: u64,
6731            ) -> crate::Result<DataBlock> {
6732                unimplemented!("Test decompressor")
6733            }
6734
6735            fn bits_per_value(&self) -> u64 {
6736                32
6737            }
6738        }
6739
6740        let rows_in_page = 100_u64;
6741        let bytes_per_value = 4_u64;
6742        let rep_start = 1000_u64;
6743        let rep_size = ((rows_in_page + 1) * bytes_per_value) as usize;
6744        let mut data = vec![0_u8; 16 * 1024];
6745        data[rep_start as usize..rep_start as usize + rep_size].fill(7);
6746        let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(data)));
6747
6748        let mut scheduler = FullZipScheduler {
6749            data_buf_position: 0,
6750            data_buf_size: 4096,
6751            rep_index: Some(FullZipRepIndexDetails {
6752                buf_position: rep_start,
6753                bytes_per_value,
6754            }),
6755            priority: 0,
6756            rows_in_page,
6757            bits_per_offset: 32,
6758            details: Arc::new(FullZipDecodeDetails {
6759                value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
6760                def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
6761                ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
6762                max_rep: 0,
6763                max_visible_def: 0,
6764            }),
6765            cached_state: None,
6766            enable_cache: true,
6767        };
6768
6769        let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
6770        let cached_data = scheduler.initialize(&io_dyn).await.unwrap();
6771        assert!(
6772            cached_data
6773                .as_arc_any()
6774                .downcast_ref::<FullZipCacheableState>()
6775                .is_some()
6776        );
6777        assert!(scheduler.cached_state.is_some());
6778        assert_eq!(
6779            io.requests(),
6780            vec![vec![
6781                rep_start..(rep_start + (rows_in_page + 1) * bytes_per_value)
6782            ]]
6783        );
6784    }
6785
6786    #[tokio::test]
6787    async fn test_fullzip_full_page_bypasses_rep_index_io() {
6788        use futures::{FutureExt, future::BoxFuture};
6789        use std::ops::Range;
6790        use std::sync::Mutex;
6791
6792        #[derive(Debug, Clone)]
6793        struct RecordingScheduler {
6794            data: bytes::Bytes,
6795            requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
6796        }
6797
6798        impl RecordingScheduler {
6799            fn new(data: bytes::Bytes) -> Self {
6800                Self {
6801                    data,
6802                    requests: Arc::new(Mutex::new(Vec::new())),
6803                }
6804            }
6805
6806            fn requests(&self) -> Vec<Vec<Range<u64>>> {
6807                self.requests.lock().unwrap().clone()
6808            }
6809        }
6810
6811        impl crate::EncodingsIo for RecordingScheduler {
6812            fn submit_request(
6813                &self,
6814                ranges: Vec<Range<u64>>,
6815                _priority: u64,
6816            ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
6817                self.requests.lock().unwrap().push(ranges.clone());
6818                let data = ranges
6819                    .into_iter()
6820                    .map(|range| self.data.slice(range.start as usize..range.end as usize))
6821                    .collect::<Vec<_>>();
6822                std::future::ready(Ok(data)).boxed()
6823            }
6824        }
6825
6826        #[derive(Debug)]
6827        struct TestFixedDecompressor;
6828
6829        impl FixedPerValueDecompressor for TestFixedDecompressor {
6830            fn decompress(
6831                &self,
6832                _data: FixedWidthDataBlock,
6833                _num_rows: u64,
6834            ) -> crate::Result<DataBlock> {
6835                unimplemented!("Test decompressor")
6836            }
6837
6838            fn bits_per_value(&self) -> u64 {
6839                32
6840            }
6841        }
6842
6843        let rows_in_page = 100_u64;
6844        let data_start = 256_u64;
6845        let data_size = 500_u64;
6846        let rep_start = 4096_u64;
6847        let bytes_per_value = 4_u64;
6848
6849        let mut bytes = vec![0_u8; 16 * 1024];
6850        for i in 0..=rows_in_page {
6851            let offset = (i * 5) as u32;
6852            let pos = rep_start as usize + (i * bytes_per_value) as usize;
6853            bytes[pos..pos + 4].copy_from_slice(&offset.to_le_bytes());
6854        }
6855        let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(bytes)));
6856
6857        let scheduler = FullZipScheduler {
6858            data_buf_position: data_start,
6859            data_buf_size: data_size,
6860            rep_index: Some(FullZipRepIndexDetails {
6861                buf_position: rep_start,
6862                bytes_per_value,
6863            }),
6864            priority: 0,
6865            rows_in_page,
6866            bits_per_offset: 32,
6867            details: Arc::new(FullZipDecodeDetails {
6868                value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
6869                def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
6870                ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
6871                max_rep: 0,
6872                max_visible_def: 0,
6873            }),
6874            cached_state: None,
6875            enable_cache: false,
6876        };
6877
6878        let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
6879        let tasks = scheduler
6880            .schedule_ranges_rep(
6881                &[0..rows_in_page],
6882                &io_dyn,
6883                FullZipRepIndexDetails {
6884                    buf_position: rep_start,
6885                    bytes_per_value,
6886                },
6887            )
6888            .unwrap();
6889
6890        let requests = io.requests();
6891        assert_eq!(requests.len(), 1);
6892        assert_eq!(requests[0], vec![data_start..(data_start + data_size)]);
6893
6894        let _ = tasks.into_iter().next().unwrap().decoder_fut.await.unwrap();
6895        let requests_after_await = io.requests();
6896        assert_eq!(
6897            requests_after_await.len(),
6898            1,
6899            "full page path should not issue rep-index I/O"
6900        );
6901    }
6902
6903    /// This test is used to reproduce fuzz test https://github.com/lancedb/lance/issues/4492
6904    #[tokio::test]
6905    async fn test_fuzz_issue_4492_empty_rep_values() {
6906        use lance_datagen::{RowCount, Seed, array, gen_batch};
6907
6908        let seed = 1823859942947654717u64;
6909        let num_rows = 2741usize;
6910
6911        // Generate the exact same data that caused the failure
6912        let batch_gen = gen_batch().with_seed(Seed::from(seed));
6913        let base_generator = array::rand_type(&DataType::FixedSizeBinary(32));
6914        let list_generator = array::rand_list_any(base_generator, false);
6915
6916        let batch = batch_gen
6917            .anon_col(list_generator)
6918            .into_batch_rows(RowCount::from(num_rows as u64))
6919            .unwrap();
6920
6921        let list_array = batch.column(0).clone();
6922
6923        // Force miniblock encoding
6924        let mut metadata = HashMap::new();
6925        metadata.insert(
6926            STRUCTURAL_ENCODING_META_KEY.to_string(),
6927            STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
6928        );
6929
6930        let test_cases = TestCases::default()
6931            .with_min_file_version(LanceFileVersion::V2_1)
6932            .with_batch_size(100)
6933            .with_range(0..num_rows.min(500) as u64)
6934            .with_indices(vec![0, num_rows as u64 / 2, (num_rows - 1) as u64]);
6935
6936        check_round_trip_encoding_of_data(vec![list_array], &test_cases, metadata).await
6937    }
6938
6939    async fn test_minichunk_size_helper(
6940        string_data: Vec<Option<String>>,
6941        minichunk_size: u64,
6942        file_version: LanceFileVersion,
6943    ) {
6944        use crate::constants::MINICHUNK_SIZE_META_KEY;
6945        use crate::testing::{TestCases, check_round_trip_encoding_of_data};
6946        use arrow_array::{ArrayRef, StringArray};
6947        use std::sync::Arc;
6948
6949        let string_array: ArrayRef = Arc::new(StringArray::from(string_data));
6950
6951        let mut metadata = HashMap::new();
6952        metadata.insert(
6953            MINICHUNK_SIZE_META_KEY.to_string(),
6954            minichunk_size.to_string(),
6955        );
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(file_version)
6963            .with_batch_size(1000);
6964
6965        check_round_trip_encoding_of_data(vec![string_array], &test_cases, metadata).await;
6966    }
6967
6968    #[tokio::test]
6969    async fn test_minichunk_size_roundtrip() {
6970        // Test that minichunk size can be configured and works correctly in round-trip encoding
6971        let mut string_data = Vec::new();
6972        for i in 0..100 {
6973            string_data.push(Some(format!("test_string_{}", i).repeat(50)));
6974        }
6975        // configure minichunk size to 64 bytes (smaller than the default 4kb) for Lance 2.1
6976        test_minichunk_size_helper(string_data, 64, LanceFileVersion::V2_1).await;
6977    }
6978
6979    #[tokio::test]
6980    async fn test_minichunk_size_128kb_v2_2() {
6981        // Test that minichunk size can be configured to 128KB and works correctly with Lance 2.2
6982        let mut string_data = Vec::new();
6983        // create a 500kb string array
6984        for i in 0..10000 {
6985            string_data.push(Some(format!("test_string_{}", i).repeat(50)));
6986        }
6987        test_minichunk_size_helper(string_data, 128 * 1024, LanceFileVersion::V2_2).await;
6988    }
6989
6990    #[tokio::test]
6991    async fn test_binary_large_minichunk_size_over_max_miniblock_values() {
6992        let mut string_data = Vec::new();
6993        // 128kb/chunk / 6 bytes (t_9999) = 21845 items per chunk
6994        for i in 0..10000 {
6995            string_data.push(Some(format!("t_{}", i)));
6996        }
6997        test_minichunk_size_helper(string_data, 128 * 1024, LanceFileVersion::V2_2).await;
6998    }
6999
7000    #[tokio::test]
7001    async fn test_large_dictionary_general_compression() {
7002        use arrow_array::{ArrayRef, StringArray};
7003        use std::collections::HashMap;
7004        use std::sync::Arc;
7005
7006        // Create large string dictionary data (>32KiB) with low cardinality
7007        // Use 100 unique strings, each 500 bytes long = 50KB dictionary
7008        let unique_values: Vec<String> = (0..100)
7009            .map(|i| format!("value_{:04}_{}", i, "x".repeat(500)))
7010            .collect();
7011
7012        // Repeat these strings many times to create a large array
7013        let repeated_strings: Vec<_> = unique_values
7014            .iter()
7015            .cycle()
7016            .take(100_000)
7017            .map(|s| Some(s.as_str()))
7018            .collect();
7019
7020        let string_array = Arc::new(StringArray::from(repeated_strings)) as ArrayRef;
7021
7022        // Configure test to use V2_2 and verify encoding
7023        let test_cases = TestCases::default()
7024            .with_min_file_version(LanceFileVersion::V2_2)
7025            .with_verify_encoding(Arc::new(|cols: &[crate::encoder::EncodedColumn], _| {
7026                assert_eq!(cols.len(), 1);
7027                let col = &cols[0];
7028
7029                // Navigate to the dictionary encoding in the page layout
7030                if let Some(PageEncoding::Structural(page_layout)) =
7031                    &col.final_pages.first().map(|p| &p.description)
7032                    && let Some(pb21::page_layout::Layout::MiniBlockLayout(mini_block)) =
7033                        &page_layout.layout
7034                    && let Some(dictionary_encoding) = &mini_block.dictionary
7035                {
7036                    match dictionary_encoding.compression.as_ref() {
7037                        Some(Compression::General(general)) => {
7038                            // Verify it's using LZ4 or Zstd
7039                            let compression = general.compression.as_ref().unwrap();
7040                            assert!(
7041                                compression.scheme()
7042                                    == pb21::CompressionScheme::CompressionAlgorithmLz4
7043                                    || compression.scheme()
7044                                        == pb21::CompressionScheme::CompressionAlgorithmZstd,
7045                                "Expected LZ4 or Zstd compression for large dictionary"
7046                            );
7047                        }
7048                        _ => panic!("Expected General compression for large dictionary"),
7049                    }
7050                }
7051            }));
7052
7053        check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await;
7054    }
7055
7056    fn dictionary_encoding_from_page(
7057        page: &crate::encoder::EncodedPage,
7058    ) -> &crate::format::pb21::CompressiveEncoding {
7059        let PageEncoding::Structural(layout) = &page.description else {
7060            panic!("Expected structural page encoding");
7061        };
7062        let pb21::page_layout::Layout::MiniBlockLayout(layout) = layout.layout.as_ref().unwrap()
7063        else {
7064            panic!("Expected mini-block layout");
7065        };
7066        layout
7067            .dictionary
7068            .as_ref()
7069            .unwrap_or_else(|| panic!("Expected dictionary encoding"))
7070    }
7071
7072    async fn encode_variable_dict_page(
7073        metadata: HashMap<String, String>,
7074    ) -> crate::encoder::EncodedPage {
7075        use arrow_array::types::Int32Type;
7076        use arrow_array::{ArrayRef, DictionaryArray, Int32Array, StringArray};
7077
7078        let values = Arc::new(StringArray::from(
7079            (0..128)
7080                .map(|i| format!("value_{i:04}_{}", "x".repeat(256)))
7081                .collect::<Vec<_>>(),
7082        )) as ArrayRef;
7083        let keys = Int32Array::from_iter_values((0..20_000).map(|i| i % 128));
7084        let dict_array =
7085            Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap()) as ArrayRef;
7086
7087        let field = arrow_schema::Field::new(
7088            "dict_col",
7089            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
7090            false,
7091        )
7092        .with_metadata(metadata);
7093
7094        encode_first_page(field, dict_array, LanceFileVersion::V2_2).await
7095    }
7096
7097    async fn encode_auto_fixed_dict_page(
7098        metadata: HashMap<String, String>,
7099    ) -> crate::encoder::EncodedPage {
7100        use arrow_array::{ArrayRef, Decimal128Array};
7101
7102        // 128-bit fixed-width values with low cardinality to trigger dictionary encoding.
7103        let values = (0..20_000)
7104            .map(|i| match i % 3 {
7105                0 => 10_i128,
7106                1 => 20_i128,
7107                _ => 30_i128,
7108            })
7109            .collect::<Vec<_>>();
7110        let decimal = Decimal128Array::from_iter_values(values)
7111            .with_precision_and_scale(38, 0)
7112            .unwrap();
7113        let decimal = Arc::new(decimal) as ArrayRef;
7114
7115        let mut field_metadata = metadata;
7116        // Strongly encourage dictionary encoding for this synthetic test data.
7117        field_metadata.insert(
7118            "lance-encoding:dict-size-ratio".to_string(),
7119            "0.99".to_string(),
7120        );
7121        let field = arrow_schema::Field::new("fixed_col", DataType::Decimal128(38, 0), false)
7122            .with_metadata(field_metadata);
7123
7124        encode_first_page(field, decimal, LanceFileVersion::V2_2).await
7125    }
7126
7127    #[tokio::test]
7128    async fn test_dict_values_general_compression_default_lz4_for_variable_dict_values() {
7129        let page = encode_variable_dict_page(HashMap::new()).await;
7130        let dictionary_encoding = dictionary_encoding_from_page(&page);
7131        let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
7132            panic!("Expected General compression for dictionary values");
7133        };
7134        let compression = general.compression.as_ref().unwrap();
7135        assert_eq!(
7136            compression.scheme(),
7137            pb21::CompressionScheme::CompressionAlgorithmLz4
7138        );
7139    }
7140
7141    #[tokio::test]
7142    async fn test_dict_values_general_compression_default_lz4_for_fixed_dict_values() {
7143        let page = encode_auto_fixed_dict_page(HashMap::new()).await;
7144        let dictionary_encoding = dictionary_encoding_from_page(&page);
7145        let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
7146            panic!("Expected General compression for dictionary values");
7147        };
7148        let compression = general.compression.as_ref().unwrap();
7149        assert_eq!(
7150            compression.scheme(),
7151            pb21::CompressionScheme::CompressionAlgorithmLz4
7152        );
7153    }
7154
7155    #[tokio::test]
7156    async fn test_dict_values_general_compression_zstd() {
7157        let mut metadata = HashMap::new();
7158        metadata.insert(
7159            DICT_VALUES_COMPRESSION_META_KEY.to_string(),
7160            "zstd".to_string(),
7161        );
7162        let page = encode_variable_dict_page(metadata).await;
7163        let dictionary_encoding = dictionary_encoding_from_page(&page);
7164        let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
7165            panic!("Expected General compression for dictionary values");
7166        };
7167        let compression = general.compression.as_ref().unwrap();
7168        assert_eq!(
7169            compression.scheme(),
7170            pb21::CompressionScheme::CompressionAlgorithmZstd
7171        );
7172    }
7173
7174    #[tokio::test]
7175    async fn test_dict_values_general_compression_none() {
7176        let mut metadata = HashMap::new();
7177        metadata.insert(
7178            DICT_VALUES_COMPRESSION_META_KEY.to_string(),
7179            "none".to_string(),
7180        );
7181        let page = encode_variable_dict_page(metadata).await;
7182        let dictionary_encoding = dictionary_encoding_from_page(&page);
7183        assert!(
7184            !matches!(
7185                dictionary_encoding.compression.as_ref(),
7186                Some(Compression::General(_))
7187            ),
7188            "Expected dictionary values to avoid General compression"
7189        );
7190    }
7191
7192    #[test]
7193    fn test_resolve_dict_values_compression_metadata_defaults_to_lz4() {
7194        let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
7195            &HashMap::new(),
7196            None,
7197            None,
7198        );
7199        assert_eq!(metadata.get(COMPRESSION_META_KEY), Some(&"lz4".to_string()),);
7200        assert!(!metadata.contains_key(COMPRESSION_LEVEL_META_KEY));
7201    }
7202
7203    #[test]
7204    fn test_resolve_dict_values_compression_metadata_metadata_overrides_env() {
7205        let field_metadata = HashMap::from([
7206            (
7207                DICT_VALUES_COMPRESSION_META_KEY.to_string(),
7208                "none".to_string(),
7209            ),
7210            (
7211                DICT_VALUES_COMPRESSION_LEVEL_META_KEY.to_string(),
7212                "7".to_string(),
7213            ),
7214        ]);
7215        let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
7216            &field_metadata,
7217            Some("zstd".to_string()),
7218            Some("3".to_string()),
7219        );
7220        assert_eq!(
7221            metadata.get(COMPRESSION_META_KEY),
7222            Some(&"none".to_string()),
7223        );
7224        assert_eq!(
7225            metadata.get(COMPRESSION_LEVEL_META_KEY),
7226            Some(&"7".to_string()),
7227        );
7228    }
7229
7230    #[test]
7231    fn test_resolve_dict_values_compression_metadata_env_fallback() {
7232        let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
7233            &HashMap::new(),
7234            Some("zstd".to_string()),
7235            Some("9".to_string()),
7236        );
7237        assert_eq!(
7238            metadata.get(COMPRESSION_META_KEY),
7239            Some(&"zstd".to_string()),
7240        );
7241        assert_eq!(
7242            metadata.get(COMPRESSION_LEVEL_META_KEY),
7243            Some(&"9".to_string()),
7244        );
7245    }
7246
7247    #[tokio::test]
7248    async fn test_dictionary_encode_int64() {
7249        use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
7250        use crate::testing::{TestCases, check_round_trip_encoding_of_data};
7251        use crate::version::LanceFileVersion;
7252        use arrow_array::{ArrayRef, Int64Array};
7253        use std::collections::HashMap;
7254        use std::sync::Arc;
7255
7256        // Low cardinality with poor RLE opportunity.
7257        let values = (0..1000)
7258            .map(|i| match i % 3 {
7259                0 => 10i64,
7260                1 => 20i64,
7261                _ => 30i64,
7262            })
7263            .collect::<Vec<_>>();
7264        let array = Arc::new(Int64Array::from(values)) as ArrayRef;
7265
7266        let mut metadata = HashMap::new();
7267        metadata.insert(
7268            STRUCTURAL_ENCODING_META_KEY.to_string(),
7269            STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
7270        );
7271        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
7272
7273        let test_cases = TestCases::default()
7274            .with_min_file_version(LanceFileVersion::V2_2)
7275            .with_batch_size(1000)
7276            .with_range(0..1000)
7277            .with_indices(vec![0, 1, 10, 999])
7278            .with_expected_encoding("dictionary");
7279
7280        check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
7281    }
7282
7283    #[tokio::test]
7284    async fn test_dictionary_encode_float64() {
7285        use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
7286        use crate::testing::{TestCases, check_round_trip_encoding_of_data};
7287        use crate::version::LanceFileVersion;
7288        use arrow_array::{ArrayRef, Float64Array};
7289        use std::collections::HashMap;
7290        use std::sync::Arc;
7291
7292        // Low cardinality with poor RLE opportunity.
7293        let values = (0..1000)
7294            .map(|i| match i % 3 {
7295                0 => 0.1f64,
7296                1 => 0.2f64,
7297                _ => 0.3f64,
7298            })
7299            .collect::<Vec<_>>();
7300        let array = Arc::new(Float64Array::from(values)) as ArrayRef;
7301
7302        let mut metadata = HashMap::new();
7303        metadata.insert(
7304            STRUCTURAL_ENCODING_META_KEY.to_string(),
7305            STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
7306        );
7307        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
7308
7309        let test_cases = TestCases::default()
7310            .with_min_file_version(LanceFileVersion::V2_2)
7311            .with_batch_size(1000)
7312            .with_range(0..1000)
7313            .with_indices(vec![0, 1, 10, 999])
7314            .with_expected_encoding("dictionary");
7315
7316        check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
7317    }
7318
7319    #[test]
7320    fn test_miniblock_dictionary_out_of_line_bitpacking_decode() {
7321        let rows = 10_000;
7322        let unique_values = 2_000;
7323
7324        let dictionary_encoding =
7325            ProtobufUtils21::out_of_line_bitpacking(64, ProtobufUtils21::flat(11, None));
7326        let layout = pb21::MiniBlockLayout {
7327            rep_compression: None,
7328            def_compression: None,
7329            value_compression: Some(ProtobufUtils21::flat(64, None)),
7330            dictionary: Some(dictionary_encoding),
7331            num_dictionary_items: unique_values,
7332            layers: vec![pb21::RepDefLayer::RepdefAllValidItem as i32],
7333            num_buffers: 1,
7334            repetition_index_depth: 0,
7335            num_items: rows,
7336            has_large_chunk: false,
7337        };
7338
7339        let buffer_offsets_and_sizes = vec![(0, 0), (0, 0), (0, 0)];
7340        let scheduler = super::MiniBlockScheduler::try_new(
7341            &buffer_offsets_and_sizes,
7342            /*priority=*/ 0,
7343            /*items_in_page=*/ rows,
7344            &layout,
7345            &DefaultDecompressionStrategy::default(),
7346        )
7347        .unwrap();
7348
7349        let dictionary = scheduler.dictionary.unwrap();
7350        assert_eq!(dictionary.num_dictionary_items, unique_values);
7351        assert_eq!(
7352            dictionary.dictionary_data_alignment,
7353            crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT
7354        );
7355    }
7356
7357    // Dictionary encoding decision tests
7358    fn create_test_fixed_data_block(
7359        num_values: u64,
7360        cardinality: u64,
7361        bits_per_value: u64,
7362    ) -> DataBlock {
7363        assert!(cardinality > 0);
7364        assert!(cardinality <= num_values);
7365        let block_info = BlockInfo::default();
7366
7367        assert_eq!(bits_per_value % 8, 0);
7368        let data = match bits_per_value {
7369            32 => {
7370                let values = (0..num_values)
7371                    .map(|i| (i % cardinality) as u32)
7372                    .collect::<Vec<_>>();
7373                crate::buffer::LanceBuffer::reinterpret_vec(values)
7374            }
7375            64 => {
7376                let values = (0..num_values).map(|i| i % cardinality).collect::<Vec<_>>();
7377                crate::buffer::LanceBuffer::reinterpret_vec(values)
7378            }
7379            128 => {
7380                let values = (0..num_values)
7381                    .map(|i| (i % cardinality) as u128)
7382                    .collect::<Vec<_>>();
7383                crate::buffer::LanceBuffer::reinterpret_vec(values)
7384            }
7385            _ => unreachable!(),
7386        };
7387        DataBlock::FixedWidth(FixedWidthDataBlock {
7388            bits_per_value,
7389            data,
7390            num_values,
7391            block_info,
7392        })
7393    }
7394
7395    /// Helper to create VariableWidth (string) test data block with exact cardinality
7396    fn create_test_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock {
7397        use arrow_array::StringArray;
7398
7399        assert!(cardinality <= num_values && cardinality > 0);
7400
7401        let mut values = Vec::with_capacity(num_values as usize);
7402        for i in 0..num_values {
7403            values.push(format!("value_{:016}", i % cardinality));
7404        }
7405
7406        let array = StringArray::from(values);
7407        DataBlock::from_array(Arc::new(array) as ArrayRef)
7408    }
7409
7410    fn create_sorted_string_array(num_values: u64, cardinality: u64) -> ArrayRef {
7411        use arrow_array::StringArray;
7412
7413        assert!(cardinality <= num_values && cardinality > 0);
7414
7415        let mut values = Vec::with_capacity(num_values as usize);
7416        for i in 0..num_values {
7417            let value_idx = i * cardinality / num_values;
7418            values.push(format!("value_{:016}", value_idx));
7419        }
7420
7421        Arc::new(StringArray::from(values)) as ArrayRef
7422    }
7423
7424    fn create_sorted_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock {
7425        DataBlock::from_array(create_sorted_string_array(num_values, cardinality))
7426    }
7427
7428    #[test]
7429    fn test_should_dictionary_encode() {
7430        use crate::constants::DICT_SIZE_RATIO_META_KEY;
7431        use lance_core::datatypes::Field as LanceField;
7432
7433        // Create data where dict encoding saves space
7434        let block = create_test_variable_width_block(1000, 10);
7435
7436        let mut metadata = HashMap::new();
7437        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7438        let arrow_field =
7439            arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7440        let field = LanceField::try_from(&arrow_field).unwrap();
7441
7442        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7443            &block,
7444            &field,
7445            LanceFileVersion::V2_1,
7446        );
7447
7448        assert!(
7449            result.is_some(),
7450            "Should use dictionary encode based on size"
7451        );
7452    }
7453
7454    #[test]
7455    fn test_block_sampling_detects_low_cardinality_in_short_sorted_runs() {
7456        let sample_count: usize = 4096;
7457        let num_values: u64 = 200_000;
7458        let cardinality: u64 = 8_000;
7459        let run_length = num_values / cardinality;
7460        let stride = num_values as usize / sample_count;
7461        assert!(
7462            stride > run_length as usize,
7463            "test must construct the stride > run_length case"
7464        );
7465
7466        let block = create_sorted_variable_width_block(num_values, cardinality);
7467        let sample_unique_ratio =
7468            PrimitiveStructuralEncoder::sample_unique_ratio(&block, sample_count).unwrap();
7469
7470        assert!(
7471            sample_unique_ratio.is_some_and(|ratio| ratio < 0.98),
7472            "sorted low-cardinality data must not be classified as near-unique"
7473        );
7474    }
7475
7476    #[test]
7477    fn test_should_dictionary_encode_sorted_low_cardinality() {
7478        use crate::constants::DICT_SIZE_RATIO_META_KEY;
7479        use lance_core::datatypes::Field as LanceField;
7480
7481        let block = create_sorted_variable_width_block(200_000, 8_000);
7482
7483        let mut metadata = HashMap::new();
7484        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7485        let arrow_field =
7486            arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7487        let field = LanceField::try_from(&arrow_field).unwrap();
7488
7489        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7490            &block,
7491            &field,
7492            LanceFileVersion::V2_2,
7493        );
7494
7495        assert!(
7496            result.is_some(),
7497            "sorted low-cardinality data should reach dictionary encoding"
7498        );
7499    }
7500
7501    #[test]
7502    fn test_should_not_dictionary_encode_sorted_high_cardinality_short_runs() {
7503        use crate::constants::DICT_SIZE_RATIO_META_KEY;
7504        use lance_core::datatypes::Field as LanceField;
7505
7506        let num_values = 200_002;
7507        let cardinality = 100_001;
7508        let block = create_sorted_variable_width_block(num_values, cardinality);
7509
7510        let mut metadata = HashMap::new();
7511        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7512        let arrow_field =
7513            arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7514        let field = LanceField::try_from(&arrow_field).unwrap();
7515
7516        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7517            &block,
7518            &field,
7519            LanceFileVersion::V2_2,
7520        );
7521
7522        assert!(
7523            result.is_none(),
7524            "sorted high-cardinality short runs should not trigger a full dictionary probe"
7525        );
7526    }
7527
7528    #[tokio::test]
7529    async fn test_encode_sorted_low_cardinality_uses_dictionary_layout() {
7530        use crate::constants::DICT_SIZE_RATIO_META_KEY;
7531
7532        let mut metadata = HashMap::new();
7533        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7534        let field = arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7535        let array = create_sorted_string_array(200_000, 8_000);
7536
7537        let page = encode_first_page(field, array, LanceFileVersion::V2_2).await;
7538        let _ = dictionary_encoding_from_page(&page);
7539    }
7540
7541    #[test]
7542    fn test_should_not_dictionary_encode_unsupported_bits() {
7543        use crate::constants::DICT_SIZE_RATIO_META_KEY;
7544        use lance_core::datatypes::Field as LanceField;
7545
7546        let block = create_test_fixed_data_block(1000, 1000, 32);
7547
7548        let mut metadata = HashMap::new();
7549        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7550        let arrow_field =
7551            arrow_schema::Field::new("test", DataType::Int32, false).with_metadata(metadata);
7552        let field = LanceField::try_from(&arrow_field).unwrap();
7553
7554        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7555            &block,
7556            &field,
7557            LanceFileVersion::V2_1,
7558        );
7559
7560        assert!(
7561            result.is_none(),
7562            "Should not use dictionary encode for unsupported bit width"
7563        );
7564    }
7565
7566    #[test]
7567    fn test_should_not_dictionary_encode_near_unique_sample() {
7568        use crate::constants::DICT_SIZE_RATIO_META_KEY;
7569        use lance_core::datatypes::Field as LanceField;
7570
7571        let num_values = 5000;
7572        let block = create_test_variable_width_block(num_values, num_values);
7573
7574        let mut metadata = HashMap::new();
7575        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "1.0".to_string());
7576        let arrow_field =
7577            arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7578        let field = LanceField::try_from(&arrow_field).unwrap();
7579
7580        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7581            &block,
7582            &field,
7583            LanceFileVersion::V2_1,
7584        );
7585
7586        assert!(
7587            result.is_none(),
7588            "Should not probe dictionary encoding for near-unique data"
7589        );
7590    }
7591
7592    #[test]
7593    fn test_v2_1_miniblock_serializes_log_num_values_15() {
7594        let miniblocks = MiniBlockCompressed {
7595            data: vec![LanceBuffer::from(vec![1_u8; 16])],
7596            chunks: vec![
7597                MiniBlockChunk {
7598                    buffer_sizes: vec![8],
7599                    log_num_values: 15,
7600                },
7601                MiniBlockChunk {
7602                    buffer_sizes: vec![8],
7603                    log_num_values: 0,
7604                },
7605            ],
7606            num_values: 32_769,
7607        };
7608
7609        let serialized =
7610            PrimitiveStructuralEncoder::serialize_miniblocks(miniblocks, None, None, false)
7611                .unwrap();
7612
7613        let chunk_metadata = serialized.metadata.borrow_to_typed_slice::<u16>();
7614        assert_eq!(chunk_metadata.len(), 2);
7615        assert_eq!(
7616            chunk_metadata[0] & 0x0F,
7617            15,
7618            "V2.1 metadata should use all 4 bits for log_num_values"
7619        );
7620    }
7621
7622    async fn encode_first_page(
7623        field: arrow_schema::Field,
7624        array: ArrayRef,
7625        version: LanceFileVersion,
7626    ) -> crate::encoder::EncodedPage {
7627        use crate::encoder::{
7628            ColumnIndexSequence, EncodingOptions, MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers,
7629            default_encoding_strategy,
7630        };
7631        use crate::repdef::RepDefBuilder;
7632
7633        let lance_field = lance_core::datatypes::Field::try_from(&field).unwrap();
7634        let encoding_strategy = default_encoding_strategy(version);
7635        let mut column_index_seq = ColumnIndexSequence::default();
7636        let encoding_options = EncodingOptions {
7637            cache_bytes_per_column: 1,
7638            max_page_bytes: 32 * 1024 * 1024,
7639            keep_original_array: true,
7640            buffer_alignment: MIN_PAGE_BUFFER_ALIGNMENT,
7641            version,
7642        };
7643
7644        let mut encoder = encoding_strategy
7645            .create_field_encoder(
7646                encoding_strategy.as_ref(),
7647                &lance_field,
7648                &mut column_index_seq,
7649                &encoding_options,
7650            )
7651            .unwrap();
7652
7653        let mut external_buffers = OutOfLineBuffers::new(0, MIN_PAGE_BUFFER_ALIGNMENT);
7654        let repdef = RepDefBuilder::default();
7655        let num_rows = array.len() as u64;
7656        let mut pages = Vec::new();
7657        for task in encoder
7658            .maybe_encode(array, &mut external_buffers, repdef, 0, num_rows)
7659            .unwrap()
7660        {
7661            pages.push(task.await.unwrap());
7662        }
7663        for task in encoder.flush(&mut external_buffers).unwrap() {
7664            pages.push(task.await.unwrap());
7665        }
7666        pages.into_iter().next().unwrap()
7667    }
7668
7669    #[tokio::test]
7670    async fn test_constant_layout_out_of_line_fixed_size_binary_v2_2() {
7671        use crate::format::pb21::page_layout::Layout;
7672
7673        let val = vec![0xABu8; 33];
7674        let arr: ArrayRef = Arc::new(
7675            arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
7676                std::iter::repeat_n(Some(val.as_slice()), 256),
7677                33,
7678            )
7679            .unwrap(),
7680        );
7681        let field = arrow_schema::Field::new("c", DataType::FixedSizeBinary(33), true);
7682        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7683
7684        let PageEncoding::Structural(layout) = &page.description else {
7685            panic!("Expected structural encoding");
7686        };
7687        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7688            panic!("Expected constant layout in slot 2");
7689        };
7690        assert!(layout.inline_value.is_none());
7691        assert_eq!(page.data.len(), 1);
7692
7693        let test_cases = TestCases::default()
7694            .with_min_file_version(LanceFileVersion::V2_2)
7695            .with_max_file_version(LanceFileVersion::V2_2)
7696            .with_page_sizes(vec![4096]);
7697        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7698    }
7699
7700    #[tokio::test]
7701    async fn test_constant_layout_out_of_line_utf8_v2_2() {
7702        use crate::format::pb21::page_layout::Layout;
7703
7704        let arr: ArrayRef = Arc::new(arrow_array::StringArray::from_iter_values(
7705            std::iter::repeat_n("hello", 512),
7706        ));
7707        let field = arrow_schema::Field::new("c", DataType::Utf8, true);
7708        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7709
7710        let PageEncoding::Structural(layout) = &page.description else {
7711            panic!("Expected structural encoding");
7712        };
7713        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7714            panic!("Expected constant layout in slot 2");
7715        };
7716        assert!(layout.inline_value.is_none());
7717        assert_eq!(page.data.len(), 1);
7718
7719        let test_cases = TestCases::default()
7720            .with_min_file_version(LanceFileVersion::V2_2)
7721            .with_max_file_version(LanceFileVersion::V2_2)
7722            .with_page_sizes(vec![4096]);
7723        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7724    }
7725
7726    #[tokio::test]
7727    async fn test_constant_layout_nullable_item_v2_2() {
7728        use crate::format::pb21::page_layout::Layout;
7729
7730        let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![
7731            Some(7),
7732            None,
7733            Some(7),
7734            None,
7735            Some(7),
7736        ]));
7737        let field = arrow_schema::Field::new("c", DataType::Int32, true);
7738        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7739
7740        let PageEncoding::Structural(layout) = &page.description else {
7741            panic!("Expected structural encoding");
7742        };
7743        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7744            panic!("Expected constant layout in slot 2");
7745        };
7746        assert!(layout.inline_value.is_some());
7747        assert_eq!(page.data.len(), 2);
7748
7749        let test_cases = TestCases::default()
7750            .with_min_file_version(LanceFileVersion::V2_2)
7751            .with_max_file_version(LanceFileVersion::V2_2)
7752            .with_page_sizes(vec![4096]);
7753        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7754    }
7755
7756    #[tokio::test]
7757    async fn test_constant_layout_list_repdef_v2_2() {
7758        use crate::format::pb21::page_layout::Layout;
7759        use arrow_array::builder::{Int32Builder, ListBuilder};
7760
7761        let mut builder = ListBuilder::new(Int32Builder::new());
7762        builder.values().append_value(7);
7763        builder.values().append_null();
7764        builder.values().append_value(7);
7765        builder.append(true);
7766
7767        builder.append(true);
7768
7769        builder.values().append_value(7);
7770        builder.append(true);
7771
7772        builder.append_null();
7773
7774        let arr: ArrayRef = Arc::new(builder.finish());
7775        let field = arrow_schema::Field::new(
7776            "c",
7777            DataType::List(Arc::new(arrow_schema::Field::new(
7778                "item",
7779                DataType::Int32,
7780                true,
7781            ))),
7782            true,
7783        );
7784        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7785
7786        let PageEncoding::Structural(layout) = &page.description else {
7787            panic!("Expected structural encoding");
7788        };
7789        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7790            panic!("Expected constant layout in slot 2");
7791        };
7792        assert!(layout.inline_value.is_some());
7793        assert_eq!(page.data.len(), 2);
7794
7795        let test_cases = TestCases::default()
7796            .with_min_file_version(LanceFileVersion::V2_2)
7797            .with_max_file_version(LanceFileVersion::V2_2)
7798            .with_page_sizes(vec![4096]);
7799        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7800    }
7801
7802    #[tokio::test]
7803    async fn test_constant_layout_fixed_size_list_not_used_v2_2() {
7804        use crate::format::pb21::page_layout::Layout;
7805        use arrow_array::builder::{FixedSizeListBuilder, Int32Builder};
7806
7807        let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 3);
7808        for _ in 0..64 {
7809            builder.values().append_value(1);
7810            builder.values().append_null();
7811            builder.values().append_value(3);
7812            builder.append(true);
7813        }
7814        let arr: ArrayRef = Arc::new(builder.finish());
7815        let field = arrow_schema::Field::new(
7816            "c",
7817            DataType::FixedSizeList(
7818                Arc::new(arrow_schema::Field::new("item", DataType::Int32, true)),
7819                3,
7820            ),
7821            true,
7822        );
7823        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7824
7825        if let PageEncoding::Structural(layout) = &page.description {
7826            assert!(
7827                !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
7828                "FixedSizeList should not use constant layout yet"
7829            );
7830        }
7831
7832        let test_cases = TestCases::default()
7833            .with_min_file_version(LanceFileVersion::V2_2)
7834            .with_max_file_version(LanceFileVersion::V2_2)
7835            .with_page_sizes(vec![4096]);
7836        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7837    }
7838
7839    #[tokio::test]
7840    async fn test_constant_layout_not_written_before_v2_2() {
7841        use crate::format::pb21::page_layout::Layout;
7842
7843        let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![7; 1024]));
7844        let field = arrow_schema::Field::new("c", DataType::Int32, true);
7845        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_1).await;
7846
7847        let PageEncoding::Structural(layout) = &page.description else {
7848            return;
7849        };
7850        assert!(
7851            !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
7852            "Should not emit constant layout before v2.2"
7853        );
7854
7855        let test_cases = TestCases::default()
7856            .with_min_file_version(LanceFileVersion::V2_1)
7857            .with_max_file_version(LanceFileVersion::V2_1)
7858            .with_page_sizes(vec![4096]);
7859        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7860    }
7861
7862    #[tokio::test]
7863    async fn test_all_null_constant_layout_still_works_v2_2() {
7864        use crate::format::pb21::page_layout::Layout;
7865
7866        let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![None, None, None]));
7867        let field = arrow_schema::Field::new("c", DataType::Int32, true);
7868        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7869
7870        let PageEncoding::Structural(layout) = &page.description else {
7871            panic!("Expected structural encoding");
7872        };
7873        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7874            panic!("Expected layout in slot 2");
7875        };
7876        assert!(layout.inline_value.is_none());
7877        assert_eq!(page.data.len(), 0);
7878
7879        let test_cases = TestCases::default()
7880            .with_min_file_version(LanceFileVersion::V2_2)
7881            .with_max_file_version(LanceFileVersion::V2_2)
7882            .with_page_sizes(vec![4096]);
7883        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7884    }
7885
7886    #[test]
7887    fn test_encode_decode_complex_all_null_vals_roundtrip() {
7888        use crate::compression::{
7889            DecompressionStrategy, DefaultCompressionStrategy, DefaultDecompressionStrategy,
7890        };
7891
7892        let values: Arc<[u16]> = Arc::from((0..2048).map(|i| (i % 5) as u16).collect::<Vec<u16>>());
7893
7894        let compression_strategy = DefaultCompressionStrategy::default();
7895        let decompression_strategy = DefaultDecompressionStrategy::default();
7896
7897        let (compressed_buf, encoding) = PrimitiveStructuralEncoder::encode_complex_all_null_vals(
7898            &values,
7899            &compression_strategy,
7900        )
7901        .unwrap();
7902
7903        let decompressor = decompression_strategy
7904            .create_block_decompressor(&encoding)
7905            .unwrap();
7906        let decompressed = decompressor
7907            .decompress(compressed_buf, values.len() as u64)
7908            .unwrap();
7909        let decompressed_fixed_width = decompressed.as_fixed_width().unwrap();
7910        assert_eq!(decompressed_fixed_width.num_values, values.len() as u64);
7911        assert_eq!(decompressed_fixed_width.bits_per_value, 16);
7912        let rep_result = decompressed_fixed_width.data.borrow_to_typed_slice::<u16>();
7913        assert_eq!(rep_result.as_ref(), values.as_ref());
7914    }
7915
7916    #[tokio::test]
7917    async fn test_complex_all_null_compression_gated_by_version() {
7918        use crate::format::pb21::page_layout::Layout;
7919        use arrow_array::ListArray;
7920
7921        let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
7922            (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
7923        );
7924        let arr: ArrayRef = Arc::new(list_array);
7925        let field = arrow_schema::Field::new(
7926            "c",
7927            DataType::List(Arc::new(arrow_schema::Field::new(
7928                "item",
7929                DataType::Int32,
7930                true,
7931            ))),
7932            true,
7933        );
7934
7935        let page_v21 = encode_first_page(field.clone(), arr.clone(), LanceFileVersion::V2_1).await;
7936        let PageEncoding::Structural(layout_v21) = &page_v21.description else {
7937            panic!("Expected structural encoding");
7938        };
7939        let Layout::ConstantLayout(layout_v21) = layout_v21.layout.as_ref().unwrap() else {
7940            panic!("Expected constant layout");
7941        };
7942        assert!(layout_v21.rep_compression.is_none());
7943        assert!(layout_v21.def_compression.is_none());
7944        assert_eq!(layout_v21.num_rep_values, 0);
7945        assert_eq!(layout_v21.num_def_values, 0);
7946
7947        let page_v22 = encode_first_page(field, arr, LanceFileVersion::V2_2).await;
7948        let PageEncoding::Structural(layout_v22) = &page_v22.description else {
7949            panic!("Expected structural encoding");
7950        };
7951        let Layout::ConstantLayout(layout_v22) = layout_v22.layout.as_ref().unwrap() else {
7952            panic!("Expected constant layout");
7953        };
7954        assert!(layout_v22.def_compression.is_some());
7955        assert!(layout_v22.num_def_values > 0);
7956    }
7957
7958    #[tokio::test]
7959    async fn test_complex_all_null_round_trip() {
7960        use arrow_array::ListArray;
7961
7962        let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
7963            (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
7964        );
7965
7966        let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_2);
7967        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
7968            .await;
7969    }
7970
7971    // https://github.com/lance-format/lance/issues/6681
7972    #[tokio::test]
7973    async fn test_sparse_boolean_list_roundtrip() {
7974        use arrow_array::builder::{BooleanBuilder, ListBuilder};
7975
7976        let mut list_builder = ListBuilder::new(BooleanBuilder::new());
7977        for i in 0..1000i32 {
7978            if i % 64 == 0 {
7979                // Alternate true/false so the array is not constant (constant path avoids the bug).
7980                list_builder.values().append_value(i % 128 == 0);
7981                list_builder.append(true);
7982            } else {
7983                list_builder.append(false);
7984            }
7985        }
7986        let list_array = Arc::new(list_builder.finish());
7987
7988        let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1);
7989        check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
7990    }
7991
7992    fn truncated_tail_details() -> std::sync::Arc<super::FullZipDecodeDetails> {
7993        use crate::compression::VariablePerValueDecompressor;
7994        use crate::encodings::physical::binary::VariableDecoder;
7995        use crate::repdef::{ControlWordParser, DefinitionInterpretation};
7996        use std::sync::Arc;
7997        Arc::new(super::FullZipDecodeDetails {
7998            value_decompressor: super::PerValueDecompressor::Variable(Arc::new(
7999                VariableDecoder::default(),
8000            )
8001                as Arc<dyn VariablePerValueDecompressor>),
8002            def_meaning: vec![DefinitionInterpretation::NullableItem].into(),
8003            ctrl_word_parser: ControlWordParser::new(0, 0),
8004            max_rep: 0,
8005            max_visible_def: 0,
8006        })
8007    }
8008
8009    fn decode_variable_full_zip(
8010        buf: Vec<u8>,
8011        bits_per_offset: u8,
8012    ) -> lance_core::Result<super::VariableFullZipDecoder> {
8013        use std::collections::VecDeque;
8014        let mut data = VecDeque::new();
8015        data.push_back(crate::buffer::LanceBuffer::from(buf));
8016        super::VariableFullZipDecoder::new(
8017            truncated_tail_details(),
8018            data,
8019            1,
8020            bits_per_offset,
8021            bits_per_offset,
8022        )
8023    }
8024
8025    /// A page whose item walk ends with a partial length prefix must surface a
8026    /// corrupt-file error rather than read past the end of the buffer.
8027    ///
8028    /// This asserts the error variant and message rather than merely expecting a
8029    /// panic: before the length prefix was bounds checked, the read was
8030    /// `get_unchecked` behind a `debug_assert!`, so a debug build panicked here
8031    /// (which a `#[should_panic]` test would have accepted as a pass) while a
8032    /// release build read up to 8 bytes out of a 4 byte allocation.
8033    #[test]
8034    fn variable_full_zip_truncated_length_prefix_is_corrupt_file() {
8035        use lance_core::Error;
8036
8037        for (bits, buf_len) in [(32u8, 3usize), (64u8, 4usize)] {
8038            let err = decode_variable_full_zip(vec![0xAA; buf_len], bits)
8039                .expect_err("a truncated length prefix must not decode");
8040            assert!(
8041                matches!(err, Error::CorruptFile { .. }),
8042                "expected CorruptFile for a {}-bit prefix with {} byte(s), got: {:?}",
8043                bits,
8044                buf_len,
8045                err
8046            );
8047            let msg = err.to_string();
8048            assert!(
8049                msg.contains("truncated length prefix"),
8050                "error should say what is wrong, got: {msg}"
8051            );
8052        }
8053    }
8054}