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