Skip to main content

lance_file/
reader.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::{
5    borrow::Cow,
6    collections::{BTreeMap, BTreeSet},
7    fmt::Debug,
8    io::Cursor,
9    ops::Range,
10    pin::Pin,
11    sync::Arc,
12};
13
14use arrow_array::RecordBatchReader;
15use arrow_schema::Schema as ArrowSchema;
16use async_trait::async_trait;
17use byteorder::{ByteOrder, LittleEndian, ReadBytesExt};
18use bytes::{Bytes, BytesMut};
19use futures::{Stream, StreamExt, stream::BoxStream};
20use lance_core::deepsize::{Context, DeepSizeOf};
21use lance_encoding::{
22    EncodingsIo,
23    decoder::{
24        ColumnInfo, DecoderConfig, DecoderPlugins, FilterExpression, PageEncoding, ReadBatchTask,
25        RequestedRows, SchedulerDecoderConfig, schedule_and_decode, schedule_and_decode_blocking,
26    },
27    encoder::EncodedBatch,
28};
29use log::debug;
30use object_store::path::Path;
31use prost::Message;
32
33use lance_core::{
34    Error, Result,
35    cache::{CacheKey, CacheKeySchema, KeyBuilder, LanceCache},
36    datatypes::{Field, Schema},
37};
38use lance_encoding::format::pb as pbenc;
39use lance_encoding::format::pb21 as pbenc21;
40use lance_io::{
41    ReadBatchParams,
42    scheduler::FileScheduler,
43    stream::{RecordBatchStream, RecordBatchStreamAdapter},
44};
45
46use crate::{
47    datatypes::{Fields, FieldsWithMeta},
48    format::{MAGIC, pb, pbfile},
49    io::LanceEncodingsIo,
50    version::ConcreteFileVersion,
51    versions,
52};
53
54pub(crate) mod structural;
55
56/// Default chunk size for reading large pages (8MiB)
57/// Pages larger than this will be split into multiple chunks during read
58pub const DEFAULT_READ_CHUNK_SIZE: u64 = 8 * 1024 * 1024;
59
60// For now, we don't use global buffers for anything other than schema.  If we
61// use these later we should make them lazily loaded and then cached once loaded.
62//
63// We store their position / length for debugging purposes
64#[derive(Debug, DeepSizeOf)]
65pub struct BufferDescriptor {
66    pub position: u64,
67    pub size: u64,
68}
69
70impl BufferDescriptor {
71    fn checked_range(&self, buffer_index: usize, file_len: u64) -> Result<Range<u64>> {
72        let end = self.position.checked_add(self.size).ok_or_else(|| {
73            Error::invalid_input_source(
74                format!(
75                    "Global buffer {} range overflows: position={}, size={}",
76                    buffer_index, self.position, self.size
77                )
78                .into(),
79            )
80        })?;
81        if self.position > file_len {
82            return Err(Error::invalid_input_source(
83                format!(
84                    "Global buffer {} position {} is outside file of size {}",
85                    buffer_index, self.position, file_len
86                )
87                .into(),
88            ));
89        }
90        if end > file_len {
91            return Err(Error::invalid_input_source(
92                format!(
93                    "Global buffer {} range {}..{} is outside file of size {}",
94                    buffer_index, self.position, end, file_len
95                )
96                .into(),
97            ));
98        }
99        Ok(self.position..end)
100    }
101}
102
103/// Statistics summarize some of the file metadata for quick summary info
104#[derive(Debug)]
105pub struct FileStatistics {
106    /// Statistics about each of the columns in the file
107    pub columns: Vec<ColumnStatistics>,
108}
109
110/// Summary information describing a column
111#[derive(Debug)]
112pub struct ColumnStatistics {
113    /// The number of pages in the column
114    pub num_pages: usize,
115    /// The total number of data & metadata bytes in the column
116    ///
117    /// This is the compressed on-disk size
118    pub size_bytes: u64,
119}
120
121// TODO: Caching
122#[derive(Debug)]
123pub struct CachedFileMetadata {
124    /// The schema of the file
125    pub file_schema: Arc<Schema>,
126    /// The column metadatas
127    pub column_metadatas: Vec<pbfile::ColumnMetadata>,
128    pub column_infos: Vec<Arc<ColumnInfo>>,
129    /// The number of rows in the file
130    pub num_rows: u64,
131    pub file_buffers: Vec<BufferDescriptor>,
132    /// The number of bytes contained in the data page section of the file
133    pub num_data_bytes: u64,
134    /// The number of bytes contained in the column metadata (not including buffers
135    /// referenced by the metadata)
136    pub num_column_metadata_bytes: u64,
137    /// The number of bytes contained in global buffers
138    pub num_global_buffer_bytes: u64,
139    /// The number of bytes contained in the CMO and GBO tables
140    pub num_footer_bytes: u64,
141    /// The major version number stored in the file footer.
142    pub major_version: u16,
143    /// The minor version number stored in the file footer.
144    pub minor_version: u16,
145    pub version: ConcreteFileVersion,
146    /// The actual total file size in bytes, as reported by the object store.
147    pub file_size_bytes: u64,
148    /// User global buffers (index >= 1) whose bytes were already captured by the
149    /// tail read that `read_all_metadata` performs at open, keyed by buffer index.
150    ///
151    /// All global buffers are laid out contiguously starting at the schema, so on
152    /// small/medium files they land inside the captured tail window. Retaining
153    /// those bytes lets `read_global_buffer` serve them with zero additional I/O.
154    /// The bytes are copied out of the tail (rather than sliced) so the much
155    /// larger tail allocation can be dropped — we only hold what we will serve.
156    ///
157    /// The schema (buffer 0) is excluded: it is already decoded at open and is
158    /// not fetched through `read_global_buffer`. Buffers that fall outside the
159    /// window (large files) are absent here and fall back to a dedicated read.
160    pub retained_global_buffers: BTreeMap<u32, Bytes>,
161}
162
163impl CachedFileMetadata {
164    /// Total file size in bytes.
165    pub fn file_size(&self) -> u64 {
166        self.file_size_bytes
167    }
168}
169
170fn column_metadata_deep_size(column_metadatas: &[pbfile::ColumnMetadata]) -> usize {
171    column_metadatas
172        .iter()
173        .map(|cm| cm.encoded_len() * 4)
174        .sum::<usize>()
175        + std::mem::size_of_val(column_metadatas)
176}
177
178impl DeepSizeOf for CachedFileMetadata {
179    fn deep_size_of_children(&self, context: &mut Context) -> usize {
180        let schema_size = self.file_schema.deep_size_of_children(context);
181
182        let buffers_size: usize = self
183            .file_buffers
184            .iter()
185            .map(|fb| fb.deep_size_of_children(context))
186            .sum();
187
188        // column_metadatas is Vec<pbfile::ColumnMetadata> (protobuf generated,
189        // does not implement DeepSizeOf). We use prost::Message::encoded_len()
190        // as a proxy for in-memory size. The decoded representation is typically
191        // several times larger than the wire format due to heap-allocated
192        // repeated/string/bytes fields, so we apply a 4x multiplier.
193        let column_metadatas_size = column_metadata_deep_size(self.column_metadatas.as_slice());
194
195        // column_infos is Vec<Arc<ColumnInfo>>. Each ColumnInfo contains
196        // page_infos (with protobuf PageEncoding), buffer offsets, and a
197        // column-level ColumnEncoding protobuf.
198        let column_infos_size = self.column_infos.deep_size_of_children(context);
199
200        // Global buffer bytes retained for zero-IO reads (copied out of the tail).
201        let retained_buffers_size = self.retained_global_buffers.deep_size_of_children(context);
202
203        schema_size
204            + buffers_size
205            + column_metadatas_size
206            + column_infos_size
207            + retained_buffers_size
208    }
209}
210
211/// Lightweight file metadata used to locate per-column metadata on demand.
212///
213/// This contains the file-level schema, row count, global buffer descriptors,
214/// and column metadata offset table. Unlike [`CachedFileMetadata`], it does not
215/// hold decoded metadata for every column.
216#[derive(Debug, DeepSizeOf)]
217pub struct FileMetadataIndex {
218    pub(crate) file_schema: Arc<Schema>,
219    pub(crate) num_rows: u64,
220    pub(crate) file_buffers: Vec<BufferDescriptor>,
221    pub(crate) column_metadata_offsets: Arc<[(u64, u64)]>,
222    pub(crate) num_columns: u32,
223    pub(crate) version: ConcreteFileVersion,
224    pub(crate) file_size_bytes: u64,
225    pub(crate) retained_global_buffers: BTreeMap<u32, Bytes>,
226    /// The part of the column metadata region that the tail read at open
227    /// captured, with the file offset of its first byte.
228    ///
229    /// Column metadata is written contiguously right before the offsets table,
230    /// so on small and medium files the whole region sits inside the tail
231    /// window and the lazy loader can decode a column without another request.
232    /// `None` when the window ends before the region starts.
233    pub(crate) retained_column_metadata: Option<(u64, Bytes)>,
234}
235
236impl FileMetadataIndex {
237    /// Returns the total size of the file in bytes.
238    pub fn file_size(&self) -> u64 {
239        self.file_size_bytes
240    }
241
242    /// Returns the number of physical columns in the file.
243    pub fn num_columns(&self) -> u32 {
244        self.num_columns
245    }
246}
247
248#[derive(Debug)]
249struct CachedColumnMetadata {
250    column_metadata: pbfile::ColumnMetadata,
251    column_info: Arc<ColumnInfo>,
252}
253
254impl DeepSizeOf for CachedColumnMetadata {
255    fn deep_size_of_children(&self, context: &mut Context) -> usize {
256        column_metadata_deep_size(std::slice::from_ref(&self.column_metadata))
257            + self.column_info.deep_size_of_children(context)
258    }
259}
260
261#[derive(Debug, Clone)]
262struct ColumnMetadataCacheKey {
263    column_index: u32,
264}
265
266impl CacheKey for ColumnMetadataCacheKey {
267    type ValueType = CachedColumnMetadata;
268
269    fn key(&self) -> Cow<'_, str> {
270        Cow::Owned(format!("column_metadata/{}", self.column_index))
271    }
272
273    fn type_name() -> &'static str {
274        "ColumnMetadata"
275    }
276
277    fn schema() -> CacheKeySchema {
278        CacheKeySchema::new("lance.file.column-metadata-key", 1)
279    }
280
281    fn write_key(&self, builder: &mut KeyBuilder) {
282        builder.write_u32(self.column_index);
283    }
284}
285
286impl CachedFileMetadata {
287    pub fn version(&self) -> ConcreteFileVersion {
288        self.version
289    }
290}
291
292/// Selecting columns from a lance file requires specifying both the
293/// index of the column and the data type of the column
294///
295/// Partly, this is because it is not strictly required that columns
296/// be read into the same type.  For example, a string column may be
297/// read as a string, large_string or string_view type.
298///
299/// A read will only succeed if the decoder for a column is capable
300/// of decoding into the requested type.
301///
302/// Note that this should generally be limited to different in-memory
303/// representations of the same semantic type.  An encoding could
304/// theoretically support "casting" (e.g. int to string, etc.) but
305/// there is little advantage in doing so here.
306///
307/// Note: in order to specify a projection the user will need some way
308/// to figure out the column indices.  In the table format we do this
309/// using field IDs and keeping track of the field id->column index mapping.
310///
311/// If users are not using the table format then they will need to figure
312/// out some way to do this themselves.
313#[derive(Debug, Clone)]
314pub struct ReaderProjection {
315    /// The data types (schema) of the selected columns.  The names
316    /// of the schema are arbitrary and ignored.
317    pub schema: Arc<Schema>,
318    /// The indices of the columns to load.
319    ///
320    /// The content of this vector depends on the file version.
321    ///
322    /// In Lance File Version 2.0 we need ids for structural fields as
323    /// well as leaf fields:
324    ///
325    ///   - Primitive: the index of the column in the schema
326    ///   - List: the index of the list column in the schema
327    ///     followed by the column indices of the children
328    ///   - FixedSizeList (of primitive): the index of the column in the schema
329    ///     (this case is not nested)
330    ///   - FixedSizeList (of non-primitive): not yet implemented
331    ///   - Dictionary: same as primitive
332    ///   - Struct: the index of the struct column in the schema
333    ///     followed by the column indices of the children
334    ///
335    ///   In other words, this should be a DFS listing of the desired schema.
336    ///
337    /// In Lance File Version 2.1 we only need ids for leaf fields.  Any structural
338    /// fields are completely transparent.
339    ///
340    /// For example, if the goal is to load:
341    ///
342    ///   x: int32
343    ///   y: `struct<z: int32, w: string>`
344    ///   z: `list<int32>`
345    ///
346    /// and the schema originally used to store the data was:
347    ///
348    ///   a: `struct<x: int32>`
349    ///   b: int64
350    ///   y: `struct<z: int32, c: int64, w: string>`
351    ///   z: `list<int32>`
352    ///
353    /// Then the column_indices should be:
354    ///
355    /// - 2.0: [1, 3, 4, 6, 7, 8]
356    /// - 2.1: [0, 2, 4, 5]
357    pub column_indices: Vec<u32>,
358}
359
360impl ReaderProjection {
361    /// Returns whether this projection is selective enough to benefit from
362    /// loading column metadata through the file's metadata index.
363    ///
364    /// The caller must already have selected a file format that supports indexed
365    /// metadata. This method only evaluates the projection shape and selectivity.
366    pub fn prefers_indexed_metadata(&self, total_columns: usize) -> bool {
367        FileMetadataProvider::projection_matches_indexed_metadata(self)
368            && self.column_indices.len().saturating_mul(4) < total_columns
369    }
370}
371
372/// File Reader Options that can control reading behaviors, such as whether to enable caching on repetition indices
373#[derive(Clone, Debug)]
374pub struct FileReaderOptions {
375    pub decoder_config: DecoderConfig,
376    /// Size of chunks when reading large pages. Pages larger than this
377    /// will be read in multiple chunks to control memory usage.
378    /// Default: 8MB (DEFAULT_READ_CHUNK_SIZE)
379    pub read_chunk_size: u64,
380    /// If set, the reader will produce batches whose total size in bytes
381    /// is approximately this value. The row-based `batch_size` remains an
382    /// independent upper bound, and the limit reached first determines the batch size.
383    ///
384    /// This can be set at the dataset level (via `ReadParams::file_reader_options`)
385    /// to provide a default for all scans, or at the scanner level (via
386    /// `Scanner::batch_size_bytes`) to override per scan.
387    pub batch_size_bytes: Option<u64>,
388}
389
390impl Default for FileReaderOptions {
391    fn default() -> Self {
392        Self {
393            decoder_config: DecoderConfig::default(),
394            read_chunk_size: DEFAULT_READ_CHUNK_SIZE,
395            batch_size_bytes: None,
396        }
397    }
398}
399
400#[derive(Debug, Clone)]
401pub(crate) struct PreparedProjection {
402    pub column_infos: Vec<Arc<ColumnInfo>>,
403    pub decoder_projection: ReaderProjection,
404}
405
406#[derive(Debug, Clone)]
407pub(crate) enum FileMetadataProvider {
408    Full(Arc<CachedFileMetadata>),
409    Indexed(Arc<FileMetadataIndex>),
410}
411
412/// Executable projection behavior selected by an exact file-version module.
413///
414/// The shared reader invokes this behavior but never interprets a version or
415/// accepted-grammar profile.
416#[async_trait]
417pub(crate) trait ReadProjection: Debug + Send + Sync {
418    fn validate_indexed(
419        &self,
420        projection: &ReaderProjection,
421        metadata_index: &FileMetadataIndex,
422    ) -> Result<()>;
423
424    fn read_length(&self, prepared: &PreparedProjection) -> Result<u64>;
425
426    async fn prepare(
427        &self,
428        metadata_provider: &FileMetadataProvider,
429        projection: &ReaderProjection,
430        io: &Arc<dyn EncodingsIo>,
431        cache: &Arc<LanceCache>,
432    ) -> Result<(PreparedProjection, u64)>;
433}
434
435#[derive(Debug, Clone)]
436pub(crate) struct DecodeEngine {
437    pub scheduler: Arc<dyn EncodingsIo>,
438    pub base_projection: ReaderProjection,
439    pub metadata_provider: FileMetadataProvider,
440    pub read_projection: Arc<dyn ReadProjection>,
441    pub decoder_plugins: Arc<DecoderPlugins>,
442    pub cache: Arc<LanceCache>,
443    pub options: FileReaderOptions,
444}
445
446/// A projection-scoped reader for a current-format Lance file.
447///
448/// This reader fixes a base projection at construction time. All later reads
449/// must stay within that projection, which lets the reader load only the column
450/// metadata needed by the base projection when opening from a [`FileMetadataIndex`].
451/// It intentionally does not expose APIs that require synchronous access to full
452/// file metadata.
453#[derive(Debug, Clone)]
454pub struct ProjectedFileReader {
455    core: DecodeEngine,
456}
457
458/// A current-format Lance file reader backed by fully decoded metadata.
459#[derive(Debug, Clone)]
460pub struct FileReader {
461    pub(crate) core: DecodeEngine,
462    pub(crate) metadata: Arc<CachedFileMetadata>,
463}
464
465pub(crate) fn tasks_to_record_batch_stream(
466    schema: Arc<Schema>,
467    tasks: Pin<Box<dyn Stream<Item = ReadBatchTask> + Send>>,
468    batch_readahead: u32,
469) -> Pin<Box<dyn RecordBatchStream>> {
470    let arrow_schema = Arc::new(ArrowSchema::from(schema.as_ref()));
471    let batches = tasks
472        .map(|task| task.task)
473        .buffered(batch_readahead as usize)
474        .boxed();
475    Box::pin(RecordBatchStreamAdapter::new(arrow_schema, batches))
476}
477
478pub(crate) enum RawFileMetadataOpen {
479    Legacy {
480        major_version: u16,
481        minor_version: u16,
482    },
483    Current {
484        version: ConcreteFileVersion,
485        metadata: RawFileMetadata,
486    },
487}
488
489pub(crate) struct RawFileMetadata {
490    pub file_schema: Arc<Schema>,
491    pub column_metadatas: Vec<pbfile::ColumnMetadata>,
492    pub num_rows: u64,
493    pub file_buffers: Vec<BufferDescriptor>,
494    pub num_data_bytes: u64,
495    pub num_column_metadata_bytes: u64,
496    pub num_global_buffer_bytes: u64,
497    pub num_footer_bytes: u64,
498    pub footer: Footer,
499    pub file_size_bytes: u64,
500    pub retained_global_buffers: BTreeMap<u32, Bytes>,
501}
502
503#[derive(Debug)]
504pub(crate) struct Footer {
505    #[allow(dead_code)]
506    pub column_meta_start: u64,
507    // We don't use this today because we always load metadata for every column
508    // and don't yet support "metadata projection"
509    #[allow(dead_code)]
510    pub column_meta_offsets_start: u64,
511    pub global_buff_offsets_start: u64,
512    pub num_global_buffers: u32,
513    pub num_columns: u32,
514    pub major_version: u16,
515    pub minor_version: u16,
516}
517
518const FOOTER_LEN: usize = 40;
519
520// Count the V2.1 physical columns required to reconstruct a projected field.
521// This is the same DFS shape consumed by `ColumnInfoIter`: ordinary structural
522// nodes are transparent and leaves contribute columns. Indexed metadata loading
523// can therefore compact any ordinary structural projection into 0..N while
524// preserving this order.
525//
526// Blob and packed-struct fields remain unsupported by indexed projection. Their
527// opaque decode semantics are handled by the existing full-metadata reader.
528fn indexed_projection_column_count(field: &Field) -> Option<usize> {
529    if field.is_blob() || field.is_packed_struct() {
530        return None;
531    }
532
533    if field.children.is_empty() {
534        return Some(1);
535    }
536
537    field.children.iter().try_fold(0usize, |count, child| {
538        count.checked_add(indexed_projection_column_count(child)?)
539    })
540}
541
542// The reader combines a projection's columns into rectangular batches, so they
543// must all have the same length.  Returns that common length, or a descriptive
544// error (naming each column's length) when they differ. Ordinary files always
545// pass; only files written with `FileWriter::write_column` whose columns ended up
546// unequal can fail, and those must be read separately.
547pub(crate) fn normalized_column_num_rows(info: &ColumnInfo) -> Result<u64> {
548    info.page_infos.iter().try_fold(0_u64, |rows, page| {
549        let page_rows = match &page.encoding {
550            PageEncoding::Planar(layout) => match layout.sparse() {
551                Some(sparse) => sparse.row_count_and_scale(page.num_rows)?.0,
552                None => page.num_rows,
553            },
554            _ => page.num_rows,
555        };
556        rows.checked_add(page_rows)
557            .ok_or_else(|| Error::invalid_input_source("Column row count overflows u64".into()))
558    })
559}
560
561pub(crate) fn verify_uniform_lengths(field_lengths: &[(&str, u64)]) -> Result<u64> {
562    let first = field_lengths.first().map_or(0, |&(_, len)| len);
563    if field_lengths.iter().all(|&(_, len)| len == first) {
564        return Ok(first);
565    }
566    let columns = field_lengths
567        .iter()
568        .map(|(name, len)| format!("{name}={len}"))
569        .collect::<Vec<_>>()
570        .join(", ");
571    Err(Error::invalid_input(format!(
572        "cannot read columns of differing lengths together ({columns}); \
573         read each column (or equal-length group) separately"
574    )))
575}
576
577impl FileReader {
578    pub(crate) fn base_projection(&self) -> &ReaderProjection {
579        &self.core.base_projection
580    }
581
582    pub(crate) fn full_projection(&self, projection: ReaderProjection) -> PreparedProjection {
583        PreparedProjection {
584            column_infos: self.metadata.column_infos.clone(),
585            decoder_projection: projection,
586        }
587    }
588
589    pub(crate) async fn read_prepared_tasks(
590        &self,
591        params: ReadBatchParams,
592        batch_size: u32,
593        prepared: PreparedProjection,
594        read_len: u64,
595        filter: FilterExpression,
596    ) -> Result<Pin<Box<dyn Stream<Item = ReadBatchTask> + Send>>> {
597        self.core
598            .read_prepared_tasks(params, batch_size, prepared, read_len, filter)
599            .await
600    }
601
602    pub fn with_scheduler(&self, scheduler: Arc<dyn EncodingsIo>) -> Self {
603        Self {
604            core: self.core.with_scheduler(scheduler),
605            metadata: self.metadata.clone(),
606        }
607    }
608
609    /// Returns a clone of this reader whose I/O is additionally recorded into
610    /// `stats`, on top of the scheduler's global accounting.
611    ///
612    /// All cached metadata is shared with `self`, so no file is re-opened and
613    /// only a few `Arc` clones are performed.  If the underlying I/O service
614    /// does not support per-scope statistics (e.g. an in-memory scheduler), the
615    /// returned reader is an ordinary, uninstrumented clone.
616    pub fn with_io_stats(
617        &self,
618        stats: Arc<dyn lance_core::utils::io_stats::IoStatsRecorder>,
619    ) -> Self {
620        match self.core.scheduler.with_io_stats(stats) {
621            Some(scheduler) => self.with_scheduler(scheduler),
622            None => self.clone(),
623        }
624    }
625
626    /// Returns a clone of this reader whose I/O merges requested byte ranges
627    /// that are at most `gap` bytes apart into one request, instead of those
628    /// within the object store's block size.
629    ///
630    /// Each range is cut back out of the merged request and decoded as
631    /// before, so reads return the same data at every gap; a wider gap reads
632    /// the bytes between the ranges too, in fewer requests. All cached
633    /// metadata is shared with `self`. If the underlying I/O service does not
634    /// coalesce requests (e.g. an in-memory scheduler), the returned reader is
635    /// an ordinary clone.
636    pub fn with_coalesce_gap(&self, gap: u64) -> Self {
637        match self.core.scheduler.with_coalesce_gap(gap) {
638            Some(scheduler) => self.with_scheduler(scheduler),
639            None => self.clone(),
640        }
641    }
642
643    pub fn num_rows(&self) -> u64 {
644        self.core.num_rows()
645    }
646
647    /// The number of rows stored in a single physical column.
648    ///
649    /// For ordinary (rectangular) files every column has the same length, equal
650    /// to [`num_rows`](Self::num_rows). Files written with
651    /// [`FileWriter::write_column`](crate::writer::FileWriter::write_column)
652    /// may have columns of differing lengths; this returns the length of one
653    /// such column, derived by summing its pages' row counts. Errors if
654    /// `column_index` is out of bounds.
655    pub fn column_num_rows(&self, column_index: usize) -> Result<u64> {
656        let column = self
657            .metadata
658            .column_metadatas
659            .get(column_index)
660            .ok_or_else(|| {
661                Error::invalid_input(format!(
662                    "column index {} is out of bounds (file has {} columns)",
663                    column_index,
664                    self.metadata.column_metadatas.len()
665                ))
666            })?;
667        Ok(column.pages.iter().map(|page| page.length).sum())
668    }
669
670    pub fn metadata(&self) -> &Arc<CachedFileMetadata> {
671        &self.metadata
672    }
673
674    fn statistics_from_column_metadata(
675        column_metadatas: &[pbfile::ColumnMetadata],
676    ) -> FileStatistics {
677        let column_stats = column_metadatas
678            .iter()
679            .map(|col_metadata| {
680                let num_pages = col_metadata.pages.len();
681                let size_bytes = col_metadata
682                    .pages
683                    .iter()
684                    .map(|page| page.buffer_sizes.iter().sum::<u64>())
685                    .sum::<u64>();
686                ColumnStatistics {
687                    num_pages,
688                    size_bytes,
689                }
690            })
691            .collect();
692
693        FileStatistics {
694            columns: column_stats,
695        }
696    }
697
698    pub fn file_statistics(&self) -> FileStatistics {
699        Self::statistics_from_column_metadata(&self.metadata().column_metadatas)
700    }
701
702    pub async fn read_global_buffer(&self, index: u32) -> Result<Bytes> {
703        self.core.read_global_buffer(index).await
704    }
705
706    async fn read_tail(scheduler: &FileScheduler) -> Result<(Bytes, u64)> {
707        let file_size = scheduler.reader().size().await? as u64;
708        let begin = if file_size < scheduler.reader().block_size() as u64 {
709            0
710        } else {
711            file_size - scheduler.reader().block_size() as u64
712        };
713        let tail_bytes = scheduler.submit_single(begin..file_size, 0).await?;
714        Ok((tail_bytes, file_size))
715    }
716
717    async fn read_range_from_tail_or_scheduler(
718        tail_bytes: &Bytes,
719        tail_offset: u64,
720        scheduler: &FileScheduler,
721        range: Range<u64>,
722    ) -> Result<Bytes> {
723        let tail_end = tail_offset + tail_bytes.len() as u64;
724        if range.start >= tail_offset && range.end <= tail_end {
725            let rel_start = (range.start - tail_offset) as usize;
726            let rel_end = (range.end - tail_offset) as usize;
727            Ok(tail_bytes.slice(rel_start..rel_end))
728        } else {
729            scheduler.submit_single(range, 0).await
730        }
731    }
732
733    fn retained_global_buffers_from_tail(
734        gbo_table: &[BufferDescriptor],
735        tail_bytes: &Bytes,
736        tail_offset: u64,
737        file_len: u64,
738    ) -> Result<BTreeMap<u32, Bytes>> {
739        let tail_end = tail_offset
740            .checked_add(tail_bytes.len() as u64)
741            .ok_or_else(|| Error::invalid_input_source("Tail byte range overflows".into()))?;
742        let mut retained_buffers = BTreeMap::new();
743        for (index, buffer) in gbo_table.iter().enumerate().skip(1) {
744            let range = buffer.checked_range(index, file_len)?;
745            if range.start >= tail_offset && range.end <= tail_end {
746                let rel_start = (range.start - tail_offset) as usize;
747                let rel_end = (range.end - tail_offset) as usize;
748                let bytes = Bytes::copy_from_slice(&tail_bytes[rel_start..rel_end]);
749                retained_buffers.insert(index as u32, bytes);
750            }
751        }
752        Ok(retained_buffers)
753    }
754
755    fn retained_column_metadata_from_tail(
756        footer: &Footer,
757        tail_bytes: &Bytes,
758        tail_offset: u64,
759    ) -> Option<(u64, Bytes)> {
760        let start = footer.column_meta_start.max(tail_offset);
761        let end = footer.column_meta_offsets_start;
762        if start >= end || end > tail_offset + tail_bytes.len() as u64 {
763            return None;
764        }
765        let rel_start = (start - tail_offset) as usize;
766        let rel_end = (end - tail_offset) as usize;
767        // Copied rather than sliced so the cached index does not pin the whole tail
768        // buffer, as `retained_global_buffers_from_tail` does.
769        Some((
770            start,
771            Bytes::copy_from_slice(&tail_bytes[rel_start..rel_end]),
772        ))
773    }
774
775    // Checks to make sure the footer is written correctly and returns the
776    // position of the file descriptor (which comes from the footer)
777    fn decode_footer(footer_bytes: &Bytes) -> Result<Footer> {
778        let len = footer_bytes.len();
779        if len < FOOTER_LEN {
780            return Err(Error::invalid_input(format!(
781                "does not have sufficient data, len: {}, bytes: {:?}",
782                len, footer_bytes
783            )));
784        }
785        let mut cursor = Cursor::new(footer_bytes.slice(len - FOOTER_LEN..));
786
787        let column_meta_start = cursor.read_u64::<LittleEndian>()?;
788        let column_meta_offsets_start = cursor.read_u64::<LittleEndian>()?;
789        let global_buff_offsets_start = cursor.read_u64::<LittleEndian>()?;
790        let num_global_buffers = cursor.read_u32::<LittleEndian>()?;
791        let num_columns = cursor.read_u32::<LittleEndian>()?;
792        let major_version = cursor.read_u16::<LittleEndian>()?;
793        let minor_version = cursor.read_u16::<LittleEndian>()?;
794
795        let magic_bytes = footer_bytes.slice(len - 4..);
796        if magic_bytes.as_ref() != MAGIC {
797            return Err(Error::invalid_input(format!(
798                "file does not appear to be a Lance file (invalid magic: {:?})",
799                MAGIC
800            )));
801        }
802        Ok(Footer {
803            column_meta_start,
804            column_meta_offsets_start,
805            global_buff_offsets_start,
806            num_global_buffers,
807            num_columns,
808            major_version,
809            minor_version,
810        })
811    }
812
813    fn current_file_version(footer: &Footer) -> Result<ConcreteFileVersion> {
814        let version =
815            ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)?;
816        match version {
817            ConcreteFileVersion::V1 => Err(Error::version_conflict(
818                "Attempt to use the lance v2 reader to read a legacy file".to_string(),
819                footer.major_version,
820                footer.minor_version,
821            )),
822            ConcreteFileVersion::V2_0
823            | ConcreteFileVersion::V2_1
824            | ConcreteFileVersion::V2_2
825            | ConcreteFileVersion::V2_3 => Ok(version),
826        }
827    }
828
829    // TODO: Once we have coalesced I/O we should only read the column metadatas that we need
830    fn read_all_column_metadata(
831        column_metadata_bytes: Bytes,
832        footer: &Footer,
833    ) -> Result<Vec<pbfile::ColumnMetadata>> {
834        let column_metadata_start = footer.column_meta_start;
835        // cmo == column_metadata_offsets
836        let cmo_table_size = 16 * footer.num_columns as usize;
837        if column_metadata_bytes.len() < cmo_table_size {
838            return Err(Error::invalid_input(format!(
839                "column metadata region has {} bytes but CMO table needs {} bytes for {} columns",
840                column_metadata_bytes.len(),
841                cmo_table_size,
842                footer.num_columns
843            )));
844        }
845        let cmo_table = column_metadata_bytes.slice(column_metadata_bytes.len() - cmo_table_size..);
846        let column_metadata_offsets = Self::decode_cmo_table(cmo_table, footer)?;
847
848        column_metadata_offsets
849            .iter()
850            .map(|(position, length)| {
851                let normalized_position = (*position - column_metadata_start) as usize;
852                let normalized_end = normalized_position + (*length as usize);
853                Ok(pbfile::ColumnMetadata::decode(
854                    &column_metadata_bytes[normalized_position..normalized_end],
855                )?)
856            })
857            .collect::<Result<Vec<_>>>()
858    }
859
860    fn decode_cmo_table(cmo_table: Bytes, footer: &Footer) -> Result<Arc<[(u64, u64)]>> {
861        let expected_size = 16 * footer.num_columns as usize;
862        if cmo_table.len() != expected_size {
863            return Err(Error::invalid_input(format!(
864                "column metadata offset table has {} bytes but expected {} bytes for {} columns",
865                cmo_table.len(),
866                expected_size,
867                footer.num_columns
868            )));
869        }
870
871        let mut offsets = Vec::with_capacity(footer.num_columns as usize);
872        for col_idx in 0..footer.num_columns {
873            let offset = (col_idx * 16) as usize;
874            let position = LittleEndian::read_u64(&cmo_table[offset..offset + 8]);
875            let length = LittleEndian::read_u64(&cmo_table[offset + 8..offset + 16]);
876            let end = position.checked_add(length).ok_or_else(|| {
877                Error::invalid_input(format!(
878                    "column metadata range overflows for column index {}, position={}, length={}",
879                    col_idx, position, length
880                ))
881            })?;
882            if position < footer.column_meta_start || end > footer.column_meta_offsets_start {
883                return Err(Error::invalid_input(format!(
884                    "column metadata range for column index {} is outside metadata region: position={}, length={}, metadata_start={}, cmo_start={}",
885                    col_idx,
886                    position,
887                    length,
888                    footer.column_meta_start,
889                    footer.column_meta_offsets_start
890                )));
891            }
892            offsets.push((position, length));
893        }
894
895        Ok(Arc::from(offsets))
896    }
897
898    async fn optimistic_tail_read(
899        data: &Bytes,
900        start_pos: u64,
901        scheduler: &FileScheduler,
902        file_len: u64,
903    ) -> Result<Bytes> {
904        let num_bytes_needed = file_len.checked_sub(start_pos).ok_or_else(|| {
905            Error::invalid_input_source(
906                format!(
907                    "Tail read position {} is outside file of size {}",
908                    start_pos, file_len
909                )
910                .into(),
911            )
912        })? as usize;
913        if data.len() >= num_bytes_needed {
914            Ok(data.slice((data.len() - num_bytes_needed)..))
915        } else {
916            let num_bytes_missing = (num_bytes_needed - data.len()) as u64;
917            let start = file_len - num_bytes_needed as u64;
918            let missing_bytes = scheduler
919                .submit_single(start..start + num_bytes_missing, 0)
920                .await?;
921            let mut combined = BytesMut::with_capacity(data.len() + num_bytes_missing as usize);
922            combined.extend(missing_bytes);
923            combined.extend(data);
924            Ok(combined.freeze())
925        }
926    }
927
928    fn do_decode_gbo_table(gbo_bytes: &Bytes, footer: &Footer) -> Result<Vec<BufferDescriptor>> {
929        let mut global_bufs_cursor = Cursor::new(gbo_bytes);
930
931        let mut global_buffers = Vec::with_capacity(footer.num_global_buffers as usize);
932        for _ in 0..footer.num_global_buffers {
933            let buf_pos = global_bufs_cursor.read_u64::<LittleEndian>()?;
934            let buf_size = global_bufs_cursor.read_u64::<LittleEndian>()?;
935            global_buffers.push(BufferDescriptor {
936                position: buf_pos,
937                size: buf_size,
938            });
939        }
940
941        Ok(global_buffers)
942    }
943
944    fn validate_gbo_table(
945        gbo_table: &[BufferDescriptor],
946        file_len: u64,
947        version: ConcreteFileVersion,
948    ) -> Result<()> {
949        versions::validate_global_buffers(version, gbo_table)?;
950        for (buffer_index, buffer) in gbo_table.iter().enumerate() {
951            buffer.checked_range(buffer_index, file_len)?;
952        }
953        Ok(())
954    }
955
956    async fn decode_gbo_table(
957        tail_bytes: &Bytes,
958        file_len: u64,
959        scheduler: &FileScheduler,
960        footer: &Footer,
961        version: ConcreteFileVersion,
962    ) -> Result<Vec<BufferDescriptor>> {
963        // This could, in theory, trigger another IOP but the GBO table should never be large
964        // enough for that to happen
965        let gbo_bytes = Self::optimistic_tail_read(
966            tail_bytes,
967            footer.global_buff_offsets_start,
968            scheduler,
969            file_len,
970        )
971        .await?;
972        let gbo_table = Self::do_decode_gbo_table(&gbo_bytes, footer)?;
973        Self::validate_gbo_table(&gbo_table, file_len, version)?;
974        Ok(gbo_table)
975    }
976
977    fn decode_schema(schema_bytes: Bytes) -> Result<(u64, lance_core::datatypes::Schema)> {
978        let file_descriptor = pb::FileDescriptor::decode(schema_bytes)?;
979        let pb_schema = file_descriptor.schema.unwrap();
980        let num_rows = file_descriptor.length;
981        let fields_with_meta = FieldsWithMeta {
982            fields: Fields(pb_schema.fields),
983            metadata: pb_schema.metadata,
984        };
985        let schema = Schema::try_from(fields_with_meta)?;
986        Ok((num_rows, schema))
987    }
988
989    pub(crate) async fn read_raw_metadata_for_dispatch(
990        scheduler: &FileScheduler,
991    ) -> Result<RawFileMetadataOpen> {
992        let (tail_bytes, file_len) = Self::read_tail(scheduler).await?;
993        let tail_offset = file_len - tail_bytes.len() as u64;
994        let footer = Self::decode_footer(&tail_bytes)?;
995        let version =
996            ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)?;
997        if version == ConcreteFileVersion::V1 {
998            return Ok(RawFileMetadataOpen::Legacy {
999                major_version: footer.major_version,
1000                minor_version: footer.minor_version,
1001            });
1002        }
1003
1004        let gbo_table =
1005            Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer, version).await?;
1006        if gbo_table.is_empty() {
1007            return Err(Error::internal(
1008                "File did not contain any global buffers, schema expected".to_string(),
1009            ));
1010        }
1011        let schema_start = gbo_table[0].position;
1012        let schema_size = gbo_table[0].size;
1013        let num_footer_bytes = file_len.checked_sub(schema_start).ok_or_else(|| {
1014            Error::invalid_input_source(
1015                format!(
1016                    "Schema position {} is outside file of size {}",
1017                    schema_start, file_len
1018                )
1019                .into(),
1020            )
1021        })?;
1022        let all_metadata_bytes =
1023            Self::optimistic_tail_read(&tail_bytes, schema_start, scheduler, file_len).await?;
1024        let schema_bytes = all_metadata_bytes.slice(0..schema_size as usize);
1025        let (num_rows, schema) = Self::decode_schema(schema_bytes)?;
1026
1027        let column_metadata_start = (footer.column_meta_start - schema_start) as usize;
1028        let column_metadata_end = (footer.global_buff_offsets_start - schema_start) as usize;
1029        let column_metadata_bytes =
1030            all_metadata_bytes.slice(column_metadata_start..column_metadata_end);
1031        let column_metadatas = Self::read_all_column_metadata(column_metadata_bytes, &footer)?;
1032
1033        let num_global_buffer_bytes = gbo_table.iter().map(|buf| buf.size).sum::<u64>();
1034        let num_data_bytes = footer.column_meta_start - num_global_buffer_bytes;
1035        let num_column_metadata_bytes = footer.global_buff_offsets_start - footer.column_meta_start;
1036        // The tail read above already pulled in any global buffer that lives within
1037        // the captured window. Copy those user buffers (index >= 1; the schema at 0
1038        // is decoded above and never fetched via read_global_buffer) out of the tail
1039        // so read_global_buffer can serve them without I/O. We copy rather than slice
1040        // so the much larger tail allocation can be released once decoding is done.
1041        let retained_global_buffers = Self::retained_global_buffers_from_tail(
1042            &gbo_table,
1043            &tail_bytes,
1044            tail_offset,
1045            file_len,
1046        )?;
1047
1048        Ok(RawFileMetadataOpen::Current {
1049            version,
1050            metadata: RawFileMetadata {
1051                file_schema: Arc::new(schema),
1052                column_metadatas,
1053                num_rows,
1054                file_buffers: gbo_table,
1055                num_data_bytes,
1056                num_column_metadata_bytes,
1057                num_global_buffer_bytes,
1058                num_footer_bytes,
1059                footer,
1060                file_size_bytes: file_len,
1061                retained_global_buffers,
1062            },
1063        })
1064    }
1065
1066    async fn read_raw_metadata_index_with_known_schema(
1067        scheduler: &FileScheduler,
1068        known_schema: Option<(Arc<Schema>, u64)>,
1069    ) -> Result<FileMetadataIndex> {
1070        let (tail_bytes, file_len) = Self::read_tail(scheduler).await?;
1071        let tail_offset = file_len - tail_bytes.len() as u64;
1072        let footer = Self::decode_footer(&tail_bytes)?;
1073
1074        let file_version = Self::current_file_version(&footer)?;
1075
1076        let gbo_table =
1077            Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer, file_version).await?;
1078        if gbo_table.is_empty() {
1079            return Err(Error::internal(
1080                "File did not contain any global buffers, schema expected".to_string(),
1081            ));
1082        }
1083        let (file_schema, num_rows) = match known_schema {
1084            Some((file_schema, num_rows)) => (file_schema, num_rows),
1085            None => {
1086                let schema_buffer = &gbo_table[0];
1087                let schema_range = schema_buffer.checked_range(0, file_len)?;
1088                let schema_bytes = Self::read_range_from_tail_or_scheduler(
1089                    &tail_bytes,
1090                    tail_offset,
1091                    scheduler,
1092                    schema_range,
1093                )
1094                .await?;
1095                let (num_rows, schema) = Self::decode_schema(schema_bytes)?;
1096                (Arc::new(schema), num_rows)
1097            }
1098        };
1099
1100        let cmo_table = Self::read_range_from_tail_or_scheduler(
1101            &tail_bytes,
1102            tail_offset,
1103            scheduler,
1104            footer.column_meta_offsets_start..footer.global_buff_offsets_start,
1105        )
1106        .await?;
1107        let column_metadata_offsets = Self::decode_cmo_table(cmo_table, &footer)?;
1108
1109        let retained_global_buffers = Self::retained_global_buffers_from_tail(
1110            &gbo_table,
1111            &tail_bytes,
1112            tail_offset,
1113            file_len,
1114        )?;
1115        let retained_column_metadata =
1116            Self::retained_column_metadata_from_tail(&footer, &tail_bytes, tail_offset);
1117
1118        Ok(FileMetadataIndex {
1119            file_schema,
1120            num_rows,
1121            file_buffers: gbo_table,
1122            column_metadata_offsets,
1123            num_columns: footer.num_columns,
1124            version: file_version,
1125            file_size_bytes: file_len,
1126            retained_global_buffers,
1127            retained_column_metadata,
1128        })
1129    }
1130
1131    /// Reads the lightweight metadata index from a file.
1132    ///
1133    /// This reads the file schema from the schema global buffer. Use
1134    /// [`Self::read_metadata_index_with_schema`] when the caller already has
1135    /// the schema and row count from a higher-level metadata source.
1136    pub(crate) async fn read_raw_metadata_index(
1137        scheduler: &FileScheduler,
1138    ) -> Result<FileMetadataIndex> {
1139        Self::read_raw_metadata_index_with_known_schema(scheduler, None).await
1140    }
1141
1142    /// Reads the metadata index without fetching the schema global buffer.
1143    ///
1144    /// Use this when the caller already has the file schema and physical row
1145    /// count from an enclosing metadata layer, such as a dataset manifest.
1146    pub(crate) async fn read_raw_metadata_index_with_schema(
1147        scheduler: &FileScheduler,
1148        file_schema: Arc<Schema>,
1149        num_rows: u64,
1150    ) -> Result<FileMetadataIndex> {
1151        Self::read_raw_metadata_index_with_known_schema(scheduler, Some((file_schema, num_rows)))
1152            .await
1153    }
1154
1155    pub(crate) fn validate_projection(
1156        projection: &ReaderProjection,
1157        metadata: &CachedFileMetadata,
1158    ) -> Result<()> {
1159        if projection.schema.fields.is_empty() {
1160            return Err(Error::invalid_input(
1161                "Attempt to read zero columns from the file, at least one column must be specified"
1162                    .to_string(),
1163            ));
1164        }
1165        let mut column_indices_seen = BTreeSet::new();
1166        for column_index in &projection.column_indices {
1167            if !column_indices_seen.insert(*column_index) {
1168                return Err(Error::invalid_input(format!(
1169                    "The projection specified the column index {} more than once",
1170                    column_index
1171                )));
1172            }
1173            if *column_index >= metadata.column_infos.len() as u32 {
1174                return Err(Error::invalid_input(format!(
1175                    "The projection specified the column index {} but there are only {} columns in the file",
1176                    column_index,
1177                    metadata.column_infos.len()
1178                )));
1179            }
1180        }
1181        Ok(())
1182    }
1183
1184    // The actual decoder needs all the column infos that make up a type.  In other words, if
1185    // the first type in the schema is Struct<i32, i32> then the decoder will need 3 column infos.
1186    //
1187    // This is a file reader concern because the file reader needs to support late projection of columns
1188    // and so it will need to figure this out anyways.
1189    //
1190    // It's a bit of a tricky process though because the number of column infos may depend on the
1191    // encoding.  Considering the above example, if we wrote it with a packed encoding, then there would
1192    // only be a single column in the file (and not 3).
1193    //
1194    // At the moment this method words because our rules are simple and we just repeat them here.  See
1195    // Self::default_projection for a similar problem.  In the future this is something the encodings
1196    // registry will need to figure out.
1197    fn collect_columns_from_projection(
1198        &self,
1199        _projection: &ReaderProjection,
1200    ) -> Result<Vec<Arc<ColumnInfo>>> {
1201        Ok(self.metadata.column_infos.clone())
1202    }
1203
1204    #[allow(clippy::too_many_arguments)]
1205    async fn do_read_range(
1206        column_infos: Vec<Arc<ColumnInfo>>,
1207        io: Arc<dyn EncodingsIo>,
1208        cache: Arc<LanceCache>,
1209        num_rows: u64,
1210        decoder_plugins: Arc<DecoderPlugins>,
1211        range: Range<u64>,
1212        batch_size: u32,
1213        projection: ReaderProjection,
1214        filter: FilterExpression,
1215        decoder_config: DecoderConfig,
1216        batch_size_bytes: Option<u64>,
1217    ) -> Result<BoxStream<'static, ReadBatchTask>> {
1218        debug!(
1219            "Reading range {:?} with batch_size {} from file with {} rows and {} columns into schema with {} columns",
1220            range,
1221            batch_size,
1222            num_rows,
1223            column_infos.len(),
1224            projection.schema.fields.len(),
1225        );
1226
1227        let config = SchedulerDecoderConfig {
1228            batch_size,
1229            cache,
1230            decoder_plugins,
1231            io,
1232            decoder_config,
1233            batch_size_bytes,
1234        };
1235
1236        let requested_rows = RequestedRows::Ranges(vec![range]);
1237
1238        schedule_and_decode(
1239            column_infos,
1240            requested_rows,
1241            filter,
1242            projection.column_indices,
1243            projection.schema,
1244            config,
1245        )
1246        .await
1247    }
1248
1249    #[allow(clippy::too_many_arguments)]
1250    async fn do_take_rows(
1251        column_infos: Vec<Arc<ColumnInfo>>,
1252        io: Arc<dyn EncodingsIo>,
1253        cache: Arc<LanceCache>,
1254        decoder_plugins: Arc<DecoderPlugins>,
1255        indices: Vec<u64>,
1256        batch_size: u32,
1257        projection: ReaderProjection,
1258        filter: FilterExpression,
1259        decoder_config: DecoderConfig,
1260        batch_size_bytes: Option<u64>,
1261    ) -> Result<BoxStream<'static, ReadBatchTask>> {
1262        debug!(
1263            "Taking {} rows spread across range {}..{} with batch_size {} from columns {:?}",
1264            indices.len(),
1265            indices[0],
1266            indices[indices.len() - 1],
1267            batch_size,
1268            column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1269        );
1270
1271        let config = SchedulerDecoderConfig {
1272            batch_size,
1273            cache,
1274            decoder_plugins,
1275            io,
1276            decoder_config,
1277            batch_size_bytes,
1278        };
1279
1280        let requested_rows = RequestedRows::Indices(indices);
1281
1282        schedule_and_decode(
1283            column_infos,
1284            requested_rows,
1285            filter,
1286            projection.column_indices,
1287            projection.schema,
1288            config,
1289        )
1290        .await
1291    }
1292
1293    #[allow(clippy::too_many_arguments)]
1294    async fn do_read_ranges(
1295        column_infos: Vec<Arc<ColumnInfo>>,
1296        io: Arc<dyn EncodingsIo>,
1297        cache: Arc<LanceCache>,
1298        decoder_plugins: Arc<DecoderPlugins>,
1299        ranges: Vec<Range<u64>>,
1300        batch_size: u32,
1301        projection: ReaderProjection,
1302        filter: FilterExpression,
1303        decoder_config: DecoderConfig,
1304        batch_size_bytes: Option<u64>,
1305    ) -> Result<BoxStream<'static, ReadBatchTask>> {
1306        let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
1307        debug!(
1308            "Taking {} ranges ({} rows) spread across range {}..{} with batch_size {} from columns {:?}",
1309            ranges.len(),
1310            num_rows,
1311            ranges[0].start,
1312            ranges[ranges.len() - 1].end,
1313            batch_size,
1314            column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1315        );
1316
1317        let config = SchedulerDecoderConfig {
1318            batch_size,
1319            cache,
1320            decoder_plugins,
1321            io,
1322            decoder_config,
1323            batch_size_bytes,
1324        };
1325
1326        let requested_rows = RequestedRows::Ranges(ranges);
1327
1328        schedule_and_decode(
1329            column_infos,
1330            requested_rows,
1331            filter,
1332            projection.column_indices,
1333            projection.schema,
1334            config,
1335        )
1336        .await
1337    }
1338
1339    fn take_rows_blocking(
1340        &self,
1341        indices: Vec<u64>,
1342        batch_size: u32,
1343        projection: ReaderProjection,
1344        filter: FilterExpression,
1345    ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1346        let column_infos = self.collect_columns_from_projection(&projection)?;
1347        debug!(
1348            "Taking {} rows spread across range {}..{} with batch_size {} from columns {:?}",
1349            indices.len(),
1350            indices[0],
1351            indices[indices.len() - 1],
1352            batch_size,
1353            column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1354        );
1355
1356        let config = SchedulerDecoderConfig {
1357            batch_size,
1358            cache: self.core.cache.clone(),
1359            decoder_plugins: self.core.decoder_plugins.clone(),
1360            io: self.core.scheduler.clone(),
1361            decoder_config: self.core.options.decoder_config.clone(),
1362            batch_size_bytes: self.core.options.batch_size_bytes,
1363        };
1364
1365        let requested_rows = RequestedRows::Indices(indices);
1366
1367        schedule_and_decode_blocking(
1368            column_infos,
1369            requested_rows,
1370            filter,
1371            projection.column_indices,
1372            projection.schema,
1373            config,
1374        )
1375    }
1376
1377    fn read_ranges_blocking(
1378        &self,
1379        ranges: Vec<Range<u64>>,
1380        batch_size: u32,
1381        projection: ReaderProjection,
1382        filter: FilterExpression,
1383    ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1384        let column_infos = self.collect_columns_from_projection(&projection)?;
1385        let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
1386        debug!(
1387            "Taking {} ranges ({} rows) spread across range {}..{} with batch_size {} from columns {:?}",
1388            ranges.len(),
1389            num_rows,
1390            ranges[0].start,
1391            ranges[ranges.len() - 1].end,
1392            batch_size,
1393            column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1394        );
1395
1396        let config = SchedulerDecoderConfig {
1397            batch_size,
1398            cache: self.core.cache.clone(),
1399            decoder_plugins: self.core.decoder_plugins.clone(),
1400            io: self.core.scheduler.clone(),
1401            decoder_config: self.core.options.decoder_config.clone(),
1402            batch_size_bytes: self.core.options.batch_size_bytes,
1403        };
1404
1405        let requested_rows = RequestedRows::Ranges(ranges);
1406
1407        schedule_and_decode_blocking(
1408            column_infos,
1409            requested_rows,
1410            filter,
1411            projection.column_indices,
1412            projection.schema,
1413            config,
1414        )
1415    }
1416
1417    fn read_range_blocking(
1418        &self,
1419        range: Range<u64>,
1420        batch_size: u32,
1421        projection: ReaderProjection,
1422        filter: FilterExpression,
1423    ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1424        let column_infos = self.collect_columns_from_projection(&projection)?;
1425        let num_rows = self.core.num_rows();
1426
1427        debug!(
1428            "Reading range {:?} with batch_size {} from file with {} rows and {} columns into schema with {} columns",
1429            range,
1430            batch_size,
1431            num_rows,
1432            column_infos.len(),
1433            projection.schema.fields.len(),
1434        );
1435
1436        let config = SchedulerDecoderConfig {
1437            batch_size,
1438            cache: self.core.cache.clone(),
1439            decoder_plugins: self.core.decoder_plugins.clone(),
1440            io: self.core.scheduler.clone(),
1441            decoder_config: self.core.options.decoder_config.clone(),
1442            batch_size_bytes: self.core.options.batch_size_bytes,
1443        };
1444
1445        let requested_rows = RequestedRows::Ranges(vec![range]);
1446
1447        schedule_and_decode_blocking(
1448            column_infos,
1449            requested_rows,
1450            filter,
1451            projection.column_indices,
1452            projection.schema,
1453            config,
1454        )
1455    }
1456
1457    pub(crate) fn read_prepared_blocking(
1458        &self,
1459        params: ReadBatchParams,
1460        batch_size: u32,
1461        prepared: PreparedProjection,
1462        read_len: u64,
1463        filter: FilterExpression,
1464    ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1465        let projection = prepared.decoder_projection;
1466        let verify_bound = |params: &ReadBatchParams, bound: u64, inclusive: bool| {
1467            if bound > read_len || (bound == read_len && inclusive) {
1468                Err(Error::invalid_input(format!(
1469                    "cannot read {params:?} from columns with {read_len} rows"
1470                )))
1471            } else {
1472                Ok(())
1473            }
1474        };
1475        match &params {
1476            ReadBatchParams::Indices(indices) => {
1477                for index in indices {
1478                    match index {
1479                        None => return Err(Error::invalid_input("Null value in indices array")),
1480                        Some(index) => verify_bound(&params, index as u64, true)?,
1481                    }
1482                }
1483                let indices = indices.iter().map(|index| index.unwrap() as u64).collect();
1484                self.take_rows_blocking(indices, batch_size, projection, filter)
1485            }
1486            ReadBatchParams::Range(range) => {
1487                verify_bound(&params, range.end as u64, false)?;
1488                self.read_range_blocking(
1489                    range.start as u64..range.end as u64,
1490                    batch_size,
1491                    projection,
1492                    filter,
1493                )
1494            }
1495            ReadBatchParams::Ranges(ranges) => {
1496                let mut ranges_u64 = Vec::with_capacity(ranges.len());
1497                for range in ranges.as_ref() {
1498                    verify_bound(&params, range.end, false)?;
1499                    ranges_u64.push(range.start..range.end);
1500                }
1501                self.read_ranges_blocking(ranges_u64, batch_size, projection, filter)
1502            }
1503            ReadBatchParams::RangeFrom(range) => {
1504                verify_bound(&params, range.start as u64, true)?;
1505                self.read_range_blocking(
1506                    range.start as u64..read_len,
1507                    batch_size,
1508                    projection,
1509                    filter,
1510                )
1511            }
1512            ReadBatchParams::RangeTo(range) => {
1513                verify_bound(&params, range.end as u64, false)?;
1514                self.read_range_blocking(0..range.end as u64, batch_size, projection, filter)
1515            }
1516            ReadBatchParams::RangeFull => {
1517                self.read_range_blocking(0..read_len, batch_size, projection, filter)
1518            }
1519        }
1520    }
1521
1522    pub fn schema(&self) -> &Arc<Schema> {
1523        self.core.schema()
1524    }
1525}
1526
1527impl FileMetadataProvider {
1528    pub(crate) fn version(&self) -> ConcreteFileVersion {
1529        match self {
1530            Self::Full(metadata) => metadata.version,
1531            Self::Indexed(metadata_index) => metadata_index.version,
1532        }
1533    }
1534
1535    pub(crate) fn num_rows(&self) -> u64 {
1536        match self {
1537            Self::Full(metadata) => metadata.num_rows,
1538            Self::Indexed(metadata_index) => metadata_index.num_rows,
1539        }
1540    }
1541
1542    pub(crate) fn schema(&self) -> &Arc<Schema> {
1543        match self {
1544            Self::Full(metadata) => &metadata.file_schema,
1545            Self::Indexed(metadata_index) => &metadata_index.file_schema,
1546        }
1547    }
1548
1549    pub(crate) fn file_buffers(&self) -> &Vec<BufferDescriptor> {
1550        match self {
1551            Self::Full(metadata) => &metadata.file_buffers,
1552            Self::Indexed(metadata_index) => &metadata_index.file_buffers,
1553        }
1554    }
1555
1556    pub(crate) fn retained_global_buffers(&self) -> &BTreeMap<u32, Bytes> {
1557        match self {
1558            Self::Full(metadata) => &metadata.retained_global_buffers,
1559            Self::Indexed(metadata_index) => &metadata_index.retained_global_buffers,
1560        }
1561    }
1562
1563    fn file_size(&self) -> u64 {
1564        match self {
1565            Self::Full(metadata) => metadata.file_size_bytes,
1566            Self::Indexed(metadata_index) => metadata_index.file_size_bytes,
1567        }
1568    }
1569
1570    pub(crate) fn file_statistics(&self) -> Option<FileStatistics> {
1571        let metadata = match self {
1572            Self::Full(metadata) => metadata,
1573            Self::Indexed(_) => return None,
1574        };
1575        Some(FileReader::statistics_from_column_metadata(
1576            &metadata.column_metadatas,
1577        ))
1578    }
1579
1580    pub(crate) fn projection_matches_indexed_metadata(projection: &ReaderProjection) -> bool {
1581        if projection.schema.fields.is_empty() {
1582            return false;
1583        }
1584
1585        projection
1586            .schema
1587            .fields
1588            .iter()
1589            .try_fold(0usize, |count, field| {
1590                count.checked_add(indexed_projection_column_count(field)?)
1591            })
1592            == Some(projection.column_indices.len())
1593    }
1594
1595    pub(crate) fn validate_indexed_projection_structure(
1596        projection: &ReaderProjection,
1597        metadata_index: &FileMetadataIndex,
1598    ) -> Result<()> {
1599        if projection.schema.fields.is_empty() {
1600            return Err(Error::invalid_input(
1601                "Attempt to read zero columns from the file, at least one column must be specified"
1602                    .to_string(),
1603            ));
1604        }
1605        let mut column_indices_seen = BTreeSet::new();
1606        for column_index in &projection.column_indices {
1607            if !column_indices_seen.insert(*column_index) {
1608                return Err(Error::invalid_input(format!(
1609                    "The projection specified the column index {} more than once",
1610                    column_index
1611                )));
1612            }
1613            if *column_index >= metadata_index.num_columns {
1614                return Err(Error::invalid_input(format!(
1615                    "The projection specified the column index {} but there are only {} columns in the file",
1616                    column_index, metadata_index.num_columns
1617                )));
1618            }
1619        }
1620        Ok(())
1621    }
1622
1623    pub(crate) fn indexed_projection_error(
1624        projection: &ReaderProjection,
1625        metadata_index: &FileMetadataIndex,
1626    ) -> Error {
1627        Error::not_supported(format!(
1628            "lazy column metadata loading requires a V2.1+ ordinary structural projection without blob or packed-struct fields whose physical-column count matches the projection; got file version {:?}, {} schema fields, and {} column indices",
1629            metadata_index.version,
1630            projection.schema.fields.len(),
1631            projection.column_indices.len()
1632        ))
1633    }
1634
1635    fn column_metadata_range(
1636        metadata_index: &FileMetadataIndex,
1637        column_index: u32,
1638    ) -> Result<Range<u64>> {
1639        let (position, length) = metadata_index
1640            .column_metadata_offsets
1641            .get(column_index as usize)
1642            .copied()
1643            .ok_or_else(|| {
1644                Error::invalid_input(format!(
1645                    "The projection specified the column index {} but there are only {} columns in the file",
1646                    column_index, metadata_index.num_columns
1647                ))
1648            })?;
1649        let end = position.checked_add(length).ok_or_else(|| {
1650            Error::invalid_input(format!(
1651                "column metadata range overflows for column index {}, position={}, length={}",
1652                column_index, position, length
1653            ))
1654        })?;
1655        Ok(position..end)
1656    }
1657
1658    /// The column metadata bytes for `range` when the tail read at open already
1659    /// holds them, so no request is needed to decode the column.
1660    fn retained_column_metadata(
1661        metadata_index: &FileMetadataIndex,
1662        range: &Range<u64>,
1663    ) -> Option<Bytes> {
1664        let (start, bytes) = metadata_index.retained_column_metadata.as_ref()?;
1665        let end = start + bytes.len() as u64;
1666        (range.start >= *start && range.end <= end)
1667            .then(|| bytes.slice((range.start - start) as usize..(range.end - start) as usize))
1668    }
1669
1670    pub(crate) async fn load_indexed_column_infos<F>(
1671        metadata_index: &FileMetadataIndex,
1672        io: &Arc<dyn EncodingsIo>,
1673        cache: &Arc<LanceCache>,
1674        column_indices: &[u32],
1675        decode_column: F,
1676    ) -> Result<Vec<Arc<ColumnInfo>>>
1677    where
1678        F: Fn(u32, &pbfile::ColumnMetadata) -> Result<Arc<ColumnInfo>>,
1679    {
1680        let mut column_infos = vec![None; column_indices.len()];
1681        let mut retained_columns = Vec::new();
1682        let mut missing_columns = Vec::new();
1683
1684        for (result_index, column_index) in column_indices.iter().copied().enumerate() {
1685            let cache_key = ColumnMetadataCacheKey { column_index };
1686            if let Some(cached) = cache.get_with_key(&cache_key).await {
1687                column_infos[result_index] = Some(cached.column_info.clone());
1688                continue;
1689            }
1690            let range = Self::column_metadata_range(metadata_index, column_index)?;
1691            match Self::retained_column_metadata(metadata_index, &range) {
1692                Some(bytes) => retained_columns.push((result_index, column_index, bytes)),
1693                None => missing_columns.push((result_index, column_index, range)),
1694            }
1695        }
1696
1697        missing_columns.sort_by_key(|(_, _, range)| range.start);
1698        let fetched_bytes = if missing_columns.is_empty() {
1699            Vec::new()
1700        } else {
1701            let ranges = missing_columns
1702                .iter()
1703                .map(|(_, _, range)| range.clone())
1704                .collect::<Vec<_>>();
1705            io.submit_request(ranges, 0).await?
1706        };
1707        let fetched_columns = missing_columns
1708            .into_iter()
1709            .zip(fetched_bytes)
1710            .map(|((result_index, column_index, _), bytes)| (result_index, column_index, bytes));
1711        for (result_index, column_index, bytes) in
1712            retained_columns.into_iter().chain(fetched_columns)
1713        {
1714            let column_metadata = pbfile::ColumnMetadata::decode(bytes)?;
1715            let column_info = decode_column(column_index, &column_metadata)?;
1716            let cached = Arc::new(CachedColumnMetadata {
1717                column_metadata,
1718                column_info: column_info.clone(),
1719            });
1720            let cache_key = ColumnMetadataCacheKey { column_index };
1721            cache.insert_with_key(&cache_key, cached).await;
1722            column_infos[result_index] = Some(column_info);
1723        }
1724
1725        column_infos
1726            .into_iter()
1727            .enumerate()
1728            .map(|(idx, column_info)| {
1729                column_info.ok_or_else(|| {
1730                    Error::internal(format!(
1731                        "lazy metadata loader did not load requested projection column at position {}",
1732                        idx
1733                    ))
1734                })
1735            })
1736            .collect()
1737    }
1738}
1739
1740impl DecodeEngine {
1741    pub(crate) fn try_new(
1742        scheduler: Arc<dyn EncodingsIo>,
1743        base_projection: ReaderProjection,
1744        decoder_plugins: Arc<DecoderPlugins>,
1745        metadata_provider: FileMetadataProvider,
1746        read_projection: Arc<dyn ReadProjection>,
1747        cache: Arc<LanceCache>,
1748        options: FileReaderOptions,
1749    ) -> Result<Self> {
1750        Ok(Self {
1751            scheduler,
1752            base_projection,
1753            metadata_provider,
1754            read_projection,
1755            decoder_plugins,
1756            cache,
1757            options,
1758        })
1759    }
1760
1761    pub(crate) fn with_scheduler(&self, scheduler: Arc<dyn EncodingsIo>) -> Self {
1762        Self {
1763            scheduler,
1764            base_projection: self.base_projection.clone(),
1765            metadata_provider: self.metadata_provider.clone(),
1766            read_projection: self.read_projection.clone(),
1767            decoder_plugins: self.decoder_plugins.clone(),
1768            cache: self.cache.clone(),
1769            options: self.options.clone(),
1770        }
1771    }
1772
1773    fn num_rows(&self) -> u64 {
1774        self.metadata_provider.num_rows()
1775    }
1776
1777    fn schema(&self) -> &Arc<Schema> {
1778        self.metadata_provider.schema()
1779    }
1780
1781    pub(crate) async fn read_global_buffer(&self, index: u32) -> Result<Bytes> {
1782        let file_buffers = self.metadata_provider.file_buffers();
1783        let buffer_desc = file_buffers.get(index as usize).ok_or_else(|| {
1784            Error::invalid_input(format!(
1785                "request for global buffer at index {} but there were only {} global buffers in the file",
1786                index,
1787                file_buffers.len()
1788            ))
1789        })?;
1790
1791        if let Some(bytes) = self.metadata_provider.retained_global_buffers().get(&index) {
1792            return Ok(bytes.clone());
1793        }
1794
1795        let bytes = self
1796            .scheduler
1797            .submit_request(
1798                vec![
1799                    buffer_desc
1800                        .checked_range(index as usize, self.metadata_provider.file_size())?,
1801                ],
1802                0,
1803            )
1804            .await?;
1805        bytes.into_iter().next().ok_or_else(|| {
1806            Error::internal(format!(
1807                "global buffer read for index {} returned no bytes",
1808                index
1809            ))
1810        })
1811    }
1812
1813    async fn read_range(
1814        &self,
1815        range: Range<u64>,
1816        batch_size: u32,
1817        prepared: PreparedProjection,
1818        filter: FilterExpression,
1819    ) -> Result<BoxStream<'static, ReadBatchTask>> {
1820        FileReader::do_read_range(
1821            prepared.column_infos,
1822            self.scheduler.clone(),
1823            self.cache.clone(),
1824            self.num_rows(),
1825            self.decoder_plugins.clone(),
1826            range,
1827            batch_size,
1828            prepared.decoder_projection,
1829            filter,
1830            self.options.decoder_config.clone(),
1831            self.options.batch_size_bytes,
1832        )
1833        .await
1834    }
1835
1836    async fn take_rows(
1837        &self,
1838        indices: Vec<u64>,
1839        batch_size: u32,
1840        prepared: PreparedProjection,
1841    ) -> Result<BoxStream<'static, ReadBatchTask>> {
1842        FileReader::do_take_rows(
1843            prepared.column_infos,
1844            self.scheduler.clone(),
1845            self.cache.clone(),
1846            self.decoder_plugins.clone(),
1847            indices,
1848            batch_size,
1849            prepared.decoder_projection,
1850            FilterExpression::no_filter(),
1851            self.options.decoder_config.clone(),
1852            self.options.batch_size_bytes,
1853        )
1854        .await
1855    }
1856
1857    async fn read_ranges(
1858        &self,
1859        ranges: Vec<Range<u64>>,
1860        batch_size: u32,
1861        prepared: PreparedProjection,
1862        filter: FilterExpression,
1863    ) -> Result<BoxStream<'static, ReadBatchTask>> {
1864        FileReader::do_read_ranges(
1865            prepared.column_infos,
1866            self.scheduler.clone(),
1867            self.cache.clone(),
1868            self.decoder_plugins.clone(),
1869            ranges,
1870            batch_size,
1871            prepared.decoder_projection,
1872            filter,
1873            self.options.decoder_config.clone(),
1874            self.options.batch_size_bytes,
1875        )
1876        .await
1877    }
1878
1879    pub(crate) async fn read_prepared_tasks(
1880        &self,
1881        params: ReadBatchParams,
1882        batch_size: u32,
1883        prepared: PreparedProjection,
1884        read_len: u64,
1885        filter: FilterExpression,
1886    ) -> Result<Pin<Box<dyn Stream<Item = ReadBatchTask> + Send>>> {
1887        let verify_bound = |params: &ReadBatchParams, bound: u64, inclusive: bool| {
1888            if bound > read_len || (bound == read_len && inclusive) {
1889                Err(Error::invalid_input(format!(
1890                    "cannot read {params:?} from columns with {read_len} rows"
1891                )))
1892            } else {
1893                Ok(())
1894            }
1895        };
1896        match &params {
1897            ReadBatchParams::Indices(indices) => {
1898                for idx in indices {
1899                    match idx {
1900                        None => {
1901                            return Err(Error::invalid_input("Null value in indices array"));
1902                        }
1903                        Some(idx) => {
1904                            verify_bound(&params, idx as u64, true)?;
1905                        }
1906                    }
1907                }
1908                let indices = indices.iter().map(|idx| idx.unwrap() as u64).collect();
1909                self.take_rows(indices, batch_size, prepared).await
1910            }
1911            ReadBatchParams::Range(range) => {
1912                verify_bound(&params, range.end as u64, false)?;
1913                self.read_range(
1914                    range.start as u64..range.end as u64,
1915                    batch_size,
1916                    prepared,
1917                    filter,
1918                )
1919                .await
1920            }
1921            ReadBatchParams::Ranges(ranges) => {
1922                let mut ranges_u64 = Vec::with_capacity(ranges.len());
1923                for range in ranges.as_ref() {
1924                    verify_bound(&params, range.end, false)?;
1925                    ranges_u64.push(range.start..range.end);
1926                }
1927                self.read_ranges(ranges_u64, batch_size, prepared, filter)
1928                    .await
1929            }
1930            ReadBatchParams::RangeFrom(range) => {
1931                verify_bound(&params, range.start as u64, true)?;
1932                self.read_range(range.start as u64..read_len, batch_size, prepared, filter)
1933                    .await
1934            }
1935            ReadBatchParams::RangeTo(range) => {
1936                verify_bound(&params, range.end as u64, false)?;
1937                self.read_range(0..range.end as u64, batch_size, prepared, filter)
1938                    .await
1939            }
1940            ReadBatchParams::RangeFull => {
1941                self.read_range(0..read_len, batch_size, prepared, filter)
1942                    .await
1943            }
1944        }
1945    }
1946}
1947
1948impl ProjectedFileReader {
1949    pub(crate) fn base_projection(&self) -> &ReaderProjection {
1950        &self.core.base_projection
1951    }
1952
1953    pub(crate) async fn read_prepared_tasks(
1954        &self,
1955        params: ReadBatchParams,
1956        batch_size: u32,
1957        prepared: PreparedProjection,
1958        read_len: u64,
1959        filter: FilterExpression,
1960    ) -> Result<Pin<Box<dyn Stream<Item = ReadBatchTask> + Send>>> {
1961        self.core
1962            .read_prepared_tasks(params, batch_size, prepared, read_len, filter)
1963            .await
1964    }
1965
1966    /// Returns a clone of this reader using a different scheduler.
1967    pub fn with_scheduler(&self, scheduler: Arc<dyn EncodingsIo>) -> Self {
1968        Self {
1969            core: self.core.with_scheduler(scheduler),
1970        }
1971    }
1972
1973    /// Returns the number of rows in the file.
1974    pub fn num_rows(&self) -> u64 {
1975        self.core.num_rows()
1976    }
1977
1978    /// Returns the file schema visible to this reader.
1979    pub fn schema(&self) -> &Arc<Schema> {
1980        self.core.schema()
1981    }
1982
1983    /// Returns file statistics when this reader has full file metadata.
1984    pub fn file_statistics(&self) -> Option<FileStatistics> {
1985        self.core.metadata_provider.file_statistics()
1986    }
1987
1988    #[cfg(test)]
1989    pub(crate) fn metadata_index(&self) -> Option<&Arc<FileMetadataIndex>> {
1990        match &self.core.metadata_provider {
1991            FileMetadataProvider::Indexed(metadata_index) => Some(metadata_index),
1992            FileMetadataProvider::Full(_) => None,
1993        }
1994    }
1995
1996    /// Reads a global buffer by index.
1997    pub async fn read_global_buffer(&self, index: u32) -> Result<Bytes> {
1998        self.core.read_global_buffer(index).await
1999    }
2000}
2001
2002impl FileReader {
2003    #[cfg(test)]
2004    fn scheduler(&self) -> Arc<dyn EncodingsIo> {
2005        self.core.scheduler.clone()
2006    }
2007
2008    pub async fn try_open(
2009        scheduler: FileScheduler,
2010        base_projection: Option<ReaderProjection>,
2011        decoder_plugins: Arc<DecoderPlugins>,
2012        cache: &LanceCache,
2013        options: FileReaderOptions,
2014    ) -> Result<Self> {
2015        match Self::try_open_for_dispatch(
2016            scheduler,
2017            base_projection,
2018            decoder_plugins,
2019            cache,
2020            options,
2021        )
2022        .await?
2023        {
2024            versions::OpenedFileReader::V1 {
2025                major_version,
2026                minor_version,
2027            } => Err(Error::version_conflict(
2028                "Attempt to use the Lance current-format reader to read a v1 file".to_string(),
2029                major_version,
2030                minor_version,
2031            )),
2032            versions::OpenedFileReader::Current(reader) => Ok(reader),
2033        }
2034    }
2035
2036    pub(crate) async fn try_open_for_dispatch(
2037        scheduler: FileScheduler,
2038        base_projection: Option<ReaderProjection>,
2039        decoder_plugins: Arc<DecoderPlugins>,
2040        cache: &LanceCache,
2041        options: FileReaderOptions,
2042    ) -> Result<versions::OpenedFileReader> {
2043        let metadata = match Self::read_raw_metadata_for_dispatch(&scheduler).await? {
2044            RawFileMetadataOpen::Legacy {
2045                major_version,
2046                minor_version,
2047            } => {
2048                return Ok(versions::OpenedFileReader::V1 {
2049                    major_version,
2050                    minor_version,
2051                });
2052            }
2053            RawFileMetadataOpen::Current { version, metadata } => {
2054                Arc::new(versions::finish_metadata(version, metadata)?)
2055            }
2056        };
2057        let path = scheduler.reader().path().clone();
2058        let io = Arc::new(
2059            LanceEncodingsIo::new(scheduler).with_read_chunk_size(options.read_chunk_size),
2060        );
2061        Self::try_open_with_file_metadata(
2062            io,
2063            path,
2064            base_projection,
2065            decoder_plugins,
2066            metadata,
2067            cache,
2068            options,
2069        )
2070        .await
2071        .map(versions::OpenedFileReader::Current)
2072    }
2073
2074    pub async fn try_open_with_file_metadata(
2075        scheduler: Arc<dyn EncodingsIo>,
2076        path: Path,
2077        base_projection: Option<ReaderProjection>,
2078        decoder_plugins: Arc<DecoderPlugins>,
2079        metadata: Arc<CachedFileMetadata>,
2080        cache: &LanceCache,
2081        options: FileReaderOptions,
2082    ) -> Result<Self> {
2083        if metadata.version == ConcreteFileVersion::V1 {
2084            return Err(Error::version_conflict(
2085                "Attempt to use the Lance current-format reader with v1 metadata".to_string(),
2086                metadata.major_version,
2087                metadata.minor_version,
2088            ));
2089        }
2090        let read_projection = versions::read_projection(metadata.version)?;
2091        let has_explicit_projection = base_projection.is_some();
2092        let base_projection = base_projection.unwrap_or_else(|| {
2093            versions::reader_projection_from_whole_schema(&metadata.file_schema, metadata.version)
2094        });
2095        if has_explicit_projection {
2096            Self::validate_projection(&base_projection, &metadata)?;
2097        }
2098        let cache = Arc::new(cache.with_key_prefix(path.as_ref()));
2099        let core = DecodeEngine::try_new(
2100            scheduler,
2101            base_projection,
2102            decoder_plugins,
2103            FileMetadataProvider::Full(metadata.clone()),
2104            read_projection,
2105            cache,
2106            options,
2107        )?;
2108        Ok(Self { core, metadata })
2109    }
2110
2111    pub async fn read_all_metadata(scheduler: &FileScheduler) -> Result<CachedFileMetadata> {
2112        match Self::read_raw_metadata_for_dispatch(scheduler).await? {
2113            RawFileMetadataOpen::Legacy {
2114                major_version,
2115                minor_version,
2116            } => Err(Error::version_conflict(
2117                "Attempt to use the Lance current-format reader to read v1 metadata".to_string(),
2118                major_version,
2119                minor_version,
2120            )),
2121            RawFileMetadataOpen::Current { version, metadata } => {
2122                versions::finish_metadata(version, metadata)
2123            }
2124        }
2125    }
2126
2127    pub async fn read_metadata_index(scheduler: &FileScheduler) -> Result<FileMetadataIndex> {
2128        let index = Self::read_raw_metadata_index(scheduler).await?;
2129        versions::finish_metadata_index(index)
2130    }
2131
2132    pub async fn read_metadata_index_with_schema(
2133        scheduler: &FileScheduler,
2134        file_schema: Arc<Schema>,
2135        num_rows: u64,
2136    ) -> Result<FileMetadataIndex> {
2137        let index =
2138            Self::read_raw_metadata_index_with_schema(scheduler, file_schema, num_rows).await?;
2139        versions::finish_metadata_index(index)
2140    }
2141
2142    pub fn version(&self) -> ConcreteFileVersion {
2143        self.metadata.version
2144    }
2145
2146    async fn prepare(&self, projection: ReaderProjection) -> Result<(PreparedProjection, u64)> {
2147        self.core
2148            .read_projection
2149            .prepare(
2150                &self.core.metadata_provider,
2151                &projection,
2152                &self.core.scheduler,
2153                &self.core.cache,
2154            )
2155            .await
2156    }
2157
2158    pub async fn read_tasks(
2159        &self,
2160        params: ReadBatchParams,
2161        batch_size: u32,
2162        projection: Option<ReaderProjection>,
2163        filter: FilterExpression,
2164    ) -> Result<Pin<Box<dyn Stream<Item = ReadBatchTask> + Send>>> {
2165        let projection = projection.unwrap_or_else(|| self.base_projection().clone());
2166        let (prepared, read_len) = self.prepare(projection).await?;
2167        self.read_prepared_tasks(params, batch_size, prepared, read_len, filter)
2168            .await
2169    }
2170
2171    pub async fn read_stream_projected(
2172        &self,
2173        params: ReadBatchParams,
2174        batch_size: u32,
2175        batch_readahead: u32,
2176        projection: ReaderProjection,
2177        filter: FilterExpression,
2178    ) -> Result<Pin<Box<dyn RecordBatchStream>>> {
2179        let schema = projection.schema.clone();
2180        let tasks = self
2181            .read_tasks(params, batch_size, Some(projection), filter)
2182            .await?;
2183        Ok(tasks_to_record_batch_stream(schema, tasks, batch_readahead))
2184    }
2185
2186    pub fn read_stream_projected_blocking(
2187        &self,
2188        params: ReadBatchParams,
2189        batch_size: u32,
2190        projection: Option<ReaderProjection>,
2191        filter: FilterExpression,
2192    ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
2193        let projection = projection.unwrap_or_else(|| self.base_projection().clone());
2194        Self::validate_projection(&projection, self.metadata())?;
2195        let prepared = self.full_projection(projection);
2196        let read_len = self.core.read_projection.read_length(&prepared)?;
2197        self.read_prepared_blocking(params, batch_size, prepared, read_len, filter)
2198    }
2199
2200    pub async fn read_stream(
2201        &self,
2202        params: ReadBatchParams,
2203        batch_size: u32,
2204        batch_readahead: u32,
2205        filter: FilterExpression,
2206    ) -> Result<Pin<Box<dyn RecordBatchStream>>> {
2207        self.read_stream_projected(
2208            params,
2209            batch_size,
2210            batch_readahead,
2211            self.base_projection().clone(),
2212            filter,
2213        )
2214        .await
2215    }
2216}
2217
2218impl ProjectedFileReader {
2219    pub async fn try_open(
2220        scheduler: FileScheduler,
2221        base_projection: Option<ReaderProjection>,
2222        decoder_plugins: Arc<DecoderPlugins>,
2223        cache: &LanceCache,
2224        options: FileReaderOptions,
2225    ) -> Result<Self> {
2226        let base_projection = base_projection.ok_or_else(|| {
2227            Error::invalid_input("ProjectedReader requires an explicit base projection")
2228        })?;
2229        let metadata_index = Arc::new(FileReader::read_metadata_index(&scheduler).await?);
2230        let path = scheduler.reader().path().clone();
2231        let io = Arc::new(
2232            LanceEncodingsIo::new(scheduler).with_read_chunk_size(options.read_chunk_size),
2233        );
2234        Self::try_open_with_metadata_index(
2235            io,
2236            path,
2237            Some(base_projection),
2238            decoder_plugins,
2239            metadata_index,
2240            cache,
2241            options,
2242        )
2243        .await
2244    }
2245
2246    pub async fn try_open_with_metadata_index(
2247        scheduler: Arc<dyn EncodingsIo>,
2248        path: Path,
2249        base_projection: Option<ReaderProjection>,
2250        decoder_plugins: Arc<DecoderPlugins>,
2251        metadata_index: Arc<FileMetadataIndex>,
2252        cache: &LanceCache,
2253        options: FileReaderOptions,
2254    ) -> Result<Self> {
2255        if metadata_index.version == ConcreteFileVersion::V1 {
2256            return Err(Error::version_conflict(
2257                "Attempt to use the Lance projected current-format reader with v1 metadata"
2258                    .to_string(),
2259                0,
2260                2,
2261            ));
2262        }
2263        let base_projection = base_projection.ok_or_else(|| {
2264            Error::invalid_input("ProjectedReader requires an explicit base projection")
2265        })?;
2266        let read_projection = versions::read_projection(metadata_index.version)?;
2267        read_projection.validate_indexed(&base_projection, &metadata_index)?;
2268        let cache = Arc::new(cache.with_key_prefix(path.as_ref()));
2269        let core = DecodeEngine::try_new(
2270            scheduler,
2271            base_projection,
2272            decoder_plugins,
2273            FileMetadataProvider::Indexed(metadata_index),
2274            read_projection,
2275            cache,
2276            options,
2277        )?;
2278        Ok(Self { core })
2279    }
2280
2281    pub async fn try_open_with_file_metadata(
2282        scheduler: Arc<dyn EncodingsIo>,
2283        path: Path,
2284        base_projection: Option<ReaderProjection>,
2285        decoder_plugins: Arc<DecoderPlugins>,
2286        metadata: Arc<CachedFileMetadata>,
2287        cache: &LanceCache,
2288        options: FileReaderOptions,
2289    ) -> Result<Self> {
2290        if metadata.version == ConcreteFileVersion::V1 {
2291            return Err(Error::version_conflict(
2292                "Attempt to use the Lance projected current-format reader with v1 metadata"
2293                    .to_string(),
2294                metadata.major_version,
2295                metadata.minor_version,
2296            ));
2297        }
2298        let read_projection = versions::read_projection(metadata.version)?;
2299        let has_explicit_projection = base_projection.is_some();
2300        let base_projection = base_projection.unwrap_or_else(|| {
2301            versions::reader_projection_from_whole_schema(&metadata.file_schema, metadata.version)
2302        });
2303        if has_explicit_projection {
2304            FileReader::validate_projection(&base_projection, &metadata)?;
2305        }
2306        let cache = Arc::new(cache.with_key_prefix(path.as_ref()));
2307        let core = DecodeEngine::try_new(
2308            scheduler,
2309            base_projection,
2310            decoder_plugins,
2311            FileMetadataProvider::Full(metadata),
2312            read_projection,
2313            cache,
2314            options,
2315        )?;
2316        Ok(Self { core })
2317    }
2318
2319    pub fn version(&self) -> ConcreteFileVersion {
2320        self.core.metadata_provider.version()
2321    }
2322
2323    pub async fn read_tasks(
2324        &self,
2325        params: ReadBatchParams,
2326        batch_size: u32,
2327        projection: Option<ReaderProjection>,
2328        filter: FilterExpression,
2329    ) -> Result<Pin<Box<dyn Stream<Item = ReadBatchTask> + Send>>> {
2330        let projection = projection.unwrap_or_else(|| self.base_projection().clone());
2331        let (prepared, read_len) = self
2332            .core
2333            .read_projection
2334            .prepare(
2335                &self.core.metadata_provider,
2336                &projection,
2337                &self.core.scheduler,
2338                &self.core.cache,
2339            )
2340            .await?;
2341        self.read_prepared_tasks(params, batch_size, prepared, read_len, filter)
2342            .await
2343    }
2344}
2345
2346/// Inspects a page and returns a String describing the page's encoding
2347pub fn describe_encoding(page: &pbfile::column_metadata::Page) -> String {
2348    if let Some(encoding) = &page.encoding {
2349        if let Some(style) = &encoding.location {
2350            match style {
2351                pbfile::encoding::Location::Indirect(indirect) => {
2352                    format!(
2353                        "IndirectEncoding(pos={},size={})",
2354                        indirect.buffer_location, indirect.buffer_length
2355                    )
2356                }
2357                pbfile::encoding::Location::Direct(direct) => {
2358                    let encoding_any =
2359                        prost_types::Any::decode(Bytes::from(direct.encoding.clone()))
2360                            .expect("failed to deserialize encoding as protobuf");
2361                    if encoding_any.type_url == "/lance.encodings.ArrayEncoding" {
2362                        let encoding = encoding_any.to_msg::<pbenc::ArrayEncoding>();
2363                        match encoding {
2364                            Ok(encoding) => {
2365                                format!("{:#?}", encoding)
2366                            }
2367                            Err(err) => {
2368                                format!("Unsupported(decode_err={})", err)
2369                            }
2370                        }
2371                    } else if encoding_any.type_url == "/lance.encodings21.PageLayout" {
2372                        let encoding = encoding_any.to_msg::<pbenc21::PageLayout>();
2373                        match encoding {
2374                            Ok(encoding) => {
2375                                format!("{:#?}", encoding)
2376                            }
2377                            Err(err) => {
2378                                format!("Unsupported(decode_err={})", err)
2379                            }
2380                        }
2381                    } else if encoding_any.type_url == "/lance.encodings23.PlanarLayout" {
2382                        match encoding_any.to_msg::<lance_encoding::format::pb23::PlanarLayout>() {
2383                            Ok(encoding) => format!("{encoding:#?}"),
2384                            Err(err) => format!("Unsupported(decode_err={err})"),
2385                        }
2386                    } else {
2387                        format!("Unrecognized(type_url={})", encoding_any.type_url)
2388                    }
2389                }
2390                pbfile::encoding::Location::None(_) => "NoEncodingDescription".to_string(),
2391            }
2392        } else {
2393            "MISSING STYLE".to_string()
2394        }
2395    } else {
2396        "MISSING".to_string()
2397    }
2398}
2399
2400pub trait EncodedBatchReaderExt {
2401    fn try_from_mini_lance(bytes: Bytes, schema: &Schema) -> Result<Self>
2402    where
2403        Self: Sized;
2404    fn try_from_self_described_lance(bytes: Bytes) -> Result<Self>
2405    where
2406        Self: Sized;
2407}
2408
2409impl EncodedBatchReaderExt for EncodedBatch {
2410    fn try_from_mini_lance(bytes: Bytes, schema: &Schema) -> Result<Self>
2411    where
2412        Self: Sized,
2413    {
2414        let footer = FileReader::decode_footer(&bytes)?;
2415        let file_version = FileReader::current_file_version(&footer)?;
2416        let projection = versions::reader_projection_from_whole_schema(schema, file_version);
2417
2418        // Next, read the metadata for the columns
2419        // This is both the column metadata and the CMO table
2420        let column_metadata_start = footer.column_meta_start as usize;
2421        let column_metadata_end = footer.global_buff_offsets_start as usize;
2422        let column_metadata_bytes = bytes.slice(column_metadata_start..column_metadata_end);
2423        let column_metadatas =
2424            FileReader::read_all_column_metadata(column_metadata_bytes, &footer)?;
2425
2426        let page_table = versions::decode_column_metadata(file_version, &column_metadatas)?;
2427
2428        Ok(Self {
2429            data: bytes,
2430            num_rows: page_table
2431                .first()
2432                .map(|col| col.page_infos.iter().map(|page| page.num_rows).sum::<u64>())
2433                .unwrap_or(0),
2434            page_table,
2435            top_level_columns: projection.column_indices,
2436            schema: Arc::new(schema.clone()),
2437        })
2438    }
2439
2440    fn try_from_self_described_lance(bytes: Bytes) -> Result<Self>
2441    where
2442        Self: Sized,
2443    {
2444        let footer = FileReader::decode_footer(&bytes)?;
2445        let file_version = FileReader::current_file_version(&footer)?;
2446
2447        let file_len = bytes.len() as u64;
2448        let gbo_table = FileReader::do_decode_gbo_table(
2449            &bytes.slice(footer.global_buff_offsets_start as usize..),
2450            &footer,
2451        )?;
2452        FileReader::validate_gbo_table(&gbo_table, file_len, file_version)?;
2453        if gbo_table.is_empty() {
2454            return Err(Error::internal(
2455                "File did not contain any global buffers, schema expected".to_string(),
2456            ));
2457        }
2458        let schema_range = gbo_table[0].checked_range(0, file_len)?;
2459        let schema_start = schema_range.start as usize;
2460        let schema_end = schema_range.end as usize;
2461
2462        let schema_bytes = bytes.slice(schema_start..schema_end);
2463        let (_, schema) = FileReader::decode_schema(schema_bytes)?;
2464        let projection = versions::reader_projection_from_whole_schema(&schema, file_version);
2465
2466        // Next, read the metadata for the columns
2467        // This is both the column metadata and the CMO table
2468        let column_metadata_start = footer.column_meta_start as usize;
2469        let column_metadata_end = footer.global_buff_offsets_start as usize;
2470        let column_metadata_bytes = bytes.slice(column_metadata_start..column_metadata_end);
2471        let column_metadatas =
2472            FileReader::read_all_column_metadata(column_metadata_bytes, &footer)?;
2473
2474        let page_table = versions::decode_column_metadata(file_version, &column_metadatas)?;
2475
2476        Ok(Self {
2477            data: bytes,
2478            num_rows: page_table
2479                .first()
2480                .map(|col| col.page_infos.iter().map(|page| page.num_rows).sum::<u64>())
2481                .unwrap_or(0),
2482            page_table,
2483            top_level_columns: projection.column_indices,
2484            schema: Arc::new(schema),
2485        })
2486    }
2487}
2488
2489#[cfg(test)]
2490mod tests {
2491    use std::{
2492        collections::{BTreeMap, HashMap},
2493        pin::Pin,
2494        sync::Arc,
2495    };
2496
2497    use arrow_array::{
2498        Array, DictionaryArray, FixedSizeListArray, Float32Array, Int8Array, Int32Array, ListArray,
2499        RecordBatch, RecordBatchIterator, StringArray, UInt32Array,
2500        cast::AsArray,
2501        types::{Float64Type, Int8Type, Int32Type},
2502    };
2503    use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
2504    use arrow_schema::{DataType, Field, Fields, Schema as ArrowSchema};
2505    use arrow_select::concat::concat_batches;
2506    use bytes::Bytes;
2507    use futures::{StreamExt, prelude::stream::TryStreamExt};
2508    use lance_arrow::{BLOB_META_KEY, RecordBatchExt};
2509    use lance_core::{ArrowResult, datatypes::Schema};
2510    use lance_datagen::{ArrayGeneratorExt, BatchCount, ByteCount, RowCount, array, gen_batch};
2511    use lance_encoding::{
2512        constants::{STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_SPARSE},
2513        decoder::{
2514            DecodeBatchScheduler, DecoderPlugins, EncodedBatchLayout, FilterExpression,
2515            PageEncoding, ReadBatchTask, decode_batch,
2516        },
2517        encoder::{EncodedBatch, EncodingOptions, encode_batch},
2518        format::metadata,
2519    };
2520    use lance_io::{scheduler::IoStats, stream::RecordBatchStream, utils::CachedFileSize};
2521    use log::debug;
2522    use rstest::rstest;
2523    use tokio::sync::mpsc;
2524
2525    use crate::reader::{
2526        EncodedBatchReaderExt, FileReader, FileReaderOptions, ProjectedFileReader, ReaderProjection,
2527    };
2528    use crate::testing::{FsFixture, WrittenFile, test_cache, write_lance_file};
2529    use crate::version::{ConcreteFileVersion, LanceFileVersion};
2530    use crate::versions;
2531    use crate::writer::{FileWriterOptions, PAGE_BUFFER_ALIGNMENT};
2532    use lance_encoding::decoder::DecoderConfig;
2533
2534    fn footer_version(bytes: &[u8]) -> (u16, u16) {
2535        let version_start = bytes.len() - 8;
2536        (
2537            u16::from_le_bytes([bytes[version_start], bytes[version_start + 1]]),
2538            u16::from_le_bytes([bytes[version_start + 2], bytes[version_start + 3]]),
2539        )
2540    }
2541
2542    #[rstest::rstest]
2543    #[case("nullable")]
2544    #[case("rle")]
2545    #[case("dictionary")]
2546    #[case("binary_range")]
2547    #[case("binary_delta")]
2548    #[case("vector_dictionary")]
2549    #[tokio::test]
2550    async fn sparse_file_writer_reader_scan_range_and_take_roundtrip(#[case] kind: &str) {
2551        let fs = FsFixture::default();
2552        let sparse_metadata = HashMap::from([(
2553            STRUCTURAL_ENCODING_META_KEY.to_string(),
2554            STRUCTURAL_ENCODING_SPARSE.to_string(),
2555        )]);
2556        let value_field =
2557            Field::new("values", DataType::Int32, true).with_metadata(sparse_metadata.clone());
2558        let item_field = Arc::new(Field::new("item", DataType::Int32, true));
2559        let list_field = Field::new("items", DataType::List(item_field.clone()), true)
2560            .with_metadata(sparse_metadata);
2561        let arrow_schema = Arc::new(ArrowSchema::new(vec![value_field, list_field]));
2562        let list = ListArray::try_new(
2563            item_field,
2564            OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 2, 2, 2, 3, 3, 5])),
2565            Arc::new(Int32Array::from(vec![
2566                Some(1),
2567                None,
2568                Some(3),
2569                Some(4),
2570                Some(5),
2571            ])),
2572            Some(NullBuffer::from(vec![true, false, true, true, true, true])),
2573        )
2574        .unwrap();
2575        let mut batch = RecordBatch::try_new(
2576            arrow_schema.clone(),
2577            vec![
2578                Arc::new(Int32Array::from(vec![
2579                    Some(10),
2580                    None,
2581                    Some(30),
2582                    Some(40),
2583                    None,
2584                    Some(60),
2585                ])),
2586                Arc::new(list),
2587            ],
2588        )
2589        .unwrap();
2590        if matches!(kind, "rle" | "dictionary") {
2591            let counts = (0..2048).map(|i| {
2592                if kind == "rle" {
2593                    i / 128 + 1
2594                } else {
2595                    [1, 31, 127][i % 3]
2596                }
2597            });
2598            let mut offsets = vec![0_i32];
2599            for count in counts {
2600                offsets.push(offsets.last().unwrap() + count as i32);
2601            }
2602            let values = Int32Array::from_iter_values(0..*offsets.last().unwrap());
2603            let list = ListArray::try_new(
2604                Arc::new(Field::new("item", DataType::Int32, true)),
2605                OffsetBuffer::new(ScalarBuffer::from(offsets)),
2606                Arc::new(values),
2607                None,
2608            )
2609            .unwrap();
2610            batch = RecordBatch::try_new(
2611                arrow_schema.clone(),
2612                vec![
2613                    Arc::new(Int32Array::from_iter_values(0..2048)),
2614                    Arc::new(list),
2615                ],
2616            )
2617            .unwrap();
2618        }
2619        if kind.starts_with("binary_") {
2620            let metadata = HashMap::from([
2621                (
2622                    STRUCTURAL_ENCODING_META_KEY.to_string(),
2623                    STRUCTURAL_ENCODING_SPARSE.to_string(),
2624                ),
2625                ("lance-encoding:compression".to_string(), "none".to_string()),
2626                (
2627                    "lance-encoding:dict-divisor".to_string(),
2628                    "100000".to_string(),
2629                ),
2630            ]);
2631            let values = Arc::new(StringArray::from_iter_values((0..2048).map(|i| {
2632                let len = if kind == "binary_range" {
2633                    8
2634                } else {
2635                    [4, 10, 5, 8][i % 4]
2636                };
2637                char::from(b'a' + (i % 26) as u8).to_string().repeat(len)
2638            })));
2639            let item =
2640                Arc::new(Field::new("item", DataType::Utf8, true).with_metadata(metadata.clone()));
2641            let lists = ListArray::try_new(
2642                item.clone(),
2643                OffsetBuffer::new(ScalarBuffer::from((0..=2048).collect::<Vec<i32>>())),
2644                values.clone(),
2645                None,
2646            )
2647            .unwrap();
2648            batch = RecordBatch::try_new(
2649                Arc::new(ArrowSchema::new(vec![
2650                    Field::new("values", DataType::Utf8, true).with_metadata(metadata.clone()),
2651                    Field::new("items", DataType::List(item), true).with_metadata(metadata),
2652                ])),
2653                vec![values, Arc::new(lists)],
2654            )
2655            .unwrap();
2656        }
2657        if kind == "vector_dictionary" {
2658            let vectors = Arc::new(FixedSizeListArray::new(
2659                Arc::new(Field::new("item", DataType::Float32, true)),
2660                64,
2661                Arc::new(Float32Array::from(
2662                    (0..2048 * 64)
2663                        .map(|i| (i % 7 != 0).then_some((i % (17 * 64)) as f32 / 10.0))
2664                        .collect::<Vec<_>>(),
2665                )),
2666                Some(NullBuffer::from(
2667                    (0..2048).map(|i| i % 8 != 0).collect::<Vec<_>>(),
2668                )),
2669            ));
2670            let metadata = HashMap::from([(
2671                STRUCTURAL_ENCODING_META_KEY.to_string(),
2672                STRUCTURAL_ENCODING_SPARSE.to_string(),
2673            )]);
2674            let item = Arc::new(
2675                Field::new("item", vectors.data_type().clone(), true)
2676                    .with_metadata(metadata.clone()),
2677            );
2678            let lists = Arc::new(ListArray::new(
2679                item.clone(),
2680                OffsetBuffer::new(ScalarBuffer::from((0..=2048).collect::<Vec<i32>>())),
2681                vectors.clone(),
2682                None,
2683            ));
2684            batch = RecordBatch::try_new(
2685                Arc::new(ArrowSchema::new(vec![
2686                    Field::new("values", vectors.data_type().clone(), true)
2687                        .with_metadata(metadata.clone()),
2688                    Field::new("items", DataType::List(item), false).with_metadata(metadata),
2689                ])),
2690                vec![vectors, lists],
2691            )
2692            .unwrap();
2693        }
2694        let input = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema());
2695        write_lance_file(
2696            input,
2697            &fs,
2698            ConcreteFileVersion::V2_3,
2699            FileWriterOptions::default(),
2700        )
2701        .await;
2702
2703        let file_scheduler = fs
2704            .scheduler
2705            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
2706            .await
2707            .unwrap();
2708        let file_reader = FileReader::try_open(
2709            file_scheduler,
2710            None,
2711            Arc::<DecoderPlugins>::default(),
2712            &test_cache(),
2713            FileReaderOptions::default(),
2714        )
2715        .await
2716        .unwrap();
2717        assert_eq!(file_reader.metadata().column_infos.len(), 2);
2718        if kind == "vector_dictionary" {
2719            for column in &file_reader.metadata().column_infos {
2720                for page in column.page_infos.iter() {
2721                    let PageEncoding::Planar(layout) = &page.encoding else {
2722                        panic!("expected Planar vector page");
2723                    };
2724                    let Some(lance_encoding::format::pb23::value_plane::Codec::Chunked(values)) =
2725                        layout.values.as_ref().and_then(|v| v.codec.as_ref())
2726                    else {
2727                        panic!("expected dictionary vector indices");
2728                    };
2729                    assert!(values.dictionary.is_some());
2730                    assert!(values.num_dictionary_items < 1024);
2731                }
2732            }
2733        }
2734        if matches!(kind, "rle" | "dictionary") {
2735            let PageEncoding::Planar(layout) =
2736                &file_reader.metadata().column_infos[1].page_infos[0].encoding
2737            else {
2738                panic!("expected structural page")
2739            };
2740            let Some(sparse) = layout.sparse() else {
2741                panic!("expected sparse page")
2742            };
2743            let counts = sparse
2744                .layers
2745                .iter()
2746                .find_map(|layer| {
2747                    if let Some(
2748                        lance_encoding::format::pb23::sparse_structural_layer::Layer::List(list),
2749                    ) = &layer.layer
2750                    {
2751                        list.counts.as_ref()
2752                    } else {
2753                        None
2754                    }
2755                })
2756                .unwrap();
2757            assert!(
2758                matches!(
2759                    (kind, &counts.compression),
2760                    (
2761                        "rle",
2762                        Some(metadata::compressive_encoding::Compression::Rle(_))
2763                    ) | (
2764                        "dictionary",
2765                        Some(metadata::compressive_encoding::Compression::Dictionary(_))
2766                    )
2767                ),
2768                "unexpected count codec: {counts:?}"
2769            );
2770        }
2771        if kind.starts_with("binary_") {
2772            for column in &file_reader.metadata().column_infos {
2773                let PageEncoding::Planar(layout) = &column.page_infos[0].encoding else {
2774                    panic!("expected structural page")
2775                };
2776                let Some(lance_encoding::format::pb23::value_plane::Codec::Chunked(values)) =
2777                    layout.values.as_ref().and_then(|v| v.codec.as_ref())
2778                else {
2779                    panic!("expected chunked values")
2780                };
2781                let Some(metadata::compressive_encoding::Compression::Variable(variable)) =
2782                    &values.compression.as_ref().unwrap().compression
2783                else {
2784                    panic!("expected variable values")
2785                };
2786                let offsets = &variable.offsets.as_ref().unwrap().compression;
2787                if kind == "binary_range" {
2788                    assert!(matches!(
2789                        offsets,
2790                        Some(metadata::compressive_encoding::Compression::Range(_))
2791                    ));
2792                } else {
2793                    assert!(matches!(
2794                        offsets,
2795                        Some(metadata::compressive_encoding::Compression::Delta(_))
2796                    ));
2797                }
2798            }
2799        }
2800        assert!(
2801            file_reader
2802                .metadata()
2803                .column_infos
2804                .iter()
2805                .flat_map(|column| column.page_infos.iter())
2806                .all(|page| {
2807                    matches!(
2808                        &page.encoding,
2809                        PageEncoding::Planar(layout)
2810                            if layout.structure.is_some() && layout.values.is_some()
2811                    )
2812                })
2813        );
2814
2815        let scan = file_reader
2816            .read_stream(
2817                lance_io::ReadBatchParams::RangeFull,
2818                1024,
2819                1,
2820                FilterExpression::no_filter(),
2821            )
2822            .await
2823            .unwrap()
2824            .try_collect::<Vec<_>>()
2825            .await
2826            .unwrap();
2827        assert_eq!(
2828            arrow_select::concat::concat_batches(&batch.schema(), &scan).unwrap(),
2829            batch
2830        );
2831
2832        let range = file_reader
2833            .read_stream(
2834                lance_io::ReadBatchParams::Range(1..5),
2835                1024,
2836                1,
2837                FilterExpression::no_filter(),
2838            )
2839            .await
2840            .unwrap()
2841            .try_collect::<Vec<_>>()
2842            .await
2843            .unwrap();
2844        assert_eq!(range, vec![batch.slice(1, 4)]);
2845
2846        let indices = UInt32Array::from(vec![0, 3, 5]);
2847        let take = file_reader
2848            .read_stream(
2849                lance_io::ReadBatchParams::Indices(indices.clone()),
2850                1024,
2851                1,
2852                FilterExpression::no_filter(),
2853            )
2854            .await
2855            .unwrap()
2856            .try_collect::<Vec<_>>()
2857            .await
2858            .unwrap();
2859        assert_eq!(take, vec![batch.take(&indices).unwrap()]);
2860    }
2861
2862    #[tokio::test]
2863    async fn full_int8_dictionary_v2_2_roundtrip() {
2864        let fs = FsFixture::default();
2865        let values = Arc::new(StringArray::from(
2866            (0..=i8::MAX)
2867                .map(|value| format!("value-{value}"))
2868                .collect::<Vec<_>>(),
2869        ));
2870        let keys = Int8Array::from((0..=i8::MAX).collect::<Vec<_>>());
2871        let dictionary = Arc::new(DictionaryArray::<Int8Type>::new(keys, values));
2872        let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
2873            "dictionary",
2874            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2875            true,
2876        )]));
2877        let batch = RecordBatch::try_new(arrow_schema.clone(), vec![dictionary]).unwrap();
2878
2879        write_lance_file(
2880            RecordBatchIterator::new([Ok(batch.clone())], arrow_schema),
2881            &fs,
2882            ConcreteFileVersion::V2_2,
2883            FileWriterOptions::default(),
2884        )
2885        .await;
2886
2887        let file_scheduler = fs
2888            .scheduler
2889            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
2890            .await
2891            .unwrap();
2892        let file_reader = FileReader::try_open(
2893            file_scheduler,
2894            None,
2895            Arc::<DecoderPlugins>::default(),
2896            &test_cache(),
2897            FileReaderOptions::default(),
2898        )
2899        .await
2900        .unwrap();
2901        let actual = file_reader
2902            .read_stream(
2903                lance_io::ReadBatchParams::RangeFull,
2904                1024,
2905                1,
2906                FilterExpression::no_filter(),
2907            )
2908            .await
2909            .unwrap()
2910            .try_collect::<Vec<_>>()
2911            .await
2912            .unwrap();
2913
2914        assert_eq!(actual, vec![batch]);
2915    }
2916
2917    async fn create_some_file(fs: &FsFixture, version: ConcreteFileVersion) -> WrittenFile {
2918        let location_type = DataType::Struct(Fields::from(vec![
2919            Field::new("x", DataType::Float64, true),
2920            Field::new("y", DataType::Float64, true),
2921        ]));
2922        let categories_type = DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)));
2923
2924        let mut reader = gen_batch()
2925            .col("score", array::rand::<Float64Type>())
2926            .col("location", array::rand_type(&location_type))
2927            .col("categories", array::rand_type(&categories_type))
2928            .col("binary", array::rand_type(&DataType::Binary));
2929        if version == ConcreteFileVersion::V2_0 {
2930            reader = reader.col("large_bin", array::rand_type(&DataType::LargeBinary));
2931        }
2932        let reader = reader.into_reader_rows(RowCount::from(1000), BatchCount::from(100));
2933
2934        write_lance_file(reader, fs, version, FileWriterOptions::default()).await
2935    }
2936
2937    async fn create_wide_direct_file(fs: &FsFixture, num_columns: usize) -> WrittenFile {
2938        let mut reader = gen_batch();
2939        for column_idx in 0..num_columns {
2940            reader = reader.col(format!("c{column_idx}"), array::step::<Int32Type>());
2941        }
2942        let reader = reader.into_reader_rows(RowCount::from(1000), BatchCount::from(100));
2943
2944        write_lance_file(
2945            reader,
2946            fs,
2947            ConcreteFileVersion::V2_1,
2948            FileWriterOptions::default(),
2949        )
2950        .await
2951    }
2952
2953    async fn create_wide_fixed_size_list_file(fs: &FsFixture, num_columns: usize) -> WrittenFile {
2954        let data_type =
2955            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4);
2956        let mut reader = gen_batch();
2957        for column_idx in 0..num_columns {
2958            reader = reader.col(
2959                format!("c{column_idx}"),
2960                array::rand_type(&data_type).with_random_nulls(0.1),
2961            );
2962        }
2963        let reader = reader.into_reader_rows(RowCount::from(64), BatchCount::from(4));
2964
2965        write_lance_file(
2966            reader,
2967            fs,
2968            ConcreteFileVersion::V2_1,
2969            FileWriterOptions::default(),
2970        )
2971        .await
2972    }
2973
2974    async fn create_wide_structural_file(fs: &FsFixture, num_groups: usize) -> WrittenFile {
2975        let struct_type = DataType::Struct(Fields::from(vec![
2976            Field::new("x", DataType::Int32, true),
2977            Field::new("y", DataType::Int32, true),
2978        ]));
2979        let list_type = DataType::List(Arc::new(Field::new("item", DataType::Int32, true)));
2980        let mut reader = gen_batch();
2981        for group_idx in 0..num_groups {
2982            reader = reader
2983                .col(
2984                    format!("s{group_idx}"),
2985                    array::rand_type(&struct_type).with_random_nulls(0.5),
2986                )
2987                .col(
2988                    format!("l{group_idx}"),
2989                    array::rand_type(&list_type).with_random_nulls(0.5),
2990                );
2991        }
2992        let reader = reader.into_reader_rows(RowCount::from(64), BatchCount::from(4));
2993
2994        write_lance_file(
2995            reader,
2996            fs,
2997            ConcreteFileVersion::V2_1,
2998            FileWriterOptions::default(),
2999        )
3000        .await
3001    }
3002
3003    type Transformer = Box<dyn Fn(&RecordBatch) -> RecordBatch>;
3004
3005    async fn verify_expected(
3006        expected: &[RecordBatch],
3007        mut actual: Pin<Box<dyn RecordBatchStream>>,
3008        read_size: u32,
3009        transform: Option<Transformer>,
3010    ) {
3011        let mut remaining = expected.iter().map(|batch| batch.num_rows()).sum::<usize>() as u32;
3012        let mut expected_iter = expected.iter().map(|batch| {
3013            if let Some(transform) = &transform {
3014                transform(batch)
3015            } else {
3016                batch.clone()
3017            }
3018        });
3019        let mut next_expected = expected_iter.next().unwrap().clone();
3020        while let Some(actual) = actual.next().await {
3021            let mut actual = actual.unwrap();
3022            let mut rows_to_verify = actual.num_rows() as u32;
3023            let expected_length = remaining.min(read_size);
3024            assert_eq!(expected_length, rows_to_verify);
3025
3026            while rows_to_verify > 0 {
3027                let next_slice_len = (next_expected.num_rows() as u32).min(rows_to_verify);
3028                assert_eq!(
3029                    next_expected.slice(0, next_slice_len as usize),
3030                    actual.slice(0, next_slice_len as usize)
3031                );
3032                remaining -= next_slice_len;
3033                rows_to_verify -= next_slice_len;
3034                if remaining > 0 {
3035                    if next_slice_len == next_expected.num_rows() as u32 {
3036                        next_expected = expected_iter.next().unwrap().clone();
3037                    } else {
3038                        next_expected = next_expected.slice(
3039                            next_slice_len as usize,
3040                            next_expected.num_rows() - next_slice_len as usize,
3041                        );
3042                    }
3043                }
3044                if rows_to_verify > 0 {
3045                    actual = actual.slice(
3046                        next_slice_len as usize,
3047                        actual.num_rows() - next_slice_len as usize,
3048                    );
3049                }
3050            }
3051        }
3052        assert_eq!(remaining, 0);
3053    }
3054
3055    async fn collect_read_tasks(
3056        tasks: Pin<Box<dyn futures::Stream<Item = ReadBatchTask> + Send>>,
3057        readahead: usize,
3058    ) -> Vec<RecordBatch> {
3059        tasks
3060            .map(|task| task.task)
3061            .buffered(readahead)
3062            .try_collect::<Vec<_>>()
3063            .await
3064            .unwrap()
3065    }
3066
3067    /// Writes `batch` to a fresh file, overwrites `patch` bytes at `patch_offset`
3068    /// into the single occurrence of `pattern`, and reads the file back with the
3069    /// default reader configuration.
3070    async fn read_file_with_mutated_bytes(
3071        version: LanceFileVersion,
3072        batch: RecordBatch,
3073        pattern: &[u8],
3074        patch_offset: usize,
3075        patch: &[u8],
3076    ) -> lance_core::Result<Vec<RecordBatch>> {
3077        let fs = FsFixture::default();
3078        let schema = batch.schema();
3079        write_lance_file(
3080            RecordBatchIterator::new(vec![Ok(batch)], schema),
3081            &fs,
3082            version.resolve(),
3083            FileWriterOptions::default(),
3084        )
3085        .await;
3086
3087        let mut bytes = fs
3088            .object_store
3089            .read_one_all(&fs.tmp_path)
3090            .await
3091            .unwrap()
3092            .to_vec();
3093        let matches = bytes
3094            .windows(pattern.len())
3095            .enumerate()
3096            .filter_map(|(position, window)| (window == pattern).then_some(position))
3097            .collect::<Vec<_>>();
3098        assert_eq!(
3099            matches.len(),
3100            1,
3101            "expected the byte pattern to appear exactly once in the file"
3102        );
3103        let patch_start = matches[0] + patch_offset;
3104        bytes[patch_start..patch_start + patch.len()].copy_from_slice(patch);
3105        fs.object_store.put(&fs.tmp_path, &bytes).await.unwrap();
3106
3107        let file_scheduler = fs
3108            .scheduler
3109            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3110            .await
3111            .unwrap();
3112        let file_reader = FileReader::try_open(
3113            file_scheduler,
3114            None,
3115            Arc::<DecoderPlugins>::default(),
3116            &test_cache(),
3117            FileReaderOptions::default(),
3118        )
3119        .await
3120        .unwrap();
3121        file_reader
3122            .read_stream(
3123                lance_io::ReadBatchParams::RangeFull,
3124                1024,
3125                16,
3126                FilterExpression::no_filter(),
3127            )
3128            .await?
3129            .try_collect::<Vec<_>>()
3130            .await
3131    }
3132
3133    #[tokio::test]
3134    async fn test_reader_rejects_excess_miniblock_row_counts() {
3135        let batch =
3136            arrow_array::record_batch!(("id", UInt64, (0..2048_u64).collect::<Vec<_>>())).unwrap();
3137        let fs = FsFixture::default();
3138        write_lance_file(
3139            RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()),
3140            &fs,
3141            ConcreteFileVersion::V2_1,
3142            FileWriterOptions::default(),
3143        )
3144        .await;
3145
3146        let mut bytes = fs
3147            .object_store
3148            .read_one_all(&fs.tmp_path)
3149            .await
3150            .unwrap()
3151            .to_vec();
3152        // V2.1 places this column's first mini-block metadata word at byte zero.
3153        // The issue's mutation makes its non-final item count exceed the page total.
3154        bytes[0] ^= 0xf7;
3155        fs.object_store.put(&fs.tmp_path, &bytes).await.unwrap();
3156
3157        let file_scheduler = fs
3158            .scheduler
3159            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3160            .await
3161            .unwrap();
3162        let file_reader = FileReader::try_open(
3163            file_scheduler,
3164            None,
3165            Arc::<DecoderPlugins>::default(),
3166            &test_cache(),
3167            FileReaderOptions::default(),
3168        )
3169        .await
3170        .unwrap();
3171        let result = file_reader
3172            .read_stream(
3173                lance_io::ReadBatchParams::RangeFull,
3174                1024,
3175                16,
3176                FilterExpression::no_filter(),
3177            )
3178            .await;
3179        let error = match result {
3180            Ok(stream) => stream
3181                .try_collect::<Vec<_>>()
3182                .await
3183                .expect_err("excess mini-block row counts must fail the read"),
3184            Err(error) => error,
3185        };
3186        assert!(
3187            matches!(error, lance_core::Error::CorruptFile { .. }),
3188            "expected CorruptFile, got: {error}"
3189        );
3190        assert!(
3191            error.to_string().contains("exceeding items_in_page"),
3192            "unexpected message: {error}"
3193        );
3194    }
3195
3196    /// A corrupt file whose variable-width offsets point outside the value bytes
3197    /// must fail with a typed error under the default reader configuration
3198    /// (`validate_on_decode` disabled) instead of materializing values outside
3199    /// the data buffer.
3200    ///
3201    /// Uses a dictionary-encoded string column because its values page stores
3202    /// the offsets verbatim, so flipping the tail offset in the file reaches the
3203    /// Arrow conversion boundary without being rejected by an intermediate
3204    /// decompressor.
3205    #[rstest]
3206    #[tokio::test]
3207    async fn test_default_reader_rejects_out_of_bounds_variable_width_offsets(
3208        #[values(LanceFileVersion::V2_1, LanceFileVersion::V2_2, LanceFileVersion::V2_3)]
3209        version: LanceFileVersion,
3210    ) {
3211        use arrow_array::{Array, DictionaryArray, Int32Array, StringArray};
3212
3213        let values = StringArray::from(vec!["alpha", "beta", "gamma"]);
3214        let indices = Int32Array::from((0..300).map(|i| i % 3).collect::<Vec<i32>>());
3215        let dictionary = DictionaryArray::new(indices, Arc::new(values));
3216        let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
3217            "category",
3218            dictionary.data_type().clone(),
3219            false,
3220        )]));
3221        let batch = RecordBatch::try_new(arrow_schema, vec![Arc::new(dictionary)]).unwrap();
3222
3223        // The dictionary values page stores the value offsets as plain
3224        // little-endian i32s ending with [5, 9, 14] (2.1 also stores the leading
3225        // zero, 2.2+ omits it).  If a future encoding change stops storing these
3226        // offsets verbatim this lookup fails loudly and the test needs a new
3227        // byte pattern.  The patch rewrites the tail offset so it points far
3228        // beyond the value bytes.
3229        let offsets_tail_pattern = [5_i32, 9, 14]
3230            .iter()
3231            .flat_map(|value| value.to_le_bytes())
3232            .collect::<Vec<u8>>();
3233        let error = read_file_with_mutated_bytes(
3234            version,
3235            batch,
3236            &offsets_tail_pattern,
3237            8,
3238            &100_000_i32.to_le_bytes(),
3239        )
3240        .await
3241        .expect_err("out-of-bounds offsets must fail the read");
3242        assert!(
3243            matches!(error, lance_core::Error::CorruptFile { .. }),
3244            "expected CorruptFile, got: {error}"
3245        );
3246        assert!(
3247            error.to_string().contains("out of bounds"),
3248            "unexpected message: {error}"
3249        );
3250    }
3251
3252    /// Storage dictionaries expand their values through `DataBlockBuilder`
3253    /// before the final Arrow layout validation.  Corrupt dictionary offsets
3254    /// must therefore fail at the append boundary instead of reaching a slice
3255    /// operation with a decreasing range.
3256    #[rstest]
3257    #[tokio::test]
3258    async fn test_default_reader_rejects_non_monotonic_storage_dictionary_offsets(
3259        #[values(LanceFileVersion::V2_1, LanceFileVersion::V2_2, LanceFileVersion::V2_3)]
3260        version: LanceFileVersion,
3261    ) {
3262        use arrow_array::StringArray;
3263
3264        let metadata = HashMap::from([
3265            (
3266                "lance-encoding:dict-size-ratio".to_string(),
3267                "0.99".to_string(),
3268            ),
3269            (
3270                "lance-encoding:dict-values-compression".to_string(),
3271                "none".to_string(),
3272            ),
3273        ]);
3274        let arrow_schema = Arc::new(ArrowSchema::new(vec![
3275            Field::new("category", DataType::Utf8, false).with_metadata(metadata),
3276        ]));
3277        let values = (0..300)
3278            .map(|index| match index % 3 {
3279                0 => "alpha",
3280                1 => "beta",
3281                _ => "gamma",
3282            })
3283            .collect::<Vec<_>>();
3284        let batch =
3285            RecordBatch::try_new(arrow_schema, vec![Arc::new(StringArray::from(values))]).unwrap();
3286
3287        let offsets_tail_pattern = [5_i32, 9, 14]
3288            .iter()
3289            .flat_map(|value| value.to_le_bytes())
3290            .collect::<Vec<u8>>();
3291        let error = read_file_with_mutated_bytes(
3292            version,
3293            batch,
3294            &offsets_tail_pattern,
3295            4,
3296            &2_i32.to_le_bytes(),
3297        )
3298        .await
3299        .expect_err("non-monotonic dictionary offsets must fail the read");
3300        assert!(
3301            matches!(error, lance_core::Error::CorruptFile { .. }),
3302            "expected CorruptFile, got: {error}"
3303        );
3304        assert!(
3305            error.to_string().contains("decreases"),
3306            "unexpected message: {error}"
3307        );
3308    }
3309
3310    /// Same contract as the test above, but for a plain (non-dictionary) string
3311    /// column: the mini-block chunk stores chunk-relative value offsets that are
3312    /// used to slice the chunk, so a corrupt tail offset must surface as a typed
3313    /// error from the chunk decompressor instead of a panic in the decode task.
3314    #[rstest]
3315    #[tokio::test]
3316    async fn test_default_reader_rejects_out_of_bounds_miniblock_offsets(
3317        #[values(LanceFileVersion::V2_1, LanceFileVersion::V2_2, LanceFileVersion::V2_3)]
3318        version: LanceFileVersion,
3319    ) {
3320        use arrow_array::StringArray;
3321
3322        let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
3323            "strings",
3324            DataType::Utf8,
3325            false,
3326        )]));
3327        let batch = RecordBatch::try_new(
3328            arrow_schema,
3329            vec![Arc::new(StringArray::from(vec!["alpha", "beta", "gamma"]))],
3330        )
3331        .unwrap();
3332
3333        // For ["alpha", "beta", "gamma"] the chunk stores LE i32 offsets
3334        // [16, 21, 25, 30] (chunk-relative: a 16-byte offsets region precedes
3335        // the value bytes).  The patch rewrites the tail offset to point far
3336        // past the chunk.
3337        let chunk_offsets_pattern = [16_i32, 21, 25, 30]
3338            .iter()
3339            .flat_map(|value| value.to_le_bytes())
3340            .collect::<Vec<u8>>();
3341        let error = read_file_with_mutated_bytes(
3342            version,
3343            batch,
3344            &chunk_offsets_pattern,
3345            12,
3346            &100_000_i32.to_le_bytes(),
3347        )
3348        .await
3349        .expect_err("an out-of-bounds chunk offset must fail the read");
3350        assert!(
3351            matches!(error, lance_core::Error::CorruptFile { .. }),
3352            "expected CorruptFile, got: {error}"
3353        );
3354        assert!(
3355            error.to_string().contains("out of bounds"),
3356            "unexpected message: {error}"
3357        );
3358    }
3359
3360    #[tokio::test]
3361    async fn test_round_trip() {
3362        let fs = FsFixture::default();
3363
3364        let WrittenFile { data, .. } = create_some_file(&fs, ConcreteFileVersion::V2_0).await;
3365
3366        let file_size = fs.object_store.size(&fs.tmp_path).await.unwrap() as usize;
3367        let footer = fs
3368            .object_store
3369            .open(&fs.tmp_path)
3370            .await
3371            .unwrap()
3372            .get_range(file_size - 8..file_size)
3373            .await
3374            .unwrap();
3375        assert_eq!(footer_version(&footer), (0, 3));
3376        assert_eq!(
3377            crate::determine_file_version(&fs.object_store, &fs.tmp_path, Some(file_size))
3378                .await
3379                .unwrap(),
3380            ConcreteFileVersion::V2_0
3381        );
3382
3383        for read_size in [32, 1024, 1024 * 1024] {
3384            let file_scheduler = fs
3385                .scheduler
3386                .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3387                .await
3388                .unwrap();
3389            let file_reader = FileReader::try_open(
3390                file_scheduler,
3391                None,
3392                Arc::<DecoderPlugins>::default(),
3393                &test_cache(),
3394                FileReaderOptions::default(),
3395            )
3396            .await
3397            .unwrap();
3398
3399            assert_eq!(
3400                (
3401                    file_reader.metadata().major_version,
3402                    file_reader.metadata().minor_version
3403                ),
3404                (0, 3)
3405            );
3406            let schema = file_reader.schema();
3407            assert_eq!(schema.metadata.get("foo").unwrap(), "bar");
3408
3409            let batch_stream = file_reader
3410                .read_stream(
3411                    lance_io::ReadBatchParams::RangeFull,
3412                    read_size,
3413                    16,
3414                    FilterExpression::no_filter(),
3415                )
3416                .await
3417                .unwrap();
3418
3419            verify_expected(&data, batch_stream, read_size, None).await;
3420        }
3421    }
3422
3423    #[rstest]
3424    #[test_log::test(tokio::test)]
3425    async fn test_encoded_batch_round_trip(
3426        // TODO: Add V2_1 (currently fails)
3427        #[values(ConcreteFileVersion::V2_0)] version: ConcreteFileVersion,
3428    ) {
3429        let data = gen_batch()
3430            .col("x", array::rand::<Int32Type>())
3431            .col("y", array::rand_utf8(ByteCount::from(16), false))
3432            .into_batch_rows(RowCount::from(10000))
3433            .unwrap();
3434
3435        let lance_schema = Arc::new(Schema::try_from(data.schema().as_ref()).unwrap());
3436
3437        let encoding_options = EncodingOptions {
3438            cache_bytes_per_column: 4096,
3439            max_page_bytes: 32 * 1024 * 1024,
3440            keep_original_array: true,
3441            buffer_alignment: 64,
3442        };
3443
3444        let encoding_strategy = crate::versions::v2_0::encoding_strategy();
3445
3446        let encoded_batch = encode_batch(
3447            &data,
3448            lance_schema.clone(),
3449            encoding_strategy.as_ref(),
3450            &encoding_options,
3451        )
3452        .await
3453        .unwrap();
3454
3455        // Test self described
3456        let bytes = versions::encode_self_described_batch(version, &encoded_batch).unwrap();
3457        assert_eq!(footer_version(&bytes), (2, 0));
3458
3459        let decoded_batch = EncodedBatch::try_from_self_described_lance(bytes).unwrap();
3460
3461        let decoded = decode_batch(
3462            &decoded_batch,
3463            &FilterExpression::no_filter(),
3464            Arc::<DecoderPlugins>::default(),
3465            false,
3466            EncodedBatchLayout::Array,
3467            None,
3468        )
3469        .await
3470        .unwrap();
3471
3472        assert_eq!(data, decoded);
3473
3474        // Test mini
3475        let bytes = versions::encode_mini_batch(version, &encoded_batch).unwrap();
3476        assert_eq!(footer_version(&bytes), (2, 0));
3477        let decoded_batch =
3478            EncodedBatch::try_from_mini_lance(bytes, lance_schema.as_ref()).unwrap();
3479        let decoded = decode_batch(
3480            &decoded_batch,
3481            &FilterExpression::no_filter(),
3482            Arc::<DecoderPlugins>::default(),
3483            false,
3484            EncodedBatchLayout::Array,
3485            None,
3486        )
3487        .await
3488        .unwrap();
3489
3490        assert_eq!(data, decoded);
3491    }
3492
3493    #[rstest]
3494    #[test_log::test(tokio::test)]
3495    async fn test_projection(
3496        #[values(
3497            ConcreteFileVersion::V2_0,
3498            ConcreteFileVersion::V2_1,
3499            ConcreteFileVersion::V2_2
3500        )]
3501        version: ConcreteFileVersion,
3502    ) {
3503        let fs = FsFixture::default();
3504
3505        let written_file = create_some_file(&fs, version).await;
3506        let file_scheduler = fs
3507            .scheduler
3508            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3509            .await
3510            .unwrap();
3511
3512        let field_id_mapping = written_file
3513            .field_id_mapping
3514            .iter()
3515            .copied()
3516            .collect::<BTreeMap<_, _>>();
3517
3518        let empty_projection = ReaderProjection {
3519            column_indices: Vec::default(),
3520            schema: Arc::new(Schema::default()),
3521        };
3522
3523        for columns in [
3524            vec!["score"],
3525            vec!["location"],
3526            vec!["categories"],
3527            vec!["score.x"],
3528            vec!["score", "categories"],
3529            vec!["score", "location"],
3530            vec!["location", "categories"],
3531            vec!["score.y", "location", "categories"],
3532        ] {
3533            debug!("Testing round trip with projection {:?}", columns);
3534            for use_field_ids in [true, false] {
3535                // We can specify the projection as part of the read operation via read_stream_projected
3536                let file_reader = FileReader::try_open(
3537                    file_scheduler.clone(),
3538                    None,
3539                    Arc::<DecoderPlugins>::default(),
3540                    &test_cache(),
3541                    FileReaderOptions::default(),
3542                )
3543                .await
3544                .unwrap();
3545
3546                let projected_schema = written_file.schema.project(&columns).unwrap();
3547                let projection = if use_field_ids {
3548                    versions::reader_projection_from_field_ids(
3549                        file_reader.metadata().version(),
3550                        &projected_schema,
3551                        &field_id_mapping,
3552                    )
3553                    .unwrap()
3554                } else {
3555                    versions::reader_projection_from_column_names(
3556                        file_reader.metadata().version(),
3557                        &written_file.schema,
3558                        &columns,
3559                    )
3560                    .unwrap()
3561                };
3562
3563                let batch_stream = file_reader
3564                    .read_stream_projected(
3565                        lance_io::ReadBatchParams::RangeFull,
3566                        1024,
3567                        16,
3568                        projection.clone(),
3569                        FilterExpression::no_filter(),
3570                    )
3571                    .await
3572                    .unwrap();
3573
3574                let projection_arrow = ArrowSchema::from(projection.schema.as_ref());
3575                verify_expected(
3576                    &written_file.data,
3577                    batch_stream,
3578                    1024,
3579                    Some(Box::new(move |batch: &RecordBatch| {
3580                        batch.project_by_schema(&projection_arrow).unwrap()
3581                    })),
3582                )
3583                .await;
3584
3585                // We can also specify the projection as a base projection when we open the file
3586                let file_reader = FileReader::try_open(
3587                    file_scheduler.clone(),
3588                    Some(projection.clone()),
3589                    Arc::<DecoderPlugins>::default(),
3590                    &test_cache(),
3591                    FileReaderOptions::default(),
3592                )
3593                .await
3594                .unwrap();
3595
3596                let batch_stream = file_reader
3597                    .read_stream(
3598                        lance_io::ReadBatchParams::RangeFull,
3599                        1024,
3600                        16,
3601                        FilterExpression::no_filter(),
3602                    )
3603                    .await
3604                    .unwrap();
3605
3606                let projection_arrow = ArrowSchema::from(projection.schema.as_ref());
3607                verify_expected(
3608                    &written_file.data,
3609                    batch_stream,
3610                    1024,
3611                    Some(Box::new(move |batch: &RecordBatch| {
3612                        batch.project_by_schema(&projection_arrow).unwrap()
3613                    })),
3614                )
3615                .await;
3616
3617                assert!(
3618                    file_reader
3619                        .read_stream_projected(
3620                            lance_io::ReadBatchParams::RangeFull,
3621                            1024,
3622                            16,
3623                            empty_projection.clone(),
3624                            FilterExpression::no_filter(),
3625                        )
3626                        .await
3627                        .is_err()
3628                );
3629            }
3630        }
3631
3632        assert!(
3633            FileReader::try_open(
3634                file_scheduler.clone(),
3635                Some(empty_projection),
3636                Arc::<DecoderPlugins>::default(),
3637                &test_cache(),
3638                FileReaderOptions::default(),
3639            )
3640            .await
3641            .is_err()
3642        );
3643
3644        let arrow_schema = ArrowSchema::new(vec![
3645            Field::new("x", DataType::Int32, true),
3646            Field::new("y", DataType::Int32, true),
3647        ]);
3648        let schema = Schema::try_from(&arrow_schema).unwrap();
3649
3650        let projection_with_dupes = ReaderProjection {
3651            column_indices: vec![0, 0],
3652            schema: Arc::new(schema),
3653        };
3654
3655        assert!(
3656            FileReader::try_open(
3657                file_scheduler.clone(),
3658                Some(projection_with_dupes),
3659                Arc::<DecoderPlugins>::default(),
3660                &test_cache(),
3661                FileReaderOptions::default(),
3662            )
3663            .await
3664            .is_err()
3665        );
3666    }
3667
3668    #[tokio::test]
3669    async fn test_lazy_reader_direct_projection_matches_eager_reader() {
3670        let fs = FsFixture::default();
3671        let written_file = create_wide_direct_file(&fs, 16).await;
3672
3673        let file_scheduler = fs
3674            .scheduler
3675            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3676            .await
3677            .unwrap();
3678        let projection = versions::reader_projection_from_column_names(
3679            ConcreteFileVersion::V2_1,
3680            &written_file.schema,
3681            &["c10"],
3682        )
3683        .unwrap();
3684
3685        let eager_reader = FileReader::try_open(
3686            file_scheduler.clone(),
3687            None,
3688            Arc::<DecoderPlugins>::default(),
3689            &test_cache(),
3690            FileReaderOptions::default(),
3691        )
3692        .await
3693        .unwrap();
3694        let expected = eager_reader
3695            .read_stream_projected(
3696                lance_io::ReadBatchParams::RangeFull,
3697                127,
3698                16,
3699                projection.clone(),
3700                FilterExpression::no_filter(),
3701            )
3702            .await
3703            .unwrap()
3704            .try_collect::<Vec<_>>()
3705            .await
3706            .unwrap();
3707
3708        let cache = test_cache();
3709        let lazy_reader = ProjectedFileReader::try_open(
3710            file_scheduler,
3711            Some(projection.clone()),
3712            Arc::<DecoderPlugins>::default(),
3713            &cache,
3714            FileReaderOptions::default(),
3715        )
3716        .await
3717        .unwrap();
3718        let tasks = lazy_reader
3719            .read_tasks(
3720                lance_io::ReadBatchParams::RangeFull,
3721                127,
3722                None,
3723                FilterExpression::no_filter(),
3724            )
3725            .await
3726            .unwrap();
3727        let actual = collect_read_tasks(tasks, 16).await;
3728
3729        assert_eq!(expected, actual);
3730    }
3731
3732    #[tokio::test]
3733    async fn test_lazy_reader_loads_only_requested_column_metadata() {
3734        let fs = FsFixture::default();
3735        let written_file = create_wide_direct_file(&fs, 512).await;
3736
3737        let projection = versions::reader_projection_from_column_names(
3738            ConcreteFileVersion::V2_1,
3739            &written_file.schema,
3740            &["c0"],
3741        )
3742        .unwrap();
3743        let file_scheduler = fs
3744            .scheduler
3745            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3746            .await
3747            .unwrap();
3748        let lazy_reader = ProjectedFileReader::try_open(
3749            file_scheduler,
3750            Some(projection.clone()),
3751            Arc::<DecoderPlugins>::default(),
3752            &test_cache(),
3753            FileReaderOptions::default(),
3754        )
3755        .await
3756        .unwrap();
3757        let selected_column = projection.column_indices[0] as usize;
3758        let requested_metadata_bytes = lazy_reader
3759            .metadata_index()
3760            .unwrap()
3761            .column_metadata_offsets[selected_column]
3762            .1;
3763        let total_metadata_bytes = lazy_reader
3764            .metadata_index()
3765            .unwrap()
3766            .column_metadata_offsets
3767            .iter()
3768            .map(|(_, length)| *length)
3769            .sum::<u64>();
3770        assert!(
3771            total_metadata_bytes > 8 * fs.object_store.block_size() as u64,
3772            "test file metadata is too small to prove lazy loading: {total_metadata_bytes} bytes"
3773        );
3774
3775        fs.object_store.io_stats_incremental();
3776        let tasks = lazy_reader
3777            .read_tasks(
3778                lance_io::ReadBatchParams::Range(0..0),
3779                1024,
3780                Some(projection.clone()),
3781                FilterExpression::no_filter(),
3782            )
3783            .await
3784            .unwrap();
3785        let batches = collect_read_tasks(tasks, 1).await;
3786        assert!(batches.is_empty());
3787
3788        let stats = fs.object_store.io_stats_incremental();
3789        assert!(
3790            stats.read_bytes < total_metadata_bytes / 2,
3791            "lazy read fetched too much metadata: read {} bytes, requested column metadata is {} bytes, total column metadata is {} bytes",
3792            stats.read_bytes,
3793            requested_metadata_bytes,
3794            total_metadata_bytes
3795        );
3796
3797        fs.object_store.io_stats_incremental();
3798        let tasks = lazy_reader
3799            .read_tasks(
3800                lance_io::ReadBatchParams::Range(0..0),
3801                1024,
3802                Some(projection),
3803                FilterExpression::no_filter(),
3804            )
3805            .await
3806            .unwrap();
3807        let batches = collect_read_tasks(tasks, 1).await;
3808        assert!(batches.is_empty());
3809
3810        let stats = fs.object_store.io_stats_incremental();
3811        assert_eq!(
3812            stats.read_iops, 0,
3813            "cached column metadata should avoid repeat metadata I/O"
3814        );
3815        assert_eq!(
3816            stats.read_bytes, 0,
3817            "cached column metadata should avoid repeat metadata reads"
3818        );
3819    }
3820
3821    #[tokio::test]
3822    async fn test_lazy_reader_serves_column_metadata_from_tail() {
3823        let fs = FsFixture::default();
3824        // Few columns: the whole column metadata region lands inside the tail
3825        // read performed at open, so decoding a column needs no request.
3826        let written_file = create_wide_direct_file(&fs, 2).await;
3827
3828        let projection = versions::reader_projection_from_column_names(
3829            ConcreteFileVersion::V2_1,
3830            &written_file.schema,
3831            &["c1"],
3832        )
3833        .unwrap();
3834        let file_scheduler = fs
3835            .scheduler
3836            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3837            .await
3838            .unwrap();
3839        let lazy_reader = ProjectedFileReader::try_open(
3840            file_scheduler,
3841            Some(projection.clone()),
3842            Arc::<DecoderPlugins>::default(),
3843            &test_cache(),
3844            FileReaderOptions::default(),
3845        )
3846        .await
3847        .unwrap();
3848
3849        fs.object_store.io_stats_incremental();
3850        let tasks = lazy_reader
3851            .read_tasks(
3852                lance_io::ReadBatchParams::Range(0..0),
3853                1024,
3854                Some(projection),
3855                FilterExpression::no_filter(),
3856            )
3857            .await
3858            .unwrap();
3859        let batches = collect_read_tasks(tasks, 1).await;
3860        assert!(batches.is_empty());
3861
3862        let stats = fs.object_store.io_stats_incremental();
3863        assert_eq!(
3864            stats.read_iops, 0,
3865            "column metadata already captured by the tail read must not be fetched again"
3866        );
3867    }
3868
3869    #[tokio::test]
3870    async fn test_lazy_reader_fetches_only_columns_outside_the_tail() {
3871        let fs = FsFixture::default();
3872        // Many columns: the column metadata region outgrows the tail read, so the
3873        // retained window starts inside it. Columns before the window need a
3874        // request, columns inside it do not.
3875        let written_file = create_wide_direct_file(&fs, 128).await;
3876
3877        let file_scheduler = fs
3878            .scheduler
3879            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3880            .await
3881            .unwrap();
3882        let metadata_index = FileReader::read_metadata_index(&file_scheduler)
3883            .await
3884            .unwrap();
3885        let (window_start, window) = metadata_index
3886            .retained_column_metadata
3887            .as_ref()
3888            .expect("the tail must reach into the column metadata region");
3889        let offsets = &metadata_index.column_metadata_offsets;
3890        let (first_position, first_length) = offsets[0];
3891        let (last_position, last_length) = offsets[offsets.len() - 1];
3892        assert!(
3893            first_position + first_length <= *window_start,
3894            "c0 must lie before the retained window to be fetched"
3895        );
3896        assert!(
3897            last_position >= *window_start
3898                && last_position + last_length <= window_start + window.len() as u64,
3899            "the last column must lie inside the retained window"
3900        );
3901
3902        let projection = versions::reader_projection_from_column_names(
3903            ConcreteFileVersion::V2_1,
3904            &written_file.schema,
3905            &["c0", "c127"],
3906        )
3907        .unwrap();
3908        let eager_reader = FileReader::try_open(
3909            file_scheduler.clone(),
3910            None,
3911            Arc::<DecoderPlugins>::default(),
3912            &test_cache(),
3913            FileReaderOptions::default(),
3914        )
3915        .await
3916        .unwrap();
3917        let expected = eager_reader
3918            .read_stream_projected(
3919                lance_io::ReadBatchParams::RangeFull,
3920                127,
3921                16,
3922                projection.clone(),
3923                FilterExpression::no_filter(),
3924            )
3925            .await
3926            .unwrap()
3927            .try_collect::<Vec<_>>()
3928            .await
3929            .unwrap();
3930
3931        let lazy_reader = ProjectedFileReader::try_open(
3932            file_scheduler,
3933            Some(projection.clone()),
3934            Arc::<DecoderPlugins>::default(),
3935            &test_cache(),
3936            FileReaderOptions::default(),
3937        )
3938        .await
3939        .unwrap();
3940        fs.object_store.io_stats_incremental();
3941        let tasks = lazy_reader
3942            .read_tasks(
3943                lance_io::ReadBatchParams::Range(0..0),
3944                127,
3945                Some(projection),
3946                FilterExpression::no_filter(),
3947            )
3948            .await
3949            .unwrap();
3950        assert!(collect_read_tasks(tasks, 1).await.is_empty());
3951        let stats = fs.object_store.io_stats_incremental();
3952        assert_eq!(
3953            stats.read_iops, 1,
3954            "only the column outside the tail window needs a request"
3955        );
3956        assert_eq!(
3957            stats.read_bytes, first_length,
3958            "the request must cover exactly the uncovered column's metadata"
3959        );
3960
3961        let tasks = lazy_reader
3962            .read_tasks(
3963                lance_io::ReadBatchParams::RangeFull,
3964                127,
3965                None,
3966                FilterExpression::no_filter(),
3967            )
3968            .await
3969            .unwrap();
3970        assert_eq!(expected, collect_read_tasks(tasks, 16).await);
3971    }
3972
3973    async fn assert_lazy_projection_matches_eager_and_reads_metadata_subset(
3974        fs: &FsFixture,
3975        projection: ReaderProjection,
3976        shape: &str,
3977    ) -> Vec<RecordBatch> {
3978        let file_scheduler = fs
3979            .scheduler
3980            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3981            .await
3982            .unwrap();
3983        let eager_reader = FileReader::try_open(
3984            file_scheduler.clone(),
3985            None,
3986            Arc::<DecoderPlugins>::default(),
3987            &test_cache(),
3988            FileReaderOptions::default(),
3989        )
3990        .await
3991        .unwrap();
3992        let expected = eager_reader
3993            .read_stream_projected(
3994                lance_io::ReadBatchParams::RangeFull,
3995                127,
3996                16,
3997                projection.clone(),
3998                FilterExpression::no_filter(),
3999            )
4000            .await
4001            .unwrap()
4002            .try_collect::<Vec<_>>()
4003            .await
4004            .unwrap();
4005
4006        let cache = test_cache();
4007        let lazy_reader = ProjectedFileReader::try_open(
4008            file_scheduler,
4009            Some(projection.clone()),
4010            Arc::<DecoderPlugins>::default(),
4011            &cache,
4012            FileReaderOptions::default(),
4013        )
4014        .await
4015        .unwrap();
4016        let metadata_index = lazy_reader.metadata_index().unwrap();
4017        let requested_metadata_bytes = projection
4018            .column_indices
4019            .iter()
4020            .map(|column_index| metadata_index.column_metadata_offsets[*column_index as usize].1)
4021            .sum::<u64>();
4022        let total_metadata_bytes = metadata_index
4023            .column_metadata_offsets
4024            .iter()
4025            .map(|(_, length)| *length)
4026            .sum::<u64>();
4027        assert!(total_metadata_bytes > requested_metadata_bytes * 8);
4028
4029        fs.object_store.io_stats_incremental();
4030        let tasks = lazy_reader
4031            .read_tasks(
4032                lance_io::ReadBatchParams::Range(0..0),
4033                127,
4034                None,
4035                FilterExpression::no_filter(),
4036            )
4037            .await
4038            .unwrap();
4039        assert!(collect_read_tasks(tasks, 1).await.is_empty());
4040        let metadata_stats = fs.object_store.io_stats_incremental();
4041        assert!(
4042            metadata_stats.read_bytes < total_metadata_bytes / 2,
4043            "lazy {shape} read fetched too much metadata: read {} bytes, requested column metadata is {} bytes, total column metadata is {} bytes",
4044            metadata_stats.read_bytes,
4045            requested_metadata_bytes,
4046            total_metadata_bytes
4047        );
4048
4049        let tasks = lazy_reader
4050            .read_tasks(
4051                lance_io::ReadBatchParams::RangeFull,
4052                127,
4053                None,
4054                FilterExpression::no_filter(),
4055            )
4056            .await
4057            .unwrap();
4058        let actual = collect_read_tasks(tasks, 16).await;
4059        assert_eq!(expected, actual);
4060        actual
4061    }
4062
4063    #[tokio::test]
4064    async fn test_lazy_reader_fixed_size_list_projection_matches_eager_reader() {
4065        let fs = FsFixture::default();
4066        let written_file = create_wide_fixed_size_list_file(&fs, 512).await;
4067        let projection = versions::reader_projection_from_column_names(
4068            ConcreteFileVersion::V2_1,
4069            &written_file.schema,
4070            &["c17", "c509"],
4071        )
4072        .unwrap();
4073        assert!(projection.prefers_indexed_metadata(512));
4074        assert_lazy_projection_matches_eager_and_reads_metadata_subset(
4075            &fs,
4076            projection,
4077            "fixed-size-list",
4078        )
4079        .await;
4080    }
4081
4082    #[tokio::test]
4083    async fn test_v2_0_rejects_indexed_metadata_reader() {
4084        let fs = FsFixture::default();
4085        let written_file = create_some_file(&fs, ConcreteFileVersion::V2_0).await;
4086        let projection = versions::reader_projection_from_column_names(
4087            ConcreteFileVersion::V2_0,
4088            &written_file.schema,
4089            &["score"],
4090        )
4091        .unwrap();
4092        assert!(projection.prefers_indexed_metadata(100));
4093
4094        let file_scheduler = fs
4095            .scheduler
4096            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4097            .await
4098            .unwrap();
4099        let err = ProjectedFileReader::try_open(
4100            file_scheduler,
4101            Some(projection),
4102            Arc::<DecoderPlugins>::default(),
4103            &test_cache(),
4104            FileReaderOptions::default(),
4105        )
4106        .await
4107        .unwrap_err();
4108        assert!(
4109            matches!(err, lance_core::Error::NotSupported { .. }),
4110            "expected V2.0 indexed metadata open to fail, got {err:?}"
4111        );
4112    }
4113
4114    #[tokio::test]
4115    async fn test_lazy_reader_nested_projection_compacts_physical_columns() {
4116        let fs = FsFixture::default();
4117        let written_file = create_wide_structural_file(&fs, 128).await;
4118        let projection = versions::reader_projection_from_column_names(
4119            ConcreteFileVersion::V2_1,
4120            &written_file.schema,
4121            &["s97.y", "l4", "s3"],
4122        )
4123        .unwrap();
4124
4125        assert_eq!(
4126            projection
4127                .schema
4128                .fields
4129                .iter()
4130                .map(|field| field.name.as_str())
4131                .collect::<Vec<_>>(),
4132            vec!["s97", "l4", "s3"]
4133        );
4134        assert_eq!(projection.schema.fields[0].children.len(), 1);
4135        assert_eq!(projection.schema.fields[0].children[0].name, "y");
4136        assert_eq!(projection.schema.fields[2].children.len(), 2);
4137        assert_eq!(projection.column_indices.len(), 4);
4138        assert!(
4139            projection
4140                .column_indices
4141                .windows(2)
4142                .any(|indices| indices[0] > indices[1]),
4143            "the projection must reorder physical columns to exercise compact remapping"
4144        );
4145        assert!(projection.prefers_indexed_metadata(128 * 4));
4146        let actual = assert_lazy_projection_matches_eager_and_reads_metadata_subset(
4147            &fs, projection, "nested",
4148        )
4149        .await;
4150        assert!(
4151            actual
4152                .iter()
4153                .flat_map(|batch| batch.columns())
4154                .any(|column| column.null_count() > 0),
4155            "the structural projection must exercise nullable arrays"
4156        );
4157    }
4158
4159    #[rstest]
4160    #[case::before_metadata_region(90, 5)]
4161    #[case::after_metadata_region(190, 20)]
4162    fn test_decode_cmo_table_rejects_out_of_range_offsets(
4163        #[case] position: u64,
4164        #[case] length: u64,
4165    ) {
4166        let mut cmo_table = [0; 16];
4167        cmo_table[0..8].copy_from_slice(&position.to_le_bytes());
4168        cmo_table[8..16].copy_from_slice(&length.to_le_bytes());
4169        let footer = super::Footer {
4170            column_meta_start: 100,
4171            column_meta_offsets_start: 200,
4172            global_buff_offsets_start: 200,
4173            num_global_buffers: 0,
4174            num_columns: 1,
4175            major_version: 2,
4176            minor_version: 1,
4177        };
4178
4179        let err = FileReader::decode_cmo_table(Bytes::copy_from_slice(&cmo_table), &footer)
4180            .expect_err("out-of-range CMO entries must be rejected");
4181        assert!(
4182            matches!(err, lance_core::Error::InvalidInput { .. }),
4183            "expected InvalidInput, got {err:?}"
4184        );
4185    }
4186
4187    #[rstest]
4188    #[case::blob(BLOB_META_KEY)]
4189    #[case::packed_struct("lance-encoding:packed")]
4190    #[tokio::test]
4191    async fn test_lazy_reader_rejects_opaque_projection(#[case] metadata_key: &str) {
4192        let fs = FsFixture::default();
4193        let written_file = create_some_file(&fs, ConcreteFileVersion::V2_1).await;
4194
4195        let ordinary_projection = versions::reader_projection_from_column_names(
4196            ConcreteFileVersion::V2_1,
4197            &written_file.schema,
4198            &["location.x"],
4199        )
4200        .unwrap();
4201        assert_eq!(ordinary_projection.schema.fields[0].children.len(), 1);
4202        assert!(ordinary_projection.prefers_indexed_metadata(100));
4203
4204        let file_scheduler = fs
4205            .scheduler
4206            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4207            .await
4208            .unwrap();
4209        let err = ProjectedFileReader::try_open(
4210            file_scheduler.clone(),
4211            None,
4212            Arc::<DecoderPlugins>::default(),
4213            &test_cache(),
4214            FileReaderOptions::default(),
4215        )
4216        .await
4217        .unwrap_err();
4218        assert!(
4219            matches!(err, lance_core::Error::InvalidInput { .. }),
4220            "expected InvalidInput, got {err:?}"
4221        );
4222
4223        let mut projection = ordinary_projection;
4224        Arc::make_mut(&mut projection.schema).fields[0]
4225            .metadata
4226            .insert(metadata_key.to_string(), "true".to_string());
4227        assert!(!projection.prefers_indexed_metadata(100));
4228
4229        let err = ProjectedFileReader::try_open(
4230            file_scheduler,
4231            Some(projection),
4232            Arc::<DecoderPlugins>::default(),
4233            &test_cache(),
4234            FileReaderOptions::default(),
4235        )
4236        .await
4237        .unwrap_err();
4238        assert!(
4239            matches!(err, lance_core::Error::NotSupported { .. }),
4240            "expected NotSupported for {metadata_key}, got {err:?}"
4241        );
4242    }
4243
4244    // The projection-length validation lives in `DecodeEngine`, shared by the
4245    // eager and the lazy (indexed) metadata providers. The indexed provider loads
4246    // only the projected columns and renumbers them 0..N, so this checks that the
4247    // renumbered `column_infos`/`column_indices` still line up for the length
4248    // check: a mismatched-length projection is rejected through
4249    // `ProjectedFileReader`, and a single short column resolves to its own length.
4250    #[tokio::test]
4251    async fn test_lazy_reader_validates_unequal_length_projection() {
4252        use arrow_array::Int32Array;
4253        use lance_io::ReadBatchParams;
4254
4255        let arrow_schema = Arc::new(ArrowSchema::new(vec![
4256            Field::new("a", DataType::Int32, true),
4257            Field::new("c", DataType::Int32, true),
4258        ]));
4259        let lance_schema = Schema::try_from(arrow_schema.as_ref()).unwrap();
4260
4261        let fs = FsFixture::default();
4262        let mut writer = versions::v2_1::create_writer(
4263            fs.object_store.create(&fs.tmp_path).await.unwrap(),
4264            lance_schema.clone(),
4265            FileWriterOptions::default(),
4266        )
4267        .unwrap();
4268        // "a" has 5 rows, "c" has 1 -- an unequal-length file.
4269        writer
4270            .write_column(0, Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])))
4271            .await
4272            .unwrap();
4273        writer
4274            .write_column(1, Arc::new(Int32Array::from(vec![100])))
4275            .await
4276            .unwrap();
4277        writer.finish().await.unwrap();
4278
4279        let file_scheduler = fs
4280            .scheduler
4281            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4282            .await
4283            .unwrap();
4284        let cache = test_cache();
4285        let open_indexed = |names: &[&str]| {
4286            let projection = versions::reader_projection_from_column_names(
4287                ConcreteFileVersion::V2_1,
4288                &lance_schema,
4289                names,
4290            )
4291            .unwrap();
4292            ProjectedFileReader::try_open(
4293                file_scheduler.clone(),
4294                Some(projection),
4295                Arc::<DecoderPlugins>::default(),
4296                &cache,
4297                FileReaderOptions::default(),
4298            )
4299        };
4300
4301        // A mismatched-length projection [a, c] (5 vs 1) is rejected at read time,
4302        // through the indexed provider's renumbered column infos.
4303        let lazy = open_indexed(&["a", "c"]).await.unwrap();
4304        // (`read_tasks` yields a stream, which is not `Debug`, so match rather
4305        // than `unwrap_err`.)
4306        let err = match lazy
4307            .read_tasks(
4308                ReadBatchParams::RangeFull,
4309                1024,
4310                None,
4311                FilterExpression::no_filter(),
4312            )
4313            .await
4314        {
4315            Ok(_) => panic!("expected the mismatched-length projection to be rejected"),
4316            Err(e) => e.to_string(),
4317        };
4318        assert!(
4319            err.contains("a=5") && err.contains("c=1"),
4320            "error should name each column's length, got: {err}"
4321        );
4322
4323        // A single short column resolves to its own length (1), not the file's
4324        // longest column.
4325        let lazy = open_indexed(&["c"]).await.unwrap();
4326        let tasks = lazy
4327            .read_tasks(
4328                ReadBatchParams::RangeFull,
4329                1024,
4330                None,
4331                FilterExpression::no_filter(),
4332            )
4333            .await
4334            .unwrap();
4335        let batches = collect_read_tasks(tasks, 16).await;
4336        let values: Vec<Option<i32>> = batches
4337            .iter()
4338            .flat_map(|b| {
4339                b.column(0)
4340                    .as_any()
4341                    .downcast_ref::<Int32Array>()
4342                    .unwrap()
4343                    .iter()
4344                    .collect::<Vec<_>>()
4345            })
4346            .collect();
4347        assert_eq!(values, vec![Some(100)]);
4348    }
4349
4350    #[test_log::test(tokio::test)]
4351    async fn test_compressing_buffer() {
4352        let fs = FsFixture::default();
4353
4354        let written_file = create_some_file(&fs, ConcreteFileVersion::V2_0).await;
4355        let file_scheduler = fs
4356            .scheduler
4357            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4358            .await
4359            .unwrap();
4360
4361        // We can specify the projection as part of the read operation via read_stream_projected
4362        let file_reader = FileReader::try_open(
4363            file_scheduler.clone(),
4364            None,
4365            Arc::<DecoderPlugins>::default(),
4366            &test_cache(),
4367            FileReaderOptions::default(),
4368        )
4369        .await
4370        .unwrap();
4371
4372        let mut projection = written_file.schema.project(&["score"]).unwrap();
4373        for field in projection.fields.iter_mut() {
4374            field
4375                .metadata
4376                .insert("lance:compression".to_string(), "zstd".to_string());
4377        }
4378        let projection = ReaderProjection {
4379            column_indices: projection.fields.iter().map(|f| f.id as u32).collect(),
4380            schema: Arc::new(projection),
4381        };
4382
4383        let batch_stream = file_reader
4384            .read_stream_projected(
4385                lance_io::ReadBatchParams::RangeFull,
4386                1024,
4387                16,
4388                projection.clone(),
4389                FilterExpression::no_filter(),
4390            )
4391            .await
4392            .unwrap();
4393
4394        let projection_arrow = Arc::new(ArrowSchema::from(projection.schema.as_ref()));
4395        verify_expected(
4396            &written_file.data,
4397            batch_stream,
4398            1024,
4399            Some(Box::new(move |batch: &RecordBatch| {
4400                batch.project_by_schema(&projection_arrow).unwrap()
4401            })),
4402        )
4403        .await;
4404    }
4405
4406    #[tokio::test]
4407    async fn test_read_all() {
4408        let fs = FsFixture::default();
4409        let WrittenFile { data, .. } = create_some_file(&fs, ConcreteFileVersion::V2_0).await;
4410        let total_rows = data.iter().map(|batch| batch.num_rows()).sum::<usize>();
4411
4412        let file_scheduler = fs
4413            .scheduler
4414            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4415            .await
4416            .unwrap();
4417        let file_reader = FileReader::try_open(
4418            file_scheduler.clone(),
4419            None,
4420            Arc::<DecoderPlugins>::default(),
4421            &test_cache(),
4422            FileReaderOptions::default(),
4423        )
4424        .await
4425        .unwrap();
4426
4427        let batches = file_reader
4428            .read_stream(
4429                lance_io::ReadBatchParams::RangeFull,
4430                total_rows as u32,
4431                16,
4432                FilterExpression::no_filter(),
4433            )
4434            .await
4435            .unwrap()
4436            .try_collect::<Vec<_>>()
4437            .await
4438            .unwrap();
4439        assert_eq!(batches.len(), 1);
4440        assert_eq!(batches[0].num_rows(), total_rows);
4441    }
4442
4443    /// The writer cuts a column into many small pages when its buffer is
4444    /// small (or, for nested columns, when the rep/def levels of a page hit
4445    /// the mini-block budget). Reading them must not cost one request per page.
4446    #[tokio::test]
4447    async fn test_read_batches_page_reads_into_one_request() {
4448        let fs = FsFixture::default();
4449        // Every batch is 40 KiB of values, so the 64 KiB buffer flushes a page
4450        // after every second batch.
4451        let reader = gen_batch()
4452            .col("x", array::step::<Int32Type>())
4453            .into_reader_rows(RowCount::from(10_000), BatchCount::from(10));
4454        write_lance_file(
4455            reader,
4456            &fs,
4457            ConcreteFileVersion::V2_1,
4458            FileWriterOptions {
4459                data_cache_bytes: Some(64 * 1024),
4460                ..Default::default()
4461            },
4462        )
4463        .await;
4464
4465        let file_scheduler = fs
4466            .scheduler
4467            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4468            .await
4469            .unwrap();
4470        let file_reader = FileReader::try_open(
4471            file_scheduler,
4472            None,
4473            Arc::<DecoderPlugins>::default(),
4474            &test_cache(),
4475            FileReaderOptions::default(),
4476        )
4477        .await
4478        .unwrap();
4479        let num_pages = file_reader.metadata().column_infos[0].page_infos.len();
4480        assert!(
4481            num_pages > 4,
4482            "the writer must cut several pages, got {num_pages}"
4483        );
4484
4485        let read_all = || async {
4486            file_reader
4487                .read_stream(
4488                    lance_io::ReadBatchParams::RangeFull,
4489                    100_000,
4490                    16,
4491                    FilterExpression::no_filter(),
4492                )
4493                .await
4494                .unwrap()
4495                .try_collect::<Vec<_>>()
4496                .await
4497                .unwrap()
4498        };
4499        // The first read also initializes the page metadata; the cache keeps
4500        // it, so the second read is data I/O only.
4501        read_all().await;
4502        fs.object_store.io_stats_incremental();
4503        let batches = read_all().await;
4504        assert_eq!(
4505            batches.iter().map(|batch| batch.num_rows()).sum::<usize>(),
4506            100_000
4507        );
4508        let stats = fs.object_store.io_stats_incremental();
4509        assert_eq!(
4510            stats.read_iops, 1,
4511            "{num_pages} adjacent pages must be read together, not one request each"
4512        );
4513    }
4514
4515    /// Read `ranges` of every column through a clone of `reader` that merges
4516    /// requested byte ranges at most `gap` bytes apart (the store's block size
4517    /// when `None`), returning the batch and the requests the read made.
4518    async fn read_ranges_with_gap(
4519        reader: &FileReader,
4520        ranges: &[std::ops::Range<u64>],
4521        gap: Option<u64>,
4522    ) -> (RecordBatch, u64) {
4523        let stats = IoStats::new();
4524        let mut reader = reader.with_io_stats(stats.recorder());
4525        if let Some(gap) = gap {
4526            reader = reader.with_coalesce_gap(gap);
4527        }
4528        let batches = reader
4529            .read_stream(
4530                lance_io::ReadBatchParams::Ranges(ranges.into()),
4531                u32::MAX,
4532                1,
4533                FilterExpression::no_filter(),
4534            )
4535            .await
4536            .unwrap()
4537            .try_collect::<Vec<_>>()
4538            .await
4539            .unwrap();
4540        let batch = concat_batches(&batches[0].schema(), &batches).unwrap();
4541        (batch, stats.snapshot().iops)
4542    }
4543
4544    /// A reader's coalescing gap changes only the requests that a read of
4545    /// scattered rows makes: every gap returns the same batch, a wider gap
4546    /// makes no more requests than a narrower one, and a gap that covers the
4547    /// file makes at most one request per column page.
4548    #[rstest]
4549    #[case::v2_0(ConcreteFileVersion::V2_0)]
4550    #[case::v2_2(ConcreteFileVersion::V2_2)]
4551    #[tokio::test]
4552    async fn test_coalesce_gap_keeps_ranges_reads(#[case] version: ConcreteFileVersion) {
4553        const PAGES: u32 = 4;
4554        const PAGE_ROWS: u64 = 1024;
4555        let fs = FsFixture::default();
4556        let codes =
4557            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::UInt8, true)), 64);
4558        let reader = gen_batch()
4559            .col("x", array::step::<Int32Type>())
4560            .col("codes", array::rand_type(&codes))
4561            .into_reader_rows(RowCount::from(PAGE_ROWS), BatchCount::from(PAGES));
4562        // Without a cache, every written batch is a page of each column.
4563        let WrittenFile { data, .. } = write_lance_file(
4564            reader,
4565            &fs,
4566            version,
4567            FileWriterOptions {
4568                data_cache_bytes: Some(1),
4569                ..Default::default()
4570            },
4571        )
4572        .await;
4573        let file_scheduler = fs
4574            .scheduler
4575            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4576            .await
4577            .unwrap();
4578        let file_reader = FileReader::try_open(
4579            file_scheduler,
4580            None,
4581            Arc::<DecoderPlugins>::default(),
4582            &test_cache(),
4583            FileReaderOptions::default(),
4584        )
4585        .await
4586        .unwrap();
4587        let num_pages = file_reader
4588            .metadata()
4589            .column_infos
4590            .iter()
4591            .map(|column| column.page_infos.len() as u64)
4592            .sum::<u64>();
4593        assert!(
4594            num_pages >= 2 * u64::from(PAGES),
4595            "the writer must cut a page of each column per batch, got {num_pages}"
4596        );
4597
4598        // Single rows 97 apart, a run across the end of the first page and
4599        // a run inside a page.
4600        let rows = (0..PAGE_ROWS * u64::from(PAGES))
4601            .filter(|row| {
4602                row.is_multiple_of(97) || (1020..1030).contains(row) || (3000..3003).contains(row)
4603            })
4604            .collect::<Vec<_>>();
4605        let mut ranges: Vec<std::ops::Range<u64>> = Vec::new();
4606        for &row in &rows {
4607            match ranges.last_mut() {
4608                Some(last) if last.end == row => last.end += 1,
4609                _ => ranges.push(row..row + 1),
4610            }
4611        }
4612        let indices = UInt32Array::from_iter_values(rows.iter().map(|&row| row as u32));
4613        let expected = concat_batches(&data[0].schema(), &data)
4614            .unwrap()
4615            .take(&indices)
4616            .unwrap();
4617
4618        // The first read also initializes the page metadata; the cache keeps
4619        // it, so the reads below are data I/O only.
4620        let (baseline, _) = read_ranges_with_gap(&file_reader, &ranges, None).await;
4621        assert_eq!(baseline.columns(), expected.columns());
4622        let mut requests = Vec::new();
4623        for gap in [0, 64 * 1024, 256 * 1024, 1024 * 1024, u64::MAX] {
4624            let (batch, gap_requests) =
4625                read_ranges_with_gap(&file_reader, &ranges, Some(gap)).await;
4626            assert_eq!(batch, baseline, "gap={gap}");
4627            requests.push(gap_requests);
4628        }
4629        assert!(
4630            requests.windows(2).all(|pair| pair[1] <= pair[0]),
4631            "{requests:?}"
4632        );
4633        let covering = requests[requests.len() - 1];
4634        assert!(
4635            covering <= num_pages,
4636            "{covering} requests for {num_pages} column pages"
4637        );
4638        if version == ConcreteFileVersion::V2_0 {
4639            // Each run of each column is its own request at gap 0.
4640            assert!(requests[0] > covering, "{requests:?}");
4641        }
4642    }
4643
4644    #[rstest]
4645    #[tokio::test]
4646    async fn test_take_batches_column_reads_into_one_request(
4647        #[values(ConcreteFileVersion::V2_0, ConcreteFileVersion::V2_1)]
4648        version: ConcreteFileVersion,
4649    ) {
4650        let fs = FsFixture::default();
4651        // 16 int32 columns of 256 rows: one 1 KiB page each, written back to
4652        // back, so one row's reads sit within the local 4 KiB block size.
4653        let mut generator = gen_batch();
4654        for i in 0..16 {
4655            generator = generator.col(format!("c{i}"), array::step::<Int32Type>());
4656        }
4657        write_lance_file(
4658            generator.into_reader_rows(RowCount::from(256), BatchCount::from(1)),
4659            &fs,
4660            version,
4661            FileWriterOptions::default(),
4662        )
4663        .await;
4664
4665        let file_scheduler = fs
4666            .scheduler
4667            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4668            .await
4669            .unwrap();
4670        let file_reader = FileReader::try_open(
4671            file_scheduler,
4672            None,
4673            Arc::<DecoderPlugins>::default(),
4674            &test_cache(),
4675            FileReaderOptions::default(),
4676        )
4677        .await
4678        .unwrap();
4679
4680        let take = |row: u32| {
4681            let file_reader = &file_reader;
4682            async move {
4683                file_reader
4684                    .read_stream(
4685                        lance_io::ReadBatchParams::Indices(UInt32Array::from(vec![row])),
4686                        1,
4687                        1,
4688                        FilterExpression::no_filter(),
4689                    )
4690                    .await
4691                    .unwrap()
4692                    .try_collect::<Vec<_>>()
4693                    .await
4694                    .unwrap()
4695            }
4696        };
4697        // The first take also initializes the page metadata; the cache keeps
4698        // it, so the second take is data I/O only.
4699        take(3).await;
4700        fs.object_store.io_stats_incremental();
4701        let batches = take(200).await;
4702        let stats = fs.object_store.io_stats_incremental();
4703        assert_eq!(
4704            stats.read_iops, 1,
4705            "one row across 16 adjacent column pages must be one request, not one per column"
4706        );
4707        assert_eq!(batches[0].num_columns(), 16);
4708        for column in batches[0].columns() {
4709            assert_eq!(column.as_primitive::<Int32Type>().value(0), 200);
4710        }
4711    }
4712
4713    #[rstest]
4714    #[tokio::test]
4715    async fn test_blocking_take(
4716        #[values(
4717            ConcreteFileVersion::V2_0,
4718            ConcreteFileVersion::V2_1,
4719            ConcreteFileVersion::V2_2
4720        )]
4721        version: ConcreteFileVersion,
4722    ) {
4723        let fs = FsFixture::default();
4724        let WrittenFile { data, schema, .. } = create_some_file(&fs, version).await;
4725        let total_rows = data.iter().map(|batch| batch.num_rows()).sum::<usize>();
4726
4727        let file_scheduler = fs
4728            .scheduler
4729            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4730            .await
4731            .unwrap();
4732        let file_reader = FileReader::try_open(
4733            file_scheduler.clone(),
4734            Some(
4735                versions::reader_projection_from_column_names(version, &schema, &["score"])
4736                    .unwrap(),
4737            ),
4738            Arc::<DecoderPlugins>::default(),
4739            &test_cache(),
4740            FileReaderOptions::default(),
4741        )
4742        .await
4743        .unwrap();
4744
4745        let batches = tokio::task::spawn_blocking(move || {
4746            file_reader
4747                .read_stream_projected_blocking(
4748                    lance_io::ReadBatchParams::Indices(UInt32Array::from(vec![0, 1, 2, 3, 4])),
4749                    total_rows as u32,
4750                    None,
4751                    FilterExpression::no_filter(),
4752                )
4753                .unwrap()
4754                .collect::<ArrowResult<Vec<_>>>()
4755                .unwrap()
4756        })
4757        .await
4758        .unwrap();
4759
4760        assert_eq!(batches.len(), 1);
4761        assert_eq!(batches[0].num_rows(), 5);
4762        assert_eq!(batches[0].num_columns(), 1);
4763    }
4764
4765    #[tokio::test(flavor = "multi_thread")]
4766    async fn test_drop_in_progress() {
4767        let fs = FsFixture::default();
4768        let WrittenFile { data, .. } = create_some_file(&fs, ConcreteFileVersion::V2_0).await;
4769        let total_rows = data.iter().map(|batch| batch.num_rows()).sum::<usize>();
4770
4771        let file_scheduler = fs
4772            .scheduler
4773            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4774            .await
4775            .unwrap();
4776        let file_reader = FileReader::try_open(
4777            file_scheduler.clone(),
4778            None,
4779            Arc::<DecoderPlugins>::default(),
4780            &test_cache(),
4781            FileReaderOptions::default(),
4782        )
4783        .await
4784        .unwrap();
4785
4786        let mut batches = file_reader
4787            .read_stream(
4788                lance_io::ReadBatchParams::RangeFull,
4789                (total_rows / 10) as u32,
4790                16,
4791                FilterExpression::no_filter(),
4792            )
4793            .await
4794            .unwrap();
4795
4796        drop(file_reader);
4797
4798        let batch = batches.next().await.unwrap().unwrap();
4799        assert!(batch.num_rows() > 0);
4800
4801        // Drop in-progress scan
4802        drop(batches);
4803    }
4804
4805    #[tokio::test]
4806    async fn drop_while_scheduling() {
4807        // This is a bit of a white-box test, pokes at the internals.  We want to
4808        // test the case where the read stream is dropped before the scheduling
4809        // thread finishes.  We can't do that in a black-box fashion because the
4810        // scheduling thread runs in the background and there is no easy way to
4811        // pause / gate it.
4812
4813        // It's a regression for a bug where the scheduling thread would panic
4814        // if the stream was dropped before it finished.
4815
4816        let fs = FsFixture::default();
4817        let written_file = create_some_file(&fs, ConcreteFileVersion::V2_0).await;
4818        let total_rows = written_file
4819            .data
4820            .iter()
4821            .map(|batch| batch.num_rows())
4822            .sum::<usize>();
4823
4824        let file_scheduler = fs
4825            .scheduler
4826            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4827            .await
4828            .unwrap();
4829        let file_reader = FileReader::try_open(
4830            file_scheduler.clone(),
4831            None,
4832            Arc::<DecoderPlugins>::default(),
4833            &test_cache(),
4834            FileReaderOptions::default(),
4835        )
4836        .await
4837        .unwrap();
4838
4839        let projection = versions::reader_projection_from_whole_schema(
4840            &written_file.schema,
4841            ConcreteFileVersion::V2_0,
4842        );
4843        let column_infos = file_reader.metadata().column_infos.clone();
4844        let mut decode_scheduler = DecodeBatchScheduler::try_new(
4845            &projection.schema,
4846            &projection.column_indices,
4847            &column_infos,
4848            &vec![],
4849            total_rows as u64,
4850            Arc::<DecoderPlugins>::default(),
4851            file_reader.scheduler(),
4852            test_cache(),
4853            &FilterExpression::no_filter(),
4854            &DecoderConfig::default(),
4855        )
4856        .await
4857        .unwrap();
4858
4859        let range = 0..total_rows as u64;
4860
4861        let (tx, rx) = mpsc::unbounded_channel();
4862
4863        // Simulate the stream / decoder being dropped
4864        drop(rx);
4865
4866        // Scheduling should not panic
4867        decode_scheduler.schedule_range(
4868            range,
4869            &FilterExpression::no_filter(),
4870            tx,
4871            file_reader.scheduler(),
4872        )
4873    }
4874
4875    #[tokio::test]
4876    async fn test_read_empty_range() {
4877        let fs = FsFixture::default();
4878        create_some_file(&fs, ConcreteFileVersion::V2_0).await;
4879
4880        let file_scheduler = fs
4881            .scheduler
4882            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4883            .await
4884            .unwrap();
4885        let file_reader = FileReader::try_open(
4886            file_scheduler.clone(),
4887            None,
4888            Arc::<DecoderPlugins>::default(),
4889            &test_cache(),
4890            FileReaderOptions::default(),
4891        )
4892        .await
4893        .unwrap();
4894
4895        // All ranges empty, no data
4896        let batches = file_reader
4897            .read_stream(
4898                lance_io::ReadBatchParams::Range(0..0),
4899                1024,
4900                16,
4901                FilterExpression::no_filter(),
4902            )
4903            .await
4904            .unwrap()
4905            .try_collect::<Vec<_>>()
4906            .await
4907            .unwrap();
4908
4909        assert_eq!(batches.len(), 0);
4910
4911        // Some ranges empty
4912        let batches = file_reader
4913            .read_stream(
4914                lance_io::ReadBatchParams::Ranges(Arc::new([0..1, 2..2])),
4915                1024,
4916                16,
4917                FilterExpression::no_filter(),
4918            )
4919            .await
4920            .unwrap()
4921            .try_collect::<Vec<_>>()
4922            .await
4923            .unwrap();
4924        assert_eq!(batches.len(), 1);
4925    }
4926
4927    async fn write_file_with_global_buffer(fs: &FsFixture, buffer: Bytes) {
4928        let lance_schema =
4929            lance_core::datatypes::Schema::try_from(&ArrowSchema::new(vec![Field::new(
4930                "foo",
4931                DataType::Int32,
4932                true,
4933            )]))
4934            .unwrap();
4935
4936        let mut file_writer = versions::v2_1::create_writer(
4937            fs.object_store.create(&fs.tmp_path).await.unwrap(),
4938            lance_schema,
4939            FileWriterOptions::default(),
4940        )
4941        .unwrap();
4942
4943        let buf_index = file_writer.add_global_buffer(buffer).await.unwrap();
4944        assert_eq!(buf_index, 1);
4945
4946        file_writer.finish().await.unwrap();
4947    }
4948
4949    #[derive(Clone, Copy, Debug)]
4950    enum MetadataReadPath {
4951        Full,
4952        Indexed,
4953    }
4954
4955    #[derive(Clone, Copy, Debug)]
4956    enum InvalidGboDescriptor {
4957        Unaligned,
4958        PastEof,
4959        Overflowing,
4960    }
4961
4962    #[rstest]
4963    #[case::full_unaligned(MetadataReadPath::Full, InvalidGboDescriptor::Unaligned, "not aligned")]
4964    #[case::full_past_eof(MetadataReadPath::Full, InvalidGboDescriptor::PastEof, "outside file")]
4965    #[case::full_overflowing(
4966        MetadataReadPath::Full,
4967        InvalidGboDescriptor::Overflowing,
4968        "overflows"
4969    )]
4970    #[case::indexed_unaligned(
4971        MetadataReadPath::Indexed,
4972        InvalidGboDescriptor::Unaligned,
4973        "not aligned"
4974    )]
4975    #[case::indexed_past_eof(
4976        MetadataReadPath::Indexed,
4977        InvalidGboDescriptor::PastEof,
4978        "outside file"
4979    )]
4980    #[case::indexed_overflowing(
4981        MetadataReadPath::Indexed,
4982        InvalidGboDescriptor::Overflowing,
4983        "overflows"
4984    )]
4985    #[tokio::test]
4986    async fn test_metadata_rejects_invalid_gbo_descriptor(
4987        #[case] read_path: MetadataReadPath,
4988        #[case] invalid_descriptor: InvalidGboDescriptor,
4989        #[case] expected_message: &str,
4990    ) {
4991        let fs = FsFixture::default();
4992        write_file_with_global_buffer(&fs, Bytes::from_static(b"hello")).await;
4993
4994        let mut file_bytes = fs
4995            .object_store
4996            .read_one_all(&fs.tmp_path)
4997            .await
4998            .unwrap()
4999            .to_vec();
5000        let file_len = file_bytes.len() as u64;
5001        let footer = FileReader::decode_footer(&Bytes::copy_from_slice(&file_bytes)).unwrap();
5002        let gbo_table_start = usize::try_from(footer.global_buff_offsets_start).unwrap();
5003        let alignment = PAGE_BUFFER_ALIGNMENT as u64;
5004        let (position, size) = match invalid_descriptor {
5005            InvalidGboDescriptor::Unaligned => (1, 0),
5006            InvalidGboDescriptor::PastEof => (((file_len + alignment) / alignment) * alignment, 0),
5007            InvalidGboDescriptor::Overflowing => (u64::MAX - (u64::MAX % alignment), alignment),
5008        };
5009        file_bytes[gbo_table_start..gbo_table_start + 8].copy_from_slice(&position.to_le_bytes());
5010        file_bytes[gbo_table_start + 8..gbo_table_start + 16].copy_from_slice(&size.to_le_bytes());
5011        fs.object_store
5012            .put(&fs.tmp_path, &file_bytes)
5013            .await
5014            .unwrap();
5015
5016        let scheduler = fs
5017            .scheduler
5018            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
5019            .await
5020            .unwrap();
5021        let error = match read_path {
5022            MetadataReadPath::Full => FileReader::read_all_metadata(&scheduler).await.map(|_| ()),
5023            MetadataReadPath::Indexed => FileReader::read_metadata_index(&scheduler)
5024                .await
5025                .map(|_| ()),
5026        }
5027        .expect_err("invalid GBO descriptor must fail before metadata I/O");
5028
5029        assert!(
5030            matches!(error, lance_core::Error::InvalidInput { .. }),
5031            "expected InvalidInput, got {error:?}"
5032        );
5033        assert!(
5034            error.to_string().contains(expected_message),
5035            "unexpected error: {error}"
5036        );
5037    }
5038
5039    /// A global buffer that fits inside the tail region captured at open is served
5040    /// from memory with no additional I/O.  A buffer larger than that window cannot
5041    /// fit and falls back to a dedicated read.  Both must round-trip correctly.
5042    #[rstest]
5043    #[case::within_tail_window(true)]
5044    #[case::outside_tail_window(false)]
5045    #[tokio::test]
5046    async fn test_read_global_buffer(#[case] within_window: bool) {
5047        let fs = FsFixture::default();
5048
5049        let block_size = fs.object_store.block_size();
5050        let buffer = if within_window {
5051            Bytes::from_static(b"hello")
5052        } else {
5053            Bytes::from(vec![7u8; 2 * block_size])
5054        };
5055        let expected_read_iops = if within_window { 0 } else { 1 };
5056
5057        write_file_with_global_buffer(&fs, buffer.clone()).await;
5058
5059        let file_scheduler = fs
5060            .scheduler
5061            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
5062            .await
5063            .unwrap();
5064        let file_reader = FileReader::try_open(
5065            file_scheduler,
5066            None,
5067            Arc::<DecoderPlugins>::default(),
5068            &test_cache(),
5069            FileReaderOptions::default(),
5070        )
5071        .await
5072        .unwrap();
5073
5074        // The user buffer should be retained only when it fits the tail window, and
5075        // the schema (buffer 0) is never retained.
5076        let retained = &file_reader.metadata().retained_global_buffers;
5077        assert!(!retained.contains_key(&0), "schema must not be retained");
5078        assert_eq!(retained.contains_key(&1), within_window);
5079
5080        // Reset the IO counters so we only measure the read_global_buffer call.
5081        fs.object_store.io_stats_incremental();
5082
5083        let buf = file_reader.read_global_buffer(1).await.unwrap();
5084        assert_eq!(buf, buffer);
5085
5086        let stats = fs.object_store.io_stats_incremental();
5087        assert_eq!(stats.read_iops, expected_read_iops);
5088    }
5089
5090    /// A file whose only global buffer is the schema (i.e. a plain data file, the
5091    /// common case) must retain nothing — there is no user buffer to serve.
5092    #[tokio::test]
5093    async fn test_read_global_buffer_no_user_buffers() {
5094        let fs = FsFixture::default();
5095        create_some_file(&fs, ConcreteFileVersion::V2_1).await;
5096
5097        let file_scheduler = fs
5098            .scheduler
5099            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
5100            .await
5101            .unwrap();
5102        let file_reader = FileReader::try_open(
5103            file_scheduler,
5104            None,
5105            Arc::<DecoderPlugins>::default(),
5106            &test_cache(),
5107            FileReaderOptions::default(),
5108        )
5109        .await
5110        .unwrap();
5111
5112        let metadata = file_reader.metadata();
5113        assert_eq!(metadata.file_buffers.len(), 1, "expected only the schema");
5114        assert!(
5115            metadata.retained_global_buffers.is_empty(),
5116            "a file with no user global buffers must retain nothing"
5117        );
5118    }
5119
5120    #[rstest]
5121    #[tokio::test]
5122    async fn test_deep_size_of_includes_column_metadata(
5123        #[values(
5124            ConcreteFileVersion::V2_0,
5125            ConcreteFileVersion::V2_1,
5126            ConcreteFileVersion::V2_2,
5127            ConcreteFileVersion::V2_3
5128        )]
5129        version: ConcreteFileVersion,
5130    ) {
5131        // Regression test: CachedFileMetadata::deep_size_of must account for
5132        // column_metadatas and column_infos, otherwise the moka cache weigher
5133        // dramatically underestimates entry sizes and never evicts, causing
5134        // unbounded memory growth on random-access workloads.
5135        use lance_core::deepsize::DeepSizeOf;
5136
5137        let fs = FsFixture::default();
5138        let _written = create_some_file(&fs, version).await;
5139        let cache = test_cache();
5140        let file_scheduler = fs
5141            .scheduler
5142            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
5143            .await
5144            .unwrap();
5145        let file_reader = FileReader::try_open(
5146            file_scheduler,
5147            None,
5148            Arc::<DecoderPlugins>::default(),
5149            &cache,
5150            FileReaderOptions::default(),
5151        )
5152        .await
5153        .unwrap();
5154
5155        let metadata = file_reader.metadata();
5156        let deep_size = metadata.deep_size_of();
5157
5158        // The file has multiple columns (score, location, categories, binary,
5159        // maybe large_bin). The deep_size_of must be substantially more than
5160        // just the schema — it should include column_metadatas + column_infos.
5161        // A naive implementation that ignores these fields reports < 1 KB;
5162        // a correct one should report at least several KB for this test file.
5163        assert!(
5164            deep_size > 1024,
5165            "deep_size_of ({deep_size}) is suspiciously small — \
5166             column_metadatas and column_infos may not be accounted for"
5167        );
5168
5169        // Verify column_metadatas is non-empty (sanity check).
5170        assert!(
5171            !metadata.column_metadatas.is_empty(),
5172            "Expected non-empty column_metadatas"
5173        );
5174
5175        // Verify the size scales with the number of columns: a file with more
5176        // columns should have a larger deep_size_of.
5177        let num_columns = metadata.column_metadatas.len();
5178        assert!(
5179            deep_size > num_columns * 50,
5180            "deep_size_of ({deep_size}) should scale with column count ({num_columns})"
5181        );
5182    }
5183
5184    #[tokio::test]
5185    async fn test_read_global_buffer_out_of_range() {
5186        let fs = FsFixture::default();
5187
5188        write_file_with_global_buffer(&fs, Bytes::from_static(b"hello")).await;
5189
5190        let file_scheduler = fs
5191            .scheduler
5192            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
5193            .await
5194            .unwrap();
5195        let file_reader = FileReader::try_open(
5196            file_scheduler,
5197            None,
5198            Arc::<DecoderPlugins>::default(),
5199            &test_cache(),
5200            FileReaderOptions::default(),
5201        )
5202        .await
5203        .unwrap();
5204
5205        // The file has two global buffers (schema at 0, "hello" at 1); index 2 is
5206        // out of range and must surface a descriptive error rather than panicking.
5207        let err = file_reader.read_global_buffer(2).await.unwrap_err();
5208        assert!(
5209            matches!(err, lance_core::Error::InvalidInput { .. }),
5210            "expected InvalidInput, got: {err:?}"
5211        );
5212        let msg = err.to_string();
5213        assert!(msg.contains('2'), "error should mention the index: {msg}");
5214    }
5215
5216    // Exercises the projection length-validation walk in isolation, feeding
5217    // synthetic per-column lengths so we can reach cases no current writer can
5218    // actually produce -- in particular a struct whose children diverge in
5219    // length, which the decoders would otherwise panic on or misread.
5220    #[rstest]
5221    fn test_validate_struct_child_lengths(#[values(false, true)] is_structural: bool) {
5222        let run = |dt: DataType, indices: &[u32], lengths: Vec<u64>| -> lance_core::Result<u64> {
5223            let arrow = ArrowSchema::new(vec![Field::new("s", dt, true)]);
5224            let schema = Schema::try_from(&arrow).unwrap();
5225            if is_structural {
5226                versions::v2_1::test_projection_length(&schema, indices, &lengths)
5227            } else {
5228                versions::v2_0::test_projection_length(&schema, indices, &lengths)
5229            }
5230        };
5231
5232        let struct_ty = || {
5233            DataType::Struct(Fields::from(vec![
5234                Field::new("a", DataType::Int32, true),
5235                Field::new("b", DataType::Int32, true),
5236            ]))
5237        };
5238
5239        // In 2.1 a struct contributes no column of its own (just its two leaves);
5240        // in 2.0 it also has its own column first.
5241        let (indices, equal, unequal): (&[u32], Vec<u64>, Vec<u64>) = if is_structural {
5242            (&[0, 1], vec![5, 5], vec![5, 3])
5243        } else {
5244            (&[0, 1, 2], vec![5, 5, 5], vec![5, 5, 3])
5245        };
5246
5247        assert_eq!(run(struct_ty(), indices, equal).unwrap(), 5);
5248
5249        let err = run(struct_ty(), indices, unequal).unwrap_err();
5250        let msg = err.to_string();
5251        assert!(
5252            msg.contains("differing lengths") && msg.contains('b'),
5253            "expected a child-length error naming 'b', got: {msg}"
5254        );
5255    }
5256
5257    #[test]
5258    fn test_validate_v2_0_unloaded_blob_projection_is_opaque() {
5259        let metadata = HashMap::from([(BLOB_META_KEY.to_string(), "true".to_string())]);
5260        let arrow = ArrowSchema::new(vec![
5261            Field::new("blob", DataType::LargeBinary, true).with_metadata(metadata),
5262        ]);
5263        let mut schema = Schema::try_from(&arrow).unwrap();
5264        schema.fields[0].unloaded_mut();
5265        let projection = ReaderProjection {
5266            schema: Arc::new(schema),
5267            column_indices: vec![0],
5268        };
5269        let rows = versions::v2_0::test_projection_length(
5270            &projection.schema,
5271            &projection.column_indices,
5272            &[3],
5273        )
5274        .unwrap();
5275
5276        assert_eq!(rows, 3);
5277    }
5278
5279    #[test]
5280    fn test_validate_length_list_and_empty_struct() {
5281        let validate = |dt: DataType,
5282                        is_structural: bool,
5283                        indices: &[u32],
5284                        lengths: Vec<u64>|
5285         -> lance_core::Result<u64> {
5286            let arrow = ArrowSchema::new(vec![Field::new("f", dt, true)]);
5287            let schema = Schema::try_from(&arrow).unwrap();
5288            if is_structural {
5289                versions::v2_1::test_projection_length(&schema, indices, &lengths)
5290            } else {
5291                versions::v2_0::test_projection_length(&schema, indices, &lengths)
5292            }
5293        };
5294
5295        // A list's items have a different cardinality than its rows; that gap
5296        // must not be flagged as a mismatch. In 2.0 the list is an offsets column
5297        // (rows) plus an items column (item count); in 2.1 it is a single column
5298        // whose page rows are the top-level row count.
5299        let list_ty = DataType::List(Arc::new(Field::new("item", DataType::Int32, true)));
5300        assert_eq!(
5301            validate(list_ty.clone(), false, &[0, 1], vec![5, 17]).unwrap(),
5302            5
5303        );
5304        assert_eq!(validate(list_ty, true, &[0], vec![5]).unwrap(), 5);
5305
5306        // A list of structs: the struct's children sit below the list boundary, so
5307        // their (item-count) lengths must NOT be compared against the list's row
5308        // count. In 2.0 each field has a column [list, struct, a, b] = [6, 6, 29,
5309        // 29]; the list resolves to its own row count (6) and the struct's longer
5310        // children are not flagged. (Regression: this previously errored because
5311        // the nested struct's children were checked against the struct's count.)
5312        let list_of_struct = DataType::List(Arc::new(Field::new(
5313            "item",
5314            DataType::Struct(Fields::from(vec![
5315                Field::new("a", DataType::Int32, true),
5316                Field::new("b", DataType::Int32, true),
5317            ])),
5318            true,
5319        )));
5320        assert_eq!(
5321            validate(
5322                list_of_struct.clone(),
5323                false,
5324                &[0, 1, 2, 3],
5325                vec![6, 6, 29, 29]
5326            )
5327            .unwrap(),
5328            6
5329        );
5330        // In 2.1 only the two leaves carry columns; still no false mismatch.
5331        assert_eq!(
5332            validate(list_of_struct, true, &[0, 1], vec![29, 29]).unwrap(),
5333            29
5334        );
5335
5336        // An empty struct contributes a single column and validates to its length.
5337        let empty_struct = DataType::Struct(Fields::empty());
5338        assert_eq!(
5339            validate(empty_struct.clone(), false, &[0], vec![9]).unwrap(),
5340            9
5341        );
5342        assert_eq!(validate(empty_struct, true, &[0], vec![9]).unwrap(), 9);
5343    }
5344}