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