1use std::{
5 borrow::Cow,
6 collections::{BTreeMap, BTreeSet},
7 io::Cursor,
8 ops::Range,
9 pin::Pin,
10 sync::Arc,
11};
12
13use arrow_array::RecordBatchReader;
14use arrow_schema::Schema as ArrowSchema;
15use byteorder::{ByteOrder, LittleEndian, ReadBytesExt};
16use bytes::{Bytes, BytesMut};
17use futures::{Stream, StreamExt, stream::BoxStream};
18use lance_core::deepsize::{Context, DeepSizeOf};
19use lance_encoding::{
20 EncodingsIo,
21 decoder::{
22 ColumnInfo, DecoderConfig, DecoderPlugins, FilterExpression, PageEncoding, PageInfo,
23 ReadBatchTask, RequestedRows, SchedulerDecoderConfig, schedule_and_decode,
24 schedule_and_decode_blocking,
25 },
26 encoder::EncodedBatch,
27 version::LanceFileVersion,
28};
29use log::debug;
30use object_store::path::Path;
31use prost::{Message, Name};
32
33use lance_core::{
34 Error, Result,
35 cache::{CacheKey, 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, MAJOR_VERSION, MINOR_VERSION, pb, pbfile},
49 io::LanceEncodingsIo,
50 version::ConcreteFileVersion,
51 writer::PAGE_BUFFER_ALIGNMENT,
52};
53
54pub const DEFAULT_READ_CHUNK_SIZE: u64 = 8 * 1024 * 1024;
57
58#[derive(Debug, DeepSizeOf)]
63pub struct BufferDescriptor {
64 pub position: u64,
65 pub size: u64,
66}
67
68#[derive(Debug)]
70pub struct FileStatistics {
71 pub columns: Vec<ColumnStatistics>,
73}
74
75#[derive(Debug)]
77pub struct ColumnStatistics {
78 pub num_pages: usize,
80 pub size_bytes: u64,
84}
85
86#[derive(Debug)]
88pub struct CachedFileMetadata {
89 pub file_schema: Arc<Schema>,
91 pub column_metadatas: Vec<pbfile::ColumnMetadata>,
93 pub column_infos: Vec<Arc<ColumnInfo>>,
94 pub num_rows: u64,
96 pub file_buffers: Vec<BufferDescriptor>,
97 pub num_data_bytes: u64,
99 pub num_column_metadata_bytes: u64,
102 pub num_global_buffer_bytes: u64,
104 pub num_footer_bytes: u64,
106 pub major_version: u16,
107 pub minor_version: u16,
108 pub file_size_bytes: u64,
110 pub retained_global_buffers: BTreeMap<u32, Bytes>,
123}
124
125impl CachedFileMetadata {
126 pub fn file_size(&self) -> u64 {
128 self.file_size_bytes
129 }
130}
131
132fn column_metadata_deep_size(column_metadatas: &[pbfile::ColumnMetadata]) -> usize {
133 column_metadatas
134 .iter()
135 .map(|cm| cm.encoded_len() * 4)
136 .sum::<usize>()
137 + std::mem::size_of_val(column_metadatas)
138}
139
140impl DeepSizeOf for CachedFileMetadata {
141 fn deep_size_of_children(&self, context: &mut Context) -> usize {
142 let schema_size = self.file_schema.deep_size_of_children(context);
143
144 let buffers_size: usize = self
145 .file_buffers
146 .iter()
147 .map(|fb| fb.deep_size_of_children(context))
148 .sum();
149
150 let column_metadatas_size = column_metadata_deep_size(self.column_metadatas.as_slice());
156
157 let column_infos_size = self.column_infos.deep_size_of_children(context);
161
162 let retained_buffers_size = self.retained_global_buffers.deep_size_of_children(context);
164
165 schema_size
166 + buffers_size
167 + column_metadatas_size
168 + column_infos_size
169 + retained_buffers_size
170 }
171}
172
173#[derive(Debug, DeepSizeOf)]
179pub struct FileMetadataIndex {
180 file_schema: Arc<Schema>,
181 num_rows: u64,
182 file_buffers: Vec<BufferDescriptor>,
183 column_metadata_offsets: Arc<[(u64, u64)]>,
184 num_columns: u32,
185 version: LanceFileVersion,
186 file_size_bytes: u64,
187 retained_global_buffers: BTreeMap<u32, Bytes>,
188}
189
190impl FileMetadataIndex {
191 pub fn file_size(&self) -> u64 {
193 self.file_size_bytes
194 }
195
196 pub fn num_columns(&self) -> u32 {
198 self.num_columns
199 }
200}
201
202#[derive(Debug)]
203struct CachedColumnMetadata {
204 column_metadata: pbfile::ColumnMetadata,
205 column_info: Arc<ColumnInfo>,
206}
207
208impl DeepSizeOf for CachedColumnMetadata {
209 fn deep_size_of_children(&self, context: &mut Context) -> usize {
210 column_metadata_deep_size(std::slice::from_ref(&self.column_metadata))
211 + self.column_info.deep_size_of_children(context)
212 }
213}
214
215#[derive(Debug, Clone)]
216struct ColumnMetadataCacheKey {
217 column_index: u32,
218}
219
220impl CacheKey for ColumnMetadataCacheKey {
221 type ValueType = CachedColumnMetadata;
222
223 fn key(&self) -> Cow<'_, str> {
224 Cow::Owned(format!("column_metadata/{}", self.column_index))
225 }
226
227 fn type_name() -> &'static str {
228 "ColumnMetadata"
229 }
230
231 fn schema() -> CacheKeySchema {
232 CacheKeySchema::new("lance.file.column-metadata-key", 1)
233 }
234
235 fn write_key(&self, builder: &mut KeyBuilder) {
236 builder.write_u32(self.column_index);
237 }
238}
239
240impl CachedFileMetadata {
241 pub fn version(&self) -> LanceFileVersion {
242 match (self.major_version, self.minor_version) {
243 (0, 3) => LanceFileVersion::V2_0,
244 (2, 0) => LanceFileVersion::V2_0,
245 (2, 1) => LanceFileVersion::V2_1,
246 (2, 2) => LanceFileVersion::V2_2,
247 (2, 3) => LanceFileVersion::V2_3,
248 _ => panic!(
249 "Unsupported version: {}.{}",
250 self.major_version, self.minor_version
251 ),
252 }
253 }
254}
255
256#[derive(Debug, Clone)]
278pub struct ReaderProjection {
279 pub schema: Arc<Schema>,
282 pub column_indices: Vec<u32>,
322}
323
324impl ReaderProjection {
325 fn from_field_ids_helper<'a>(
326 file_version: LanceFileVersion,
327 fields: impl Iterator<Item = &'a Field>,
328 field_id_to_column_index: &BTreeMap<u32, u32>,
329 column_indices: &mut Vec<u32>,
330 ) -> Result<()> {
331 for field in fields {
332 let is_structural = file_version >= LanceFileVersion::V2_1;
333 let (contributes, recurse) = field_column_shape(field, is_structural);
334 if contributes
337 && let Some(column_idx) = field_id_to_column_index.get(&(field.id as u32)).copied()
338 {
339 column_indices.push(column_idx);
340 }
341 if recurse {
342 Self::from_field_ids_helper(
343 file_version,
344 field.children.iter(),
345 field_id_to_column_index,
346 column_indices,
347 )?;
348 }
349 }
350 Ok(())
351 }
352
353 pub fn from_field_ids(
358 file_version: LanceFileVersion,
359 schema: &Schema,
360 field_id_to_column_index: &BTreeMap<u32, u32>,
361 ) -> Result<Self> {
362 let mut column_indices = Vec::new();
363 Self::from_field_ids_helper(
364 file_version,
365 schema.fields.iter(),
366 field_id_to_column_index,
367 &mut column_indices,
368 )?;
369 let projection = Self {
370 schema: Arc::new(schema.clone()),
371 column_indices,
372 };
373 Ok(projection)
374 }
375
376 pub fn from_whole_schema(schema: &Schema, version: LanceFileVersion) -> Self {
384 let schema = Arc::new(schema.clone());
385 let is_structural = version >= LanceFileVersion::V2_1;
386 let mut column_indices = vec![];
387 let mut curr_column_idx = 0;
388 let mut packed_struct_fields_num = 0;
389 for field in schema.fields_pre_order() {
390 if packed_struct_fields_num > 0 {
391 packed_struct_fields_num -= 1;
392 continue;
393 }
394 if field.is_packed_struct() {
395 column_indices.push(curr_column_idx);
396 curr_column_idx += 1;
397 packed_struct_fields_num = field.children.len();
398 } else if field.children.is_empty() || !is_structural {
399 column_indices.push(curr_column_idx);
400 curr_column_idx += 1;
401 }
402 }
403 Self {
404 schema,
405 column_indices,
406 }
407 }
408
409 pub fn from_column_names(
416 file_version: LanceFileVersion,
417 schema: &Schema,
418 column_names: &[&str],
419 ) -> Result<Self> {
420 let field_id_to_column_index = schema
421 .fields_pre_order()
422 .filter(|field| {
425 file_version < LanceFileVersion::V2_1 || field.is_leaf() || field.is_packed_struct()
426 })
427 .enumerate()
428 .map(|(idx, field)| (field.id as u32, idx as u32))
429 .collect::<BTreeMap<_, _>>();
430 let projected = schema.project(column_names)?;
431 let mut column_indices = Vec::new();
432 Self::from_field_ids_helper(
433 file_version,
434 projected.fields.iter(),
435 &field_id_to_column_index,
436 &mut column_indices,
437 )?;
438 Ok(Self {
439 schema: Arc::new(projected),
440 column_indices,
441 })
442 }
443}
444
445#[derive(Clone, Debug)]
447pub struct FileReaderOptions {
448 pub decoder_config: DecoderConfig,
449 pub read_chunk_size: u64,
453 pub batch_size_bytes: Option<u64>,
460}
461
462impl Default for FileReaderOptions {
463 fn default() -> Self {
464 Self {
465 decoder_config: DecoderConfig::default(),
466 read_chunk_size: DEFAULT_READ_CHUNK_SIZE,
467 batch_size_bytes: None,
468 }
469 }
470}
471
472#[derive(Debug, Clone)]
473struct PreparedProjection {
474 column_infos: Vec<Arc<ColumnInfo>>,
475 decoder_projection: ReaderProjection,
476}
477
478#[derive(Debug, Clone)]
479enum FileMetadataProvider {
480 Full(Arc<CachedFileMetadata>),
481 Indexed(Arc<FileMetadataIndex>),
482}
483
484#[derive(Debug, Clone)]
485struct FileReadCore {
486 scheduler: Arc<dyn EncodingsIo>,
487 base_projection: ReaderProjection,
488 metadata_provider: FileMetadataProvider,
489 decoder_plugins: Arc<DecoderPlugins>,
490 cache: Arc<LanceCache>,
491 options: FileReaderOptions,
492}
493
494#[derive(Debug, Clone)]
502pub struct ProjectedFileReader {
503 core: FileReadCore,
504}
505
506#[derive(Debug, Clone)]
508pub struct FileReader {
509 core: FileReadCore,
510 metadata: Arc<CachedFileMetadata>,
511}
512#[derive(Debug)]
513struct Footer {
514 #[allow(dead_code)]
515 column_meta_start: u64,
516 #[allow(dead_code)]
519 column_meta_offsets_start: u64,
520 global_buff_offsets_start: u64,
521 num_global_buffers: u32,
522 num_columns: u32,
523 major_version: u16,
524 minor_version: u16,
525}
526
527const FOOTER_LEN: usize = 40;
528
529fn field_column_shape(field: &Field, is_structural: bool) -> (bool, bool) {
539 if field.is_blob() || field.is_packed_struct() {
540 return (true, false);
541 }
542
543 let contributes = !is_structural || field.children.is_empty();
544 let recurse = !field.children.is_empty();
545 (contributes, recurse)
546}
547
548fn indexed_projection_column_count(field: &Field) -> Option<usize> {
557 if field.is_blob() || field.is_packed_struct() {
558 return None;
559 }
560
561 let (contributes, recurse) = field_column_shape(field, true);
562 let initial = usize::from(contributes);
563 if !recurse {
564 return Some(initial);
565 }
566
567 field.children.iter().try_fold(initial, |count, child| {
568 count.checked_add(indexed_projection_column_count(child)?)
569 })
570}
571
572fn children_share_parent_length(field: &Field) -> bool {
577 field.logical_type.is_struct()
578}
579
580fn validate_field_length<F: Fn(usize) -> Result<u64>>(
593 field: &Field,
594 is_structural: bool,
595 comparable: bool,
596 column_indices: &[u32],
597 cursor: &mut usize,
598 column_len: &F,
599) -> Result<u64> {
600 let (contributes, recurse) = field_column_shape(field, is_structural);
601 let mut field_rows: Option<u64> = None;
602 if contributes {
603 let column = *column_indices.get(*cursor).ok_or_else(|| {
604 Error::invalid_input(format!(
605 "projection supplied fewer column indices than its fields require \
606 (ran out at field '{}')",
607 field.name
608 ))
609 })?;
610 *cursor += 1;
611 field_rows = Some(column_len(column as usize)?);
612 }
613 if recurse {
614 let enforce_children = comparable && children_share_parent_length(field);
619 for child in &field.children {
620 let child_rows = validate_field_length(
621 child,
622 is_structural,
623 enforce_children,
624 column_indices,
625 cursor,
626 column_len,
627 )?;
628 let expected = *field_rows.get_or_insert(child_rows);
631 if enforce_children && child_rows != expected {
632 return Err(Error::invalid_input(format!(
633 "cannot read field '{}': its children have differing lengths \
634 (child '{}' has {} rows, but the field has {}); a struct's \
635 children must all have the same length",
636 field.name, child.name, child_rows, expected
637 )));
638 }
639 }
640 }
641 field_rows.ok_or_else(|| {
642 Error::invalid_input(format!(
643 "projected field '{}' maps to no columns",
644 field.name
645 ))
646 })
647}
648
649fn verify_uniform_lengths(field_lengths: &[(&str, u64)]) -> Result<u64> {
655 let first = field_lengths.first().map_or(0, |&(_, len)| len);
656 if field_lengths.iter().all(|&(_, len)| len == first) {
657 return Ok(first);
658 }
659 let columns = field_lengths
660 .iter()
661 .map(|(name, len)| format!("{name}={len}"))
662 .collect::<Vec<_>>()
663 .join(", ");
664 Err(Error::invalid_input(format!(
665 "cannot read columns of differing lengths together ({columns}); \
666 read each column (or equal-length group) separately"
667 )))
668}
669
670impl FileReader {
671 pub fn with_scheduler(&self, scheduler: Arc<dyn EncodingsIo>) -> Self {
672 Self {
673 core: self.core.with_scheduler(scheduler),
674 metadata: self.metadata.clone(),
675 }
676 }
677
678 pub fn with_io_stats(
686 &self,
687 stats: Arc<dyn lance_core::utils::io_stats::IoStatsRecorder>,
688 ) -> Self {
689 match self.core.scheduler.with_io_stats(stats) {
690 Some(scheduler) => self.with_scheduler(scheduler),
691 None => self.clone(),
692 }
693 }
694
695 pub fn num_rows(&self) -> u64 {
696 self.core.num_rows()
697 }
698
699 pub fn column_num_rows(&self, column_index: usize) -> Result<u64> {
708 let column = self
709 .metadata
710 .column_metadatas
711 .get(column_index)
712 .ok_or_else(|| {
713 Error::invalid_input(format!(
714 "column index {} is out of bounds (file has {} columns)",
715 column_index,
716 self.metadata.column_metadatas.len()
717 ))
718 })?;
719 Ok(column.pages.iter().map(|page| page.length).sum())
720 }
721
722 pub fn metadata(&self) -> &Arc<CachedFileMetadata> {
723 &self.metadata
724 }
725
726 fn statistics_from_column_metadata(
727 column_metadatas: &[pbfile::ColumnMetadata],
728 ) -> FileStatistics {
729 let column_stats = column_metadatas
730 .iter()
731 .map(|col_metadata| {
732 let num_pages = col_metadata.pages.len();
733 let size_bytes = col_metadata
734 .pages
735 .iter()
736 .map(|page| page.buffer_sizes.iter().sum::<u64>())
737 .sum::<u64>();
738 ColumnStatistics {
739 num_pages,
740 size_bytes,
741 }
742 })
743 .collect();
744
745 FileStatistics {
746 columns: column_stats,
747 }
748 }
749
750 pub fn file_statistics(&self) -> FileStatistics {
751 Self::statistics_from_column_metadata(&self.metadata().column_metadatas)
752 }
753
754 pub async fn read_global_buffer(&self, index: u32) -> Result<Bytes> {
755 self.core.read_global_buffer(index).await
756 }
757
758 async fn read_tail(scheduler: &FileScheduler) -> Result<(Bytes, u64)> {
759 let file_size = scheduler.reader().size().await? as u64;
760 let begin = if file_size < scheduler.reader().block_size() as u64 {
761 0
762 } else {
763 file_size - scheduler.reader().block_size() as u64
764 };
765 let tail_bytes = scheduler.submit_single(begin..file_size, 0).await?;
766 Ok((tail_bytes, file_size))
767 }
768
769 async fn read_range_from_tail_or_scheduler(
770 tail_bytes: &Bytes,
771 tail_offset: u64,
772 scheduler: &FileScheduler,
773 range: Range<u64>,
774 ) -> Result<Bytes> {
775 let tail_end = tail_offset + tail_bytes.len() as u64;
776 if range.start >= tail_offset && range.end <= tail_end {
777 let rel_start = (range.start - tail_offset) as usize;
778 let rel_end = (range.end - tail_offset) as usize;
779 Ok(tail_bytes.slice(rel_start..rel_end))
780 } else {
781 scheduler.submit_single(range, 0).await
782 }
783 }
784
785 fn retained_global_buffers_from_tail(
786 gbo_table: &[BufferDescriptor],
787 tail_bytes: &Bytes,
788 tail_offset: u64,
789 ) -> BTreeMap<u32, Bytes> {
790 let tail_end = tail_offset + tail_bytes.len() as u64;
791 gbo_table
792 .iter()
793 .enumerate()
794 .skip(1)
795 .filter_map(|(index, buffer)| {
796 let start = buffer.position;
797 let end = buffer.position + buffer.size;
798 if start >= tail_offset && end <= tail_end {
799 let rel_start = (start - tail_offset) as usize;
800 let rel_end = (end - tail_offset) as usize;
801 let bytes = Bytes::copy_from_slice(&tail_bytes[rel_start..rel_end]);
802 Some((index as u32, bytes))
803 } else {
804 None
805 }
806 })
807 .collect()
808 }
809
810 fn decode_footer(footer_bytes: &Bytes) -> Result<Footer> {
813 let len = footer_bytes.len();
814 if len < FOOTER_LEN {
815 return Err(Error::invalid_input(format!(
816 "does not have sufficient data, len: {}, bytes: {:?}",
817 len, footer_bytes
818 )));
819 }
820 let mut cursor = Cursor::new(footer_bytes.slice(len - FOOTER_LEN..));
821
822 let column_meta_start = cursor.read_u64::<LittleEndian>()?;
823 let column_meta_offsets_start = cursor.read_u64::<LittleEndian>()?;
824 let global_buff_offsets_start = cursor.read_u64::<LittleEndian>()?;
825 let num_global_buffers = cursor.read_u32::<LittleEndian>()?;
826 let num_columns = cursor.read_u32::<LittleEndian>()?;
827 let major_version = cursor.read_u16::<LittleEndian>()?;
828 let minor_version = cursor.read_u16::<LittleEndian>()?;
829
830 if major_version == MAJOR_VERSION as u16 && minor_version == MINOR_VERSION as u16 {
831 return Err(Error::version_conflict(
832 "Attempt to use the lance v2 reader to read a legacy file".to_string(),
833 major_version,
834 minor_version,
835 ));
836 }
837
838 let magic_bytes = footer_bytes.slice(len - 4..);
839 if magic_bytes.as_ref() != MAGIC {
840 return Err(Error::invalid_input(format!(
841 "file does not appear to be a Lance file (invalid magic: {:?})",
842 MAGIC
843 )));
844 }
845 Ok(Footer {
846 column_meta_start,
847 column_meta_offsets_start,
848 global_buff_offsets_start,
849 num_global_buffers,
850 num_columns,
851 major_version,
852 minor_version,
853 })
854 }
855
856 fn read_all_column_metadata(
858 column_metadata_bytes: Bytes,
859 footer: &Footer,
860 ) -> Result<Vec<pbfile::ColumnMetadata>> {
861 let column_metadata_start = footer.column_meta_start;
862 let cmo_table_size = 16 * footer.num_columns as usize;
864 if column_metadata_bytes.len() < cmo_table_size {
865 return Err(Error::invalid_input(format!(
866 "column metadata region has {} bytes but CMO table needs {} bytes for {} columns",
867 column_metadata_bytes.len(),
868 cmo_table_size,
869 footer.num_columns
870 )));
871 }
872 let cmo_table = column_metadata_bytes.slice(column_metadata_bytes.len() - cmo_table_size..);
873 let column_metadata_offsets = Self::decode_cmo_table(cmo_table, footer)?;
874
875 column_metadata_offsets
876 .iter()
877 .map(|(position, length)| {
878 let normalized_position = (*position - column_metadata_start) as usize;
879 let normalized_end = normalized_position + (*length as usize);
880 Ok(pbfile::ColumnMetadata::decode(
881 &column_metadata_bytes[normalized_position..normalized_end],
882 )?)
883 })
884 .collect::<Result<Vec<_>>>()
885 }
886
887 fn decode_cmo_table(cmo_table: Bytes, footer: &Footer) -> Result<Arc<[(u64, u64)]>> {
888 let expected_size = 16 * footer.num_columns as usize;
889 if cmo_table.len() != expected_size {
890 return Err(Error::invalid_input(format!(
891 "column metadata offset table has {} bytes but expected {} bytes for {} columns",
892 cmo_table.len(),
893 expected_size,
894 footer.num_columns
895 )));
896 }
897
898 let mut offsets = Vec::with_capacity(footer.num_columns as usize);
899 for col_idx in 0..footer.num_columns {
900 let offset = (col_idx * 16) as usize;
901 let position = LittleEndian::read_u64(&cmo_table[offset..offset + 8]);
902 let length = LittleEndian::read_u64(&cmo_table[offset + 8..offset + 16]);
903 let end = position.checked_add(length).ok_or_else(|| {
904 Error::invalid_input(format!(
905 "column metadata range overflows for column index {}, position={}, length={}",
906 col_idx, position, length
907 ))
908 })?;
909 if position < footer.column_meta_start || end > footer.column_meta_offsets_start {
910 return Err(Error::invalid_input(format!(
911 "column metadata range for column index {} is outside metadata region: position={}, length={}, metadata_start={}, cmo_start={}",
912 col_idx,
913 position,
914 length,
915 footer.column_meta_start,
916 footer.column_meta_offsets_start
917 )));
918 }
919 offsets.push((position, length));
920 }
921
922 Ok(Arc::from(offsets))
923 }
924
925 async fn optimistic_tail_read(
926 data: &Bytes,
927 start_pos: u64,
928 scheduler: &FileScheduler,
929 file_len: u64,
930 ) -> Result<Bytes> {
931 let num_bytes_needed = (file_len - start_pos) as usize;
932 if data.len() >= num_bytes_needed {
933 Ok(data.slice((data.len() - num_bytes_needed)..))
934 } else {
935 let num_bytes_missing = (num_bytes_needed - data.len()) as u64;
936 let start = file_len - num_bytes_needed as u64;
937 let missing_bytes = scheduler
938 .submit_single(start..start + num_bytes_missing, 0)
939 .await?;
940 let mut combined = BytesMut::with_capacity(data.len() + num_bytes_missing as usize);
941 combined.extend(missing_bytes);
942 combined.extend(data);
943 Ok(combined.freeze())
944 }
945 }
946
947 fn do_decode_gbo_table(
948 gbo_bytes: &Bytes,
949 footer: &Footer,
950 version: LanceFileVersion,
951 ) -> Result<Vec<BufferDescriptor>> {
952 let mut global_bufs_cursor = Cursor::new(gbo_bytes);
953
954 let mut global_buffers = Vec::with_capacity(footer.num_global_buffers as usize);
955 for buffer_index in 0..footer.num_global_buffers {
956 let buf_pos = global_bufs_cursor.read_u64::<LittleEndian>()?;
957 if version >= LanceFileVersion::V2_1 && buf_pos % PAGE_BUFFER_ALIGNMENT as u64 != 0 {
958 return Err(Error::invalid_input_source(
959 format!(
960 "Global buffer {} position {} is not aligned to {} bytes",
961 buffer_index, buf_pos, PAGE_BUFFER_ALIGNMENT
962 )
963 .into(),
964 ));
965 }
966 let buf_size = global_bufs_cursor.read_u64::<LittleEndian>()?;
967 global_buffers.push(BufferDescriptor {
968 position: buf_pos,
969 size: buf_size,
970 });
971 }
972
973 Ok(global_buffers)
974 }
975
976 async fn decode_gbo_table(
977 tail_bytes: &Bytes,
978 file_len: u64,
979 scheduler: &FileScheduler,
980 footer: &Footer,
981 version: LanceFileVersion,
982 ) -> Result<Vec<BufferDescriptor>> {
983 let gbo_bytes = Self::optimistic_tail_read(
986 tail_bytes,
987 footer.global_buff_offsets_start,
988 scheduler,
989 file_len,
990 )
991 .await?;
992 Self::do_decode_gbo_table(&gbo_bytes, footer, version)
993 }
994
995 fn decode_schema(schema_bytes: Bytes) -> Result<(u64, lance_core::datatypes::Schema)> {
996 let file_descriptor = pb::FileDescriptor::decode(schema_bytes)?;
997 let pb_schema = file_descriptor.schema.unwrap();
998 let num_rows = file_descriptor.length;
999 let fields_with_meta = FieldsWithMeta {
1000 fields: Fields(pb_schema.fields),
1001 metadata: pb_schema.metadata,
1002 };
1003 let schema = Schema::try_from(fields_with_meta)?;
1004 Ok((num_rows, schema))
1005 }
1006
1007 pub async fn read_all_metadata(scheduler: &FileScheduler) -> Result<CachedFileMetadata> {
1021 let (tail_bytes, file_len) = Self::read_tail(scheduler).await?;
1023 let tail_offset = file_len - tail_bytes.len() as u64;
1024 let footer = Self::decode_footer(&tail_bytes)?;
1025
1026 let file_version: LanceFileVersion =
1027 ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)?
1028 .into();
1029
1030 let gbo_table =
1031 Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer, file_version).await?;
1032 if gbo_table.is_empty() {
1033 return Err(Error::internal(
1034 "File did not contain any global buffers, schema expected".to_string(),
1035 ));
1036 }
1037 let schema_start = gbo_table[0].position;
1038 let schema_size = gbo_table[0].size;
1039
1040 let num_footer_bytes = file_len - schema_start;
1041
1042 let all_metadata_bytes =
1045 Self::optimistic_tail_read(&tail_bytes, schema_start, scheduler, file_len).await?;
1046
1047 let schema_bytes = all_metadata_bytes.slice(0..schema_size as usize);
1048 let (num_rows, schema) = Self::decode_schema(schema_bytes)?;
1049
1050 let column_metadata_start = (footer.column_meta_start - schema_start) as usize;
1053 let column_metadata_end = (footer.global_buff_offsets_start - schema_start) as usize;
1054 let column_metadata_bytes =
1055 all_metadata_bytes.slice(column_metadata_start..column_metadata_end);
1056 let column_metadatas = Self::read_all_column_metadata(column_metadata_bytes, &footer)?;
1057
1058 let num_global_buffer_bytes = gbo_table.iter().map(|buf| buf.size).sum::<u64>();
1059 let num_data_bytes = footer.column_meta_start - num_global_buffer_bytes;
1060 let num_column_metadata_bytes = footer.global_buff_offsets_start - footer.column_meta_start;
1061
1062 let column_infos = Self::meta_to_col_infos(column_metadatas.as_slice(), file_version)?;
1063
1064 let retained_global_buffers =
1070 Self::retained_global_buffers_from_tail(&gbo_table, &tail_bytes, tail_offset);
1071
1072 Ok(CachedFileMetadata {
1073 file_schema: Arc::new(schema),
1074 column_metadatas,
1075 column_infos,
1076 num_rows,
1077 num_data_bytes,
1078 num_column_metadata_bytes,
1079 num_global_buffer_bytes,
1080 num_footer_bytes,
1081 file_buffers: gbo_table,
1082 major_version: footer.major_version,
1083 minor_version: footer.minor_version,
1084 file_size_bytes: file_len,
1085 retained_global_buffers,
1086 })
1087 }
1088
1089 async fn read_metadata_index_with_known_schema(
1090 scheduler: &FileScheduler,
1091 known_schema: Option<(Arc<Schema>, u64)>,
1092 ) -> Result<FileMetadataIndex> {
1093 let (tail_bytes, file_len) = Self::read_tail(scheduler).await?;
1094 let tail_offset = file_len - tail_bytes.len() as u64;
1095 let footer = Self::decode_footer(&tail_bytes)?;
1096
1097 let file_version: LanceFileVersion =
1098 ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)?
1099 .into();
1100
1101 let gbo_table =
1102 Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer, file_version).await?;
1103 if gbo_table.is_empty() {
1104 return Err(Error::internal(
1105 "File did not contain any global buffers, schema expected".to_string(),
1106 ));
1107 }
1108 let (file_schema, num_rows) = match known_schema {
1109 Some((file_schema, num_rows)) => (file_schema, num_rows),
1110 None => {
1111 let schema_buffer = &gbo_table[0];
1112 let schema_bytes = Self::read_range_from_tail_or_scheduler(
1113 &tail_bytes,
1114 tail_offset,
1115 scheduler,
1116 schema_buffer.position..schema_buffer.position + schema_buffer.size,
1117 )
1118 .await?;
1119 let (num_rows, schema) = Self::decode_schema(schema_bytes)?;
1120 (Arc::new(schema), num_rows)
1121 }
1122 };
1123
1124 let cmo_table = Self::read_range_from_tail_or_scheduler(
1125 &tail_bytes,
1126 tail_offset,
1127 scheduler,
1128 footer.column_meta_offsets_start..footer.global_buff_offsets_start,
1129 )
1130 .await?;
1131 let column_metadata_offsets = Self::decode_cmo_table(cmo_table, &footer)?;
1132
1133 let retained_global_buffers =
1134 Self::retained_global_buffers_from_tail(&gbo_table, &tail_bytes, tail_offset);
1135
1136 Ok(FileMetadataIndex {
1137 file_schema,
1138 num_rows,
1139 file_buffers: gbo_table,
1140 column_metadata_offsets,
1141 num_columns: footer.num_columns,
1142 version: file_version,
1143 file_size_bytes: file_len,
1144 retained_global_buffers,
1145 })
1146 }
1147
1148 pub async fn read_metadata_index(scheduler: &FileScheduler) -> Result<FileMetadataIndex> {
1154 Self::read_metadata_index_with_known_schema(scheduler, None).await
1155 }
1156
1157 pub async fn read_metadata_index_with_schema(
1162 scheduler: &FileScheduler,
1163 file_schema: Arc<Schema>,
1164 num_rows: u64,
1165 ) -> Result<FileMetadataIndex> {
1166 Self::read_metadata_index_with_known_schema(scheduler, Some((file_schema, num_rows))).await
1167 }
1168
1169 fn fetch_encoding<M: Default + Name + Sized>(encoding: &pbfile::Encoding) -> Result<M> {
1170 match &encoding.location {
1171 Some(pbfile::encoding::Location::Indirect(_)) => Err(Error::invalid_input_source(
1172 "Indirect file encodings are not supported".into(),
1173 )),
1174 Some(pbfile::encoding::Location::Direct(encoding)) => {
1175 let encoding_buf = Bytes::from(encoding.encoding.clone());
1176 let encoding_any = prost_types::Any::decode(encoding_buf).map_err(|error| {
1177 Error::invalid_input_source(
1178 format!("Invalid direct {} encoding envelope: {error}", M::NAME).into(),
1179 )
1180 })?;
1181 encoding_any.to_msg::<M>().map_err(|error| {
1182 Error::invalid_input_source(
1183 format!("Invalid direct {} encoding: {error}", M::NAME).into(),
1184 )
1185 })
1186 }
1187 Some(pbfile::encoding::Location::None(_)) => Err(Error::invalid_input_source(
1188 format!("Missing {} encoding description", M::NAME).into(),
1189 )),
1190 None => Err(Error::invalid_input_source(
1191 format!("Missing {} encoding location", M::NAME).into(),
1192 )),
1193 }
1194 }
1195
1196 fn meta_to_col_infos(
1197 column_metadatas: &[pbfile::ColumnMetadata],
1198 file_version: LanceFileVersion,
1199 ) -> Result<Vec<Arc<ColumnInfo>>> {
1200 column_metadatas
1201 .iter()
1202 .enumerate()
1203 .map(|(col_idx, col_meta)| {
1204 let col_idx = u32::try_from(col_idx).map_err(|_| {
1205 Error::invalid_input_source("File has more than u32::MAX columns".into())
1206 })?;
1207 Self::meta_to_col_info(col_idx, col_meta, file_version)
1208 })
1209 .collect()
1210 }
1211
1212 fn meta_to_col_info(
1213 col_idx: u32,
1214 col_meta: &pbfile::ColumnMetadata,
1215 file_version: LanceFileVersion,
1216 ) -> Result<Arc<ColumnInfo>> {
1217 let page_infos = col_meta
1218 .pages
1219 .iter()
1220 .enumerate()
1221 .map(|(page_idx, page)| {
1222 let num_rows = page.length;
1223 let encoding = match file_version {
1224 LanceFileVersion::V2_0 => {
1225 PageEncoding::Legacy(Self::fetch_encoding::<pbenc::ArrayEncoding>(
1226 page.encoding.as_ref().ok_or_else(|| {
1227 Error::invalid_input_source(
1228 format!(
1229 "Column {} page {} is missing its encoding",
1230 col_idx, page_idx
1231 )
1232 .into(),
1233 )
1234 })?,
1235 )?)
1236 }
1237 _ => {
1238 PageEncoding::Structural(Self::fetch_encoding::<pbenc21::PageLayout>(
1239 page.encoding.as_ref().ok_or_else(|| {
1240 Error::invalid_input_source(
1241 format!(
1242 "Column {} page {} is missing its encoding",
1243 col_idx, page_idx
1244 )
1245 .into(),
1246 )
1247 })?,
1248 )?)
1249 }
1250 };
1251 if page.buffer_offsets.len() != page.buffer_sizes.len() {
1252 return Err(Error::invalid_input_source(
1253 format!(
1254 "Column {} page {} has {} buffer offsets but {} buffer sizes",
1255 col_idx,
1256 page_idx,
1257 page.buffer_offsets.len(),
1258 page.buffer_sizes.len()
1259 )
1260 .into(),
1261 ));
1262 }
1263 let buffer_offsets_and_sizes = Arc::from(
1264 page.buffer_offsets
1265 .iter()
1266 .zip(page.buffer_sizes.iter())
1267 .map(|(offset, size)| -> Result<_> {
1268 if file_version >= LanceFileVersion::V2_1
1269 && offset % PAGE_BUFFER_ALIGNMENT as u64 != 0
1270 {
1271 return Err(Error::invalid_input_source(
1272 format!(
1273 "Column {} page {} buffer offset {} is not aligned to {} bytes",
1274 col_idx, page_idx, offset, PAGE_BUFFER_ALIGNMENT
1275 )
1276 .into(),
1277 ));
1278 }
1279 Ok((*offset, *size))
1280 })
1281 .collect::<Result<Vec<_>>>()?,
1282 );
1283 Ok(PageInfo {
1284 buffer_offsets_and_sizes,
1285 encoding,
1286 num_rows,
1287 priority: page.priority,
1288 })
1289 })
1290 .collect::<Result<Vec<_>>>()?;
1291 if col_meta.buffer_offsets.len() != col_meta.buffer_sizes.len() {
1292 return Err(Error::invalid_input_source(
1293 format!(
1294 "Column {} has {} buffer offsets but {} buffer sizes",
1295 col_idx,
1296 col_meta.buffer_offsets.len(),
1297 col_meta.buffer_sizes.len()
1298 )
1299 .into(),
1300 ));
1301 }
1302 let buffer_offsets_and_sizes = Arc::from(
1303 col_meta
1304 .buffer_offsets
1305 .iter()
1306 .zip(col_meta.buffer_sizes.iter())
1307 .map(|(offset, size)| (*offset, *size))
1308 .collect::<Vec<_>>(),
1309 );
1310 Ok(Arc::new(ColumnInfo {
1311 index: col_idx,
1312 page_infos: Arc::from(page_infos),
1313 buffer_offsets_and_sizes,
1314 encoding: Self::fetch_encoding(col_meta.encoding.as_ref().ok_or_else(|| {
1315 Error::invalid_input_source(
1316 format!("Column {} is missing its encoding", col_idx).into(),
1317 )
1318 })?)?,
1319 }))
1320 }
1321
1322 fn validate_projection(
1323 projection: &ReaderProjection,
1324 metadata: &CachedFileMetadata,
1325 ) -> Result<()> {
1326 if projection.schema.fields.is_empty() {
1327 return Err(Error::invalid_input(
1328 "Attempt to read zero columns from the file, at least one column must be specified"
1329 .to_string(),
1330 ));
1331 }
1332 let mut column_indices_seen = BTreeSet::new();
1333 for column_index in &projection.column_indices {
1334 if !column_indices_seen.insert(*column_index) {
1335 return Err(Error::invalid_input(format!(
1336 "The projection specified the column index {} more than once",
1337 column_index
1338 )));
1339 }
1340 if *column_index >= metadata.column_infos.len() as u32 {
1341 return Err(Error::invalid_input(format!(
1342 "The projection specified the column index {} but there are only {} columns in the file",
1343 column_index,
1344 metadata.column_infos.len()
1345 )));
1346 }
1347 }
1348 Ok(())
1349 }
1350
1351 pub async fn try_open(
1358 scheduler: FileScheduler,
1359 base_projection: Option<ReaderProjection>,
1360 decoder_plugins: Arc<DecoderPlugins>,
1361 cache: &LanceCache,
1362 options: FileReaderOptions,
1363 ) -> Result<Self> {
1364 let file_metadata = Arc::new(Self::read_all_metadata(&scheduler).await?);
1365 let path = scheduler.reader().path().clone();
1366
1367 let encodings_io =
1369 LanceEncodingsIo::new(scheduler).with_read_chunk_size(options.read_chunk_size);
1370
1371 Self::try_open_with_file_metadata(
1372 Arc::new(encodings_io),
1373 path,
1374 base_projection,
1375 decoder_plugins,
1376 file_metadata,
1377 cache,
1378 options,
1379 )
1380 .await
1381 }
1382
1383 pub async fn try_open_with_file_metadata(
1390 scheduler: Arc<dyn EncodingsIo>,
1391 path: Path,
1392 base_projection: Option<ReaderProjection>,
1393 decoder_plugins: Arc<DecoderPlugins>,
1394 file_metadata: Arc<CachedFileMetadata>,
1395 cache: &LanceCache,
1396 options: FileReaderOptions,
1397 ) -> Result<Self> {
1398 let cache = Arc::new(cache.with_key_prefix(path.as_ref()));
1399 let core = FileReadCore::try_new(
1400 scheduler,
1401 base_projection,
1402 decoder_plugins,
1403 FileMetadataProvider::Full(file_metadata.clone()),
1404 cache,
1405 options,
1406 )?;
1407 Ok(Self {
1408 core,
1409 metadata: file_metadata,
1410 })
1411 }
1412
1413 fn collect_columns_from_projection(
1427 &self,
1428 _projection: &ReaderProjection,
1429 ) -> Result<Vec<Arc<ColumnInfo>>> {
1430 Ok(self.metadata.column_infos.clone())
1431 }
1432
1433 #[allow(clippy::too_many_arguments)]
1434 async fn do_read_range(
1435 column_infos: Vec<Arc<ColumnInfo>>,
1436 io: Arc<dyn EncodingsIo>,
1437 cache: Arc<LanceCache>,
1438 num_rows: u64,
1439 decoder_plugins: Arc<DecoderPlugins>,
1440 range: Range<u64>,
1441 batch_size: u32,
1442 projection: ReaderProjection,
1443 filter: FilterExpression,
1444 decoder_config: DecoderConfig,
1445 batch_size_bytes: Option<u64>,
1446 ) -> Result<BoxStream<'static, ReadBatchTask>> {
1447 debug!(
1448 "Reading range {:?} with batch_size {} from file with {} rows and {} columns into schema with {} columns",
1449 range,
1450 batch_size,
1451 num_rows,
1452 column_infos.len(),
1453 projection.schema.fields.len(),
1454 );
1455
1456 let config = SchedulerDecoderConfig {
1457 batch_size,
1458 cache,
1459 decoder_plugins,
1460 io,
1461 decoder_config,
1462 batch_size_bytes,
1463 };
1464
1465 let requested_rows = RequestedRows::Ranges(vec![range]);
1466
1467 schedule_and_decode(
1468 column_infos,
1469 requested_rows,
1470 filter,
1471 projection.column_indices,
1472 projection.schema,
1473 config,
1474 )
1475 .await
1476 }
1477
1478 #[allow(clippy::too_many_arguments)]
1479 async fn do_take_rows(
1480 column_infos: Vec<Arc<ColumnInfo>>,
1481 io: Arc<dyn EncodingsIo>,
1482 cache: Arc<LanceCache>,
1483 decoder_plugins: Arc<DecoderPlugins>,
1484 indices: Vec<u64>,
1485 batch_size: u32,
1486 projection: ReaderProjection,
1487 filter: FilterExpression,
1488 decoder_config: DecoderConfig,
1489 batch_size_bytes: Option<u64>,
1490 ) -> Result<BoxStream<'static, ReadBatchTask>> {
1491 debug!(
1492 "Taking {} rows spread across range {}..{} with batch_size {} from columns {:?}",
1493 indices.len(),
1494 indices[0],
1495 indices[indices.len() - 1],
1496 batch_size,
1497 column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1498 );
1499
1500 let config = SchedulerDecoderConfig {
1501 batch_size,
1502 cache,
1503 decoder_plugins,
1504 io,
1505 decoder_config,
1506 batch_size_bytes,
1507 };
1508
1509 let requested_rows = RequestedRows::Indices(indices);
1510
1511 schedule_and_decode(
1512 column_infos,
1513 requested_rows,
1514 filter,
1515 projection.column_indices,
1516 projection.schema,
1517 config,
1518 )
1519 .await
1520 }
1521
1522 #[allow(clippy::too_many_arguments)]
1523 async fn do_read_ranges(
1524 column_infos: Vec<Arc<ColumnInfo>>,
1525 io: Arc<dyn EncodingsIo>,
1526 cache: Arc<LanceCache>,
1527 decoder_plugins: Arc<DecoderPlugins>,
1528 ranges: Vec<Range<u64>>,
1529 batch_size: u32,
1530 projection: ReaderProjection,
1531 filter: FilterExpression,
1532 decoder_config: DecoderConfig,
1533 batch_size_bytes: Option<u64>,
1534 ) -> Result<BoxStream<'static, ReadBatchTask>> {
1535 let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
1536 debug!(
1537 "Taking {} ranges ({} rows) spread across range {}..{} with batch_size {} from columns {:?}",
1538 ranges.len(),
1539 num_rows,
1540 ranges[0].start,
1541 ranges[ranges.len() - 1].end,
1542 batch_size,
1543 column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1544 );
1545
1546 let config = SchedulerDecoderConfig {
1547 batch_size,
1548 cache,
1549 decoder_plugins,
1550 io,
1551 decoder_config,
1552 batch_size_bytes,
1553 };
1554
1555 let requested_rows = RequestedRows::Ranges(ranges);
1556
1557 schedule_and_decode(
1558 column_infos,
1559 requested_rows,
1560 filter,
1561 projection.column_indices,
1562 projection.schema,
1563 config,
1564 )
1565 .await
1566 }
1567
1568 pub async fn read_tasks(
1592 &self,
1593 params: ReadBatchParams,
1594 batch_size: u32,
1595 projection: Option<ReaderProjection>,
1596 filter: FilterExpression,
1597 ) -> Result<Pin<Box<dyn Stream<Item = ReadBatchTask> + Send>>> {
1598 self.core
1599 .read_tasks(params, batch_size, projection, filter)
1600 .await
1601 }
1602
1603 pub async fn read_stream_projected(
1633 &self,
1634 params: ReadBatchParams,
1635 batch_size: u32,
1636 batch_readahead: u32,
1637 projection: ReaderProjection,
1638 filter: FilterExpression,
1639 ) -> Result<Pin<Box<dyn RecordBatchStream>>> {
1640 let arrow_schema = Arc::new(ArrowSchema::from(projection.schema.as_ref()));
1641 let tasks_stream = self
1642 .read_tasks(params, batch_size, Some(projection), filter)
1643 .await?;
1644 let batch_stream = tasks_stream
1645 .map(|task| task.task)
1646 .buffered(batch_readahead as usize)
1647 .boxed();
1648 Ok(Box::pin(RecordBatchStreamAdapter::new(
1649 arrow_schema,
1650 batch_stream,
1651 )))
1652 }
1653
1654 fn take_rows_blocking(
1655 &self,
1656 indices: Vec<u64>,
1657 batch_size: u32,
1658 projection: ReaderProjection,
1659 filter: FilterExpression,
1660 ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1661 let column_infos = self.collect_columns_from_projection(&projection)?;
1662 debug!(
1663 "Taking {} rows spread across range {}..{} with batch_size {} from columns {:?}",
1664 indices.len(),
1665 indices[0],
1666 indices[indices.len() - 1],
1667 batch_size,
1668 column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1669 );
1670
1671 let config = SchedulerDecoderConfig {
1672 batch_size,
1673 cache: self.core.cache.clone(),
1674 decoder_plugins: self.core.decoder_plugins.clone(),
1675 io: self.core.scheduler.clone(),
1676 decoder_config: self.core.options.decoder_config.clone(),
1677 batch_size_bytes: self.core.options.batch_size_bytes,
1678 };
1679
1680 let requested_rows = RequestedRows::Indices(indices);
1681
1682 schedule_and_decode_blocking(
1683 column_infos,
1684 requested_rows,
1685 filter,
1686 projection.column_indices,
1687 projection.schema,
1688 config,
1689 )
1690 }
1691
1692 fn read_ranges_blocking(
1693 &self,
1694 ranges: Vec<Range<u64>>,
1695 batch_size: u32,
1696 projection: ReaderProjection,
1697 filter: FilterExpression,
1698 ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1699 let column_infos = self.collect_columns_from_projection(&projection)?;
1700 let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
1701 debug!(
1702 "Taking {} ranges ({} rows) spread across range {}..{} with batch_size {} from columns {:?}",
1703 ranges.len(),
1704 num_rows,
1705 ranges[0].start,
1706 ranges[ranges.len() - 1].end,
1707 batch_size,
1708 column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1709 );
1710
1711 let config = SchedulerDecoderConfig {
1712 batch_size,
1713 cache: self.core.cache.clone(),
1714 decoder_plugins: self.core.decoder_plugins.clone(),
1715 io: self.core.scheduler.clone(),
1716 decoder_config: self.core.options.decoder_config.clone(),
1717 batch_size_bytes: self.core.options.batch_size_bytes,
1718 };
1719
1720 let requested_rows = RequestedRows::Ranges(ranges);
1721
1722 schedule_and_decode_blocking(
1723 column_infos,
1724 requested_rows,
1725 filter,
1726 projection.column_indices,
1727 projection.schema,
1728 config,
1729 )
1730 }
1731
1732 fn read_range_blocking(
1733 &self,
1734 range: Range<u64>,
1735 batch_size: u32,
1736 projection: ReaderProjection,
1737 filter: FilterExpression,
1738 ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1739 let column_infos = self.collect_columns_from_projection(&projection)?;
1740 let num_rows = self.core.num_rows();
1741
1742 debug!(
1743 "Reading range {:?} with batch_size {} from file with {} rows and {} columns into schema with {} columns",
1744 range,
1745 batch_size,
1746 num_rows,
1747 column_infos.len(),
1748 projection.schema.fields.len(),
1749 );
1750
1751 let config = SchedulerDecoderConfig {
1752 batch_size,
1753 cache: self.core.cache.clone(),
1754 decoder_plugins: self.core.decoder_plugins.clone(),
1755 io: self.core.scheduler.clone(),
1756 decoder_config: self.core.options.decoder_config.clone(),
1757 batch_size_bytes: self.core.options.batch_size_bytes,
1758 };
1759
1760 let requested_rows = RequestedRows::Ranges(vec![range]);
1761
1762 schedule_and_decode_blocking(
1763 column_infos,
1764 requested_rows,
1765 filter,
1766 projection.column_indices,
1767 projection.schema,
1768 config,
1769 )
1770 }
1771
1772 pub fn read_stream_projected_blocking(
1784 &self,
1785 params: ReadBatchParams,
1786 batch_size: u32,
1787 projection: Option<ReaderProjection>,
1788 filter: FilterExpression,
1789 ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1790 let projection = projection.unwrap_or_else(|| self.core.base_projection.clone());
1791 Self::validate_projection(&projection, &self.metadata)?;
1792 let prepared = PreparedProjection {
1798 column_infos: self.metadata.column_infos.clone(),
1799 decoder_projection: projection.clone(),
1800 };
1801 let read_len = self.core.prepared_read_length(&prepared)?;
1802 let verify_bound = |params: &ReadBatchParams, bound: u64, inclusive: bool| {
1803 if bound > read_len || (bound == read_len && inclusive) {
1804 Err(Error::invalid_input(format!(
1805 "cannot read {params:?} from columns with {read_len} rows"
1806 )))
1807 } else {
1808 Ok(())
1809 }
1810 };
1811 match ¶ms {
1812 ReadBatchParams::Indices(indices) => {
1813 for idx in indices {
1814 match idx {
1815 None => {
1816 return Err(Error::invalid_input("Null value in indices array"));
1817 }
1818 Some(idx) => {
1819 verify_bound(¶ms, idx as u64, true)?;
1820 }
1821 }
1822 }
1823 let indices = indices.iter().map(|idx| idx.unwrap() as u64).collect();
1824 self.take_rows_blocking(indices, batch_size, projection, filter)
1825 }
1826 ReadBatchParams::Range(range) => {
1827 verify_bound(¶ms, range.end as u64, false)?;
1828 self.read_range_blocking(
1829 range.start as u64..range.end as u64,
1830 batch_size,
1831 projection,
1832 filter,
1833 )
1834 }
1835 ReadBatchParams::Ranges(ranges) => {
1836 let mut ranges_u64 = Vec::with_capacity(ranges.len());
1837 for range in ranges.as_ref() {
1838 verify_bound(¶ms, range.end, false)?;
1839 ranges_u64.push(range.start..range.end);
1840 }
1841 self.read_ranges_blocking(ranges_u64, batch_size, projection, filter)
1842 }
1843 ReadBatchParams::RangeFrom(range) => {
1844 verify_bound(¶ms, range.start as u64, true)?;
1845 self.read_range_blocking(
1846 range.start as u64..read_len,
1847 batch_size,
1848 projection,
1849 filter,
1850 )
1851 }
1852 ReadBatchParams::RangeTo(range) => {
1853 verify_bound(¶ms, range.end as u64, false)?;
1854 self.read_range_blocking(0..range.end as u64, batch_size, projection, filter)
1855 }
1856 ReadBatchParams::RangeFull => {
1857 self.read_range_blocking(0..read_len, batch_size, projection, filter)
1858 }
1859 }
1860 }
1861
1862 pub async fn read_stream(
1874 &self,
1875 params: ReadBatchParams,
1876 batch_size: u32,
1877 batch_readahead: u32,
1878 filter: FilterExpression,
1879 ) -> Result<Pin<Box<dyn RecordBatchStream>>> {
1880 self.read_stream_projected(
1881 params,
1882 batch_size,
1883 batch_readahead,
1884 self.core.base_projection.clone(),
1885 filter,
1886 )
1887 .await
1888 }
1889
1890 pub fn schema(&self) -> &Arc<Schema> {
1891 self.core.schema()
1892 }
1893}
1894
1895impl FileMetadataProvider {
1896 fn version(&self) -> LanceFileVersion {
1897 match self {
1898 Self::Full(metadata) => metadata.version(),
1899 Self::Indexed(metadata_index) => metadata_index.version,
1900 }
1901 }
1902
1903 fn num_rows(&self) -> u64 {
1904 match self {
1905 Self::Full(metadata) => metadata.num_rows,
1906 Self::Indexed(metadata_index) => metadata_index.num_rows,
1907 }
1908 }
1909
1910 fn schema(&self) -> &Arc<Schema> {
1911 match self {
1912 Self::Full(metadata) => &metadata.file_schema,
1913 Self::Indexed(metadata_index) => &metadata_index.file_schema,
1914 }
1915 }
1916
1917 fn file_buffers(&self) -> &Vec<BufferDescriptor> {
1918 match self {
1919 Self::Full(metadata) => &metadata.file_buffers,
1920 Self::Indexed(metadata_index) => &metadata_index.file_buffers,
1921 }
1922 }
1923
1924 fn retained_global_buffers(&self) -> &BTreeMap<u32, Bytes> {
1925 match self {
1926 Self::Full(metadata) => &metadata.retained_global_buffers,
1927 Self::Indexed(metadata_index) => &metadata_index.retained_global_buffers,
1928 }
1929 }
1930
1931 fn file_statistics(&self) -> Option<FileStatistics> {
1932 let metadata = match self {
1933 Self::Full(metadata) => metadata,
1934 Self::Indexed(_) => return None,
1935 };
1936 Some(FileReader::statistics_from_column_metadata(
1937 &metadata.column_metadatas,
1938 ))
1939 }
1940
1941 fn supports_indexed_projection(
1942 projection: &ReaderProjection,
1943 version: LanceFileVersion,
1944 ) -> bool {
1945 if version < LanceFileVersion::V2_1 || projection.schema.fields.is_empty() {
1946 return false;
1947 }
1948
1949 projection
1950 .schema
1951 .fields
1952 .iter()
1953 .try_fold(0usize, |count, field| {
1954 count.checked_add(indexed_projection_column_count(field)?)
1955 })
1956 == Some(projection.column_indices.len())
1957 }
1958
1959 fn validate_indexed_projection(
1960 projection: &ReaderProjection,
1961 metadata_index: &FileMetadataIndex,
1962 ) -> Result<()> {
1963 if projection.schema.fields.is_empty() {
1964 return Err(Error::invalid_input(
1965 "Attempt to read zero columns from the file, at least one column must be specified"
1966 .to_string(),
1967 ));
1968 }
1969 let mut column_indices_seen = BTreeSet::new();
1970 for column_index in &projection.column_indices {
1971 if !column_indices_seen.insert(*column_index) {
1972 return Err(Error::invalid_input(format!(
1973 "The projection specified the column index {} more than once",
1974 column_index
1975 )));
1976 }
1977 if *column_index >= metadata_index.num_columns {
1978 return Err(Error::invalid_input(format!(
1979 "The projection specified the column index {} but there are only {} columns in the file",
1980 column_index, metadata_index.num_columns
1981 )));
1982 }
1983 }
1984 if !Self::supports_indexed_projection(projection, metadata_index.version) {
1985 return Err(Error::not_supported(format!(
1986 "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",
1987 metadata_index.version,
1988 projection.schema.fields.len(),
1989 projection.column_indices.len()
1990 )));
1991 }
1992 Ok(())
1993 }
1994
1995 fn validate_projection(&self, projection: &ReaderProjection) -> Result<()> {
1996 match self {
1997 Self::Full(metadata) => FileReader::validate_projection(projection, metadata),
1998 Self::Indexed(metadata_index) => {
1999 Self::validate_indexed_projection(projection, metadata_index)
2000 }
2001 }
2002 }
2003
2004 fn column_metadata_range(
2005 metadata_index: &FileMetadataIndex,
2006 column_index: u32,
2007 ) -> Result<Range<u64>> {
2008 let (position, length) = metadata_index
2009 .column_metadata_offsets
2010 .get(column_index as usize)
2011 .copied()
2012 .ok_or_else(|| {
2013 Error::invalid_input(format!(
2014 "The projection specified the column index {} but there are only {} columns in the file",
2015 column_index, metadata_index.num_columns
2016 ))
2017 })?;
2018 let end = position.checked_add(length).ok_or_else(|| {
2019 Error::invalid_input(format!(
2020 "column metadata range overflows for column index {}, position={}, length={}",
2021 column_index, position, length
2022 ))
2023 })?;
2024 Ok(position..end)
2025 }
2026
2027 async fn load_indexed_column_infos(
2028 metadata_index: &FileMetadataIndex,
2029 io: &Arc<dyn EncodingsIo>,
2030 cache: &Arc<LanceCache>,
2031 column_indices: &[u32],
2032 ) -> Result<Vec<Arc<ColumnInfo>>> {
2033 let mut column_infos = vec![None; column_indices.len()];
2034 let mut missing_columns = Vec::new();
2035
2036 for (result_index, column_index) in column_indices.iter().copied().enumerate() {
2037 let cache_key = ColumnMetadataCacheKey { column_index };
2038 if let Some(cached) = cache.get_with_key(&cache_key).await {
2039 column_infos[result_index] = Some(cached.column_info.clone());
2040 } else {
2041 let range = Self::column_metadata_range(metadata_index, column_index)?;
2042 missing_columns.push((result_index, column_index, range));
2043 }
2044 }
2045
2046 missing_columns.sort_by_key(|(_, _, range)| range.start);
2047 if !missing_columns.is_empty() {
2048 let ranges = missing_columns
2049 .iter()
2050 .map(|(_, _, range)| range.clone())
2051 .collect::<Vec<_>>();
2052 let metadata_bytes = io.submit_request(ranges, 0).await?;
2053 for ((result_index, column_index, _), bytes) in
2054 missing_columns.into_iter().zip(metadata_bytes)
2055 {
2056 let column_metadata = pbfile::ColumnMetadata::decode(bytes)?;
2057 let column_info = FileReader::meta_to_col_info(
2058 column_index,
2059 &column_metadata,
2060 metadata_index.version,
2061 )?;
2062 let cached = Arc::new(CachedColumnMetadata {
2063 column_metadata,
2064 column_info: column_info.clone(),
2065 });
2066 let cache_key = ColumnMetadataCacheKey { column_index };
2067 cache.insert_with_key(&cache_key, cached).await;
2068 column_infos[result_index] = Some(column_info);
2069 }
2070 }
2071
2072 column_infos
2073 .into_iter()
2074 .enumerate()
2075 .map(|(idx, column_info)| {
2076 column_info.ok_or_else(|| {
2077 Error::internal(format!(
2078 "lazy metadata loader did not load requested projection column at position {}",
2079 idx
2080 ))
2081 })
2082 })
2083 .collect()
2084 }
2085
2086 async fn prepare_projection(
2087 &self,
2088 projection: &ReaderProjection,
2089 io: &Arc<dyn EncodingsIo>,
2090 cache: &Arc<LanceCache>,
2091 ) -> Result<PreparedProjection> {
2092 self.validate_projection(projection)?;
2093 match self {
2094 Self::Full(metadata) => Ok(PreparedProjection {
2095 column_infos: metadata.column_infos.clone(),
2096 decoder_projection: projection.clone(),
2097 }),
2098 Self::Indexed(metadata_index) => {
2099 let column_infos = Self::load_indexed_column_infos(
2100 metadata_index,
2101 io,
2102 cache,
2103 &projection.column_indices,
2104 )
2105 .await?;
2106 let decoder_projection = ReaderProjection {
2107 schema: projection.schema.clone(),
2108 column_indices: (0..projection.column_indices.len())
2109 .map(|idx| idx as u32)
2110 .collect(),
2111 };
2112 Ok(PreparedProjection {
2113 column_infos,
2114 decoder_projection,
2115 })
2116 }
2117 }
2118 }
2119}
2120
2121impl FileReadCore {
2122 fn try_new(
2123 scheduler: Arc<dyn EncodingsIo>,
2124 base_projection: Option<ReaderProjection>,
2125 decoder_plugins: Arc<DecoderPlugins>,
2126 metadata_provider: FileMetadataProvider,
2127 cache: Arc<LanceCache>,
2128 options: FileReaderOptions,
2129 ) -> Result<Self> {
2130 if let Some(base_projection) = base_projection.as_ref() {
2131 metadata_provider.validate_projection(base_projection)?;
2132 }
2133 let base_projection = base_projection.unwrap_or(ReaderProjection::from_whole_schema(
2134 metadata_provider.schema().as_ref(),
2135 metadata_provider.version(),
2136 ));
2137 Ok(Self {
2138 scheduler,
2139 base_projection,
2140 metadata_provider,
2141 decoder_plugins,
2142 cache,
2143 options,
2144 })
2145 }
2146
2147 fn with_scheduler(&self, scheduler: Arc<dyn EncodingsIo>) -> Self {
2148 Self {
2149 scheduler,
2150 base_projection: self.base_projection.clone(),
2151 metadata_provider: self.metadata_provider.clone(),
2152 decoder_plugins: self.decoder_plugins.clone(),
2153 cache: self.cache.clone(),
2154 options: self.options.clone(),
2155 }
2156 }
2157
2158 fn version(&self) -> LanceFileVersion {
2159 self.metadata_provider.version()
2160 }
2161
2162 fn num_rows(&self) -> u64 {
2163 self.metadata_provider.num_rows()
2164 }
2165
2166 fn schema(&self) -> &Arc<Schema> {
2167 self.metadata_provider.schema()
2168 }
2169
2170 async fn read_global_buffer(&self, index: u32) -> Result<Bytes> {
2171 let file_buffers = self.metadata_provider.file_buffers();
2172 let buffer_desc = file_buffers.get(index as usize).ok_or_else(|| {
2173 Error::invalid_input(format!(
2174 "request for global buffer at index {} but there were only {} global buffers in the file",
2175 index,
2176 file_buffers.len()
2177 ))
2178 })?;
2179
2180 if let Some(bytes) = self.metadata_provider.retained_global_buffers().get(&index) {
2181 return Ok(bytes.clone());
2182 }
2183
2184 let bytes = self
2185 .scheduler
2186 .submit_request(
2187 vec![buffer_desc.position..buffer_desc.position + buffer_desc.size],
2188 0,
2189 )
2190 .await?;
2191 bytes.into_iter().next().ok_or_else(|| {
2192 Error::internal(format!(
2193 "global buffer read for index {} returned no bytes",
2194 index
2195 ))
2196 })
2197 }
2198
2199 fn prepared_read_length(&self, prepared: &PreparedProjection) -> Result<u64> {
2213 let is_structural = self.version() >= LanceFileVersion::V2_1;
2214 let column_infos = &prepared.column_infos;
2215 let column_len = |column: usize| -> Result<u64> {
2216 let info = column_infos.get(column).ok_or_else(|| {
2217 Error::invalid_input(format!(
2218 "projection references column index {} but only {} columns are available",
2219 column,
2220 column_infos.len()
2221 ))
2222 })?;
2223 info.page_infos.iter().try_fold(0_u64, |rows, page| {
2224 let page_rows = match &page.encoding {
2225 PageEncoding::Structural(layout) => match &layout.layout {
2226 Some(pbenc21::page_layout::Layout::SparseLayout(sparse)) => sparse
2227 .structural_layers
2228 .first()
2229 .and_then(|layer| layer.layer.as_ref())
2230 .map_or(page.num_rows, |layer| match layer {
2231 pbenc21::sparse_structural_layer::Layer::Validity(layer) => {
2232 layer.num_slots
2233 }
2234 pbenc21::sparse_structural_layer::Layer::List(layer) => {
2235 layer.num_slots
2236 }
2237 pbenc21::sparse_structural_layer::Layer::FixedSizeList(layer) => {
2238 layer.num_slots
2239 }
2240 }),
2241 _ => page.num_rows,
2242 },
2243 _ => page.num_rows,
2244 };
2245 rows.checked_add(page_rows).ok_or_else(|| {
2246 Error::invalid_input_source("Column row count overflows u64".into())
2247 })
2248 })
2249 };
2250 let column_indices = &prepared.decoder_projection.column_indices;
2251 let fields = &prepared.decoder_projection.schema.fields;
2252 let mut cursor = 0usize;
2253 let mut field_lengths = Vec::with_capacity(fields.len());
2254 for field in fields {
2255 let rows = validate_field_length(
2256 field,
2257 is_structural,
2258 true,
2259 column_indices,
2260 &mut cursor,
2261 &column_len,
2262 )?;
2263 field_lengths.push((field.name.as_str(), rows));
2264 }
2265 if cursor != column_indices.len() {
2266 return Err(Error::invalid_input(format!(
2267 "projection supplied {} column indices but its fields require {}",
2268 column_indices.len(),
2269 cursor
2270 )));
2271 }
2272 verify_uniform_lengths(&field_lengths)
2273 }
2274
2275 async fn read_range(
2276 &self,
2277 range: Range<u64>,
2278 batch_size: u32,
2279 prepared: PreparedProjection,
2280 filter: FilterExpression,
2281 ) -> Result<BoxStream<'static, ReadBatchTask>> {
2282 FileReader::do_read_range(
2283 prepared.column_infos,
2284 self.scheduler.clone(),
2285 self.cache.clone(),
2286 self.num_rows(),
2287 self.decoder_plugins.clone(),
2288 range,
2289 batch_size,
2290 prepared.decoder_projection,
2291 filter,
2292 self.options.decoder_config.clone(),
2293 self.options.batch_size_bytes,
2294 )
2295 .await
2296 }
2297
2298 async fn take_rows(
2299 &self,
2300 indices: Vec<u64>,
2301 batch_size: u32,
2302 prepared: PreparedProjection,
2303 ) -> Result<BoxStream<'static, ReadBatchTask>> {
2304 FileReader::do_take_rows(
2305 prepared.column_infos,
2306 self.scheduler.clone(),
2307 self.cache.clone(),
2308 self.decoder_plugins.clone(),
2309 indices,
2310 batch_size,
2311 prepared.decoder_projection,
2312 FilterExpression::no_filter(),
2313 self.options.decoder_config.clone(),
2314 self.options.batch_size_bytes,
2315 )
2316 .await
2317 }
2318
2319 async fn read_ranges(
2320 &self,
2321 ranges: Vec<Range<u64>>,
2322 batch_size: u32,
2323 prepared: PreparedProjection,
2324 filter: FilterExpression,
2325 ) -> Result<BoxStream<'static, ReadBatchTask>> {
2326 FileReader::do_read_ranges(
2327 prepared.column_infos,
2328 self.scheduler.clone(),
2329 self.cache.clone(),
2330 self.decoder_plugins.clone(),
2331 ranges,
2332 batch_size,
2333 prepared.decoder_projection,
2334 filter,
2335 self.options.decoder_config.clone(),
2336 self.options.batch_size_bytes,
2337 )
2338 .await
2339 }
2340
2341 async fn read_tasks(
2342 &self,
2343 params: ReadBatchParams,
2344 batch_size: u32,
2345 projection: Option<ReaderProjection>,
2346 filter: FilterExpression,
2347 ) -> Result<Pin<Box<dyn Stream<Item = ReadBatchTask> + Send>>> {
2348 let projection = projection.unwrap_or_else(|| self.base_projection.clone());
2349 let prepared = self
2350 .metadata_provider
2351 .prepare_projection(&projection, &self.scheduler, &self.cache)
2352 .await?;
2353 let read_len = self.prepared_read_length(&prepared)?;
2360 let verify_bound = |params: &ReadBatchParams, bound: u64, inclusive: bool| {
2361 if bound > read_len || (bound == read_len && inclusive) {
2362 Err(Error::invalid_input(format!(
2363 "cannot read {params:?} from columns with {read_len} rows"
2364 )))
2365 } else {
2366 Ok(())
2367 }
2368 };
2369 match ¶ms {
2370 ReadBatchParams::Indices(indices) => {
2371 for idx in indices {
2372 match idx {
2373 None => {
2374 return Err(Error::invalid_input("Null value in indices array"));
2375 }
2376 Some(idx) => {
2377 verify_bound(¶ms, idx as u64, true)?;
2378 }
2379 }
2380 }
2381 let indices = indices.iter().map(|idx| idx.unwrap() as u64).collect();
2382 self.take_rows(indices, batch_size, prepared).await
2383 }
2384 ReadBatchParams::Range(range) => {
2385 verify_bound(¶ms, range.end as u64, false)?;
2386 self.read_range(
2387 range.start as u64..range.end as u64,
2388 batch_size,
2389 prepared,
2390 filter,
2391 )
2392 .await
2393 }
2394 ReadBatchParams::Ranges(ranges) => {
2395 let mut ranges_u64 = Vec::with_capacity(ranges.len());
2396 for range in ranges.as_ref() {
2397 verify_bound(¶ms, range.end, false)?;
2398 ranges_u64.push(range.start..range.end);
2399 }
2400 self.read_ranges(ranges_u64, batch_size, prepared, filter)
2401 .await
2402 }
2403 ReadBatchParams::RangeFrom(range) => {
2404 verify_bound(¶ms, range.start as u64, true)?;
2405 self.read_range(range.start as u64..read_len, batch_size, prepared, filter)
2406 .await
2407 }
2408 ReadBatchParams::RangeTo(range) => {
2409 verify_bound(¶ms, range.end as u64, false)?;
2410 self.read_range(0..range.end as u64, batch_size, prepared, filter)
2411 .await
2412 }
2413 ReadBatchParams::RangeFull => {
2414 self.read_range(0..read_len, batch_size, prepared, filter)
2415 .await
2416 }
2417 }
2418 }
2419}
2420
2421impl ProjectedFileReader {
2422 pub async fn try_open(
2427 scheduler: FileScheduler,
2428 base_projection: Option<ReaderProjection>,
2429 decoder_plugins: Arc<DecoderPlugins>,
2430 cache: &LanceCache,
2431 options: FileReaderOptions,
2432 ) -> Result<Self> {
2433 let base_projection = Self::require_indexed_base_projection(base_projection)?;
2434 let metadata_index = Arc::new(FileReader::read_metadata_index(&scheduler).await?);
2435 let path = scheduler.reader().path().clone();
2436 let encodings_io =
2437 LanceEncodingsIo::new(scheduler).with_read_chunk_size(options.read_chunk_size);
2438 Self::try_open_with_metadata_index(
2439 Arc::new(encodings_io),
2440 path,
2441 Some(base_projection),
2442 decoder_plugins,
2443 metadata_index,
2444 cache,
2445 options,
2446 )
2447 .await
2448 }
2449
2450 pub async fn try_open_with_metadata_index(
2456 scheduler: Arc<dyn EncodingsIo>,
2457 path: Path,
2458 base_projection: Option<ReaderProjection>,
2459 decoder_plugins: Arc<DecoderPlugins>,
2460 metadata_index: Arc<FileMetadataIndex>,
2461 cache: &LanceCache,
2462 options: FileReaderOptions,
2463 ) -> Result<Self> {
2464 let base_projection = Self::require_indexed_base_projection(base_projection)?;
2465 let cache = Arc::new(cache.with_key_prefix(path.as_ref()));
2466 let core = FileReadCore::try_new(
2467 scheduler,
2468 Some(base_projection),
2469 decoder_plugins,
2470 FileMetadataProvider::Indexed(metadata_index),
2471 cache,
2472 options,
2473 )?;
2474 Ok(Self { core })
2475 }
2476
2477 fn require_indexed_base_projection(
2478 base_projection: Option<ReaderProjection>,
2479 ) -> Result<ReaderProjection> {
2480 base_projection.ok_or_else(|| {
2481 Error::invalid_input("ProjectedFileReader requires an explicit base projection")
2482 })
2483 }
2484
2485 pub async fn try_open_with_file_metadata(
2486 scheduler: Arc<dyn EncodingsIo>,
2487 path: Path,
2488 base_projection: Option<ReaderProjection>,
2489 decoder_plugins: Arc<DecoderPlugins>,
2490 file_metadata: Arc<CachedFileMetadata>,
2491 cache: &LanceCache,
2492 options: FileReaderOptions,
2493 ) -> Result<Self> {
2494 let cache = Arc::new(cache.with_key_prefix(path.as_ref()));
2495 let core = FileReadCore::try_new(
2496 scheduler,
2497 base_projection,
2498 decoder_plugins,
2499 FileMetadataProvider::Full(file_metadata),
2500 cache,
2501 options,
2502 )?;
2503 Ok(Self { core })
2504 }
2505
2506 pub fn supports_projection(projection: &ReaderProjection, version: LanceFileVersion) -> bool {
2508 FileMetadataProvider::supports_indexed_projection(projection, version)
2509 }
2510
2511 pub fn with_scheduler(&self, scheduler: Arc<dyn EncodingsIo>) -> Self {
2513 Self {
2514 core: self.core.with_scheduler(scheduler),
2515 }
2516 }
2517
2518 pub fn version(&self) -> LanceFileVersion {
2520 self.core.version()
2521 }
2522
2523 pub fn num_rows(&self) -> u64 {
2525 self.core.num_rows()
2526 }
2527
2528 pub fn schema(&self) -> &Arc<Schema> {
2530 self.core.schema()
2531 }
2532
2533 pub fn file_statistics(&self) -> Option<FileStatistics> {
2535 self.core.metadata_provider.file_statistics()
2536 }
2537
2538 #[cfg(test)]
2539 fn metadata_index(&self) -> Option<&Arc<FileMetadataIndex>> {
2540 match &self.core.metadata_provider {
2541 FileMetadataProvider::Indexed(metadata_index) => Some(metadata_index),
2542 FileMetadataProvider::Full(_) => None,
2543 }
2544 }
2545
2546 pub async fn read_global_buffer(&self, index: u32) -> Result<Bytes> {
2548 self.core.read_global_buffer(index).await
2549 }
2550
2551 pub async fn read_tasks(
2553 &self,
2554 params: ReadBatchParams,
2555 batch_size: u32,
2556 projection: Option<ReaderProjection>,
2557 filter: FilterExpression,
2558 ) -> Result<Pin<Box<dyn Stream<Item = ReadBatchTask> + Send>>> {
2559 self.core
2560 .read_tasks(params, batch_size, projection, filter)
2561 .await
2562 }
2563}
2564
2565pub fn describe_encoding(page: &pbfile::column_metadata::Page) -> String {
2567 if let Some(encoding) = &page.encoding {
2568 if let Some(style) = &encoding.location {
2569 match style {
2570 pbfile::encoding::Location::Indirect(indirect) => {
2571 format!(
2572 "IndirectEncoding(pos={},size={})",
2573 indirect.buffer_location, indirect.buffer_length
2574 )
2575 }
2576 pbfile::encoding::Location::Direct(direct) => {
2577 let encoding_any =
2578 prost_types::Any::decode(Bytes::from(direct.encoding.clone()))
2579 .expect("failed to deserialize encoding as protobuf");
2580 if encoding_any.type_url == "/lance.encodings.ArrayEncoding" {
2581 let encoding = encoding_any.to_msg::<pbenc::ArrayEncoding>();
2582 match encoding {
2583 Ok(encoding) => {
2584 format!("{:#?}", encoding)
2585 }
2586 Err(err) => {
2587 format!("Unsupported(decode_err={})", err)
2588 }
2589 }
2590 } else if encoding_any.type_url == "/lance.encodings21.PageLayout" {
2591 let encoding = encoding_any.to_msg::<pbenc21::PageLayout>();
2592 match encoding {
2593 Ok(encoding) => {
2594 format!("{:#?}", encoding)
2595 }
2596 Err(err) => {
2597 format!("Unsupported(decode_err={})", err)
2598 }
2599 }
2600 } else {
2601 format!("Unrecognized(type_url={})", encoding_any.type_url)
2602 }
2603 }
2604 pbfile::encoding::Location::None(_) => "NoEncodingDescription".to_string(),
2605 }
2606 } else {
2607 "MISSING STYLE".to_string()
2608 }
2609 } else {
2610 "MISSING".to_string()
2611 }
2612}
2613
2614pub trait EncodedBatchReaderExt {
2615 fn try_from_mini_lance(
2616 bytes: Bytes,
2617 schema: &Schema,
2618 version: LanceFileVersion,
2619 ) -> Result<Self>
2620 where
2621 Self: Sized;
2622 fn try_from_self_described_lance(bytes: Bytes) -> Result<Self>
2623 where
2624 Self: Sized;
2625}
2626
2627impl EncodedBatchReaderExt for EncodedBatch {
2628 fn try_from_mini_lance(
2629 bytes: Bytes,
2630 schema: &Schema,
2631 file_version: LanceFileVersion,
2632 ) -> Result<Self>
2633 where
2634 Self: Sized,
2635 {
2636 let projection = ReaderProjection::from_whole_schema(schema, file_version);
2637 let footer = FileReader::decode_footer(&bytes)?;
2638
2639 let column_metadata_start = footer.column_meta_start as usize;
2642 let column_metadata_end = footer.global_buff_offsets_start as usize;
2643 let column_metadata_bytes = bytes.slice(column_metadata_start..column_metadata_end);
2644 let column_metadatas =
2645 FileReader::read_all_column_metadata(column_metadata_bytes, &footer)?;
2646
2647 let file_version: LanceFileVersion =
2648 ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)?
2649 .into();
2650
2651 let page_table = FileReader::meta_to_col_infos(&column_metadatas, file_version)?;
2652
2653 Ok(Self {
2654 data: bytes,
2655 num_rows: page_table
2656 .first()
2657 .map(|col| col.page_infos.iter().map(|page| page.num_rows).sum::<u64>())
2658 .unwrap_or(0),
2659 page_table,
2660 top_level_columns: projection.column_indices,
2661 schema: Arc::new(schema.clone()),
2662 })
2663 }
2664
2665 fn try_from_self_described_lance(bytes: Bytes) -> Result<Self>
2666 where
2667 Self: Sized,
2668 {
2669 let footer = FileReader::decode_footer(&bytes)?;
2670 let file_version: LanceFileVersion =
2671 ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)?
2672 .into();
2673
2674 let gbo_table = FileReader::do_decode_gbo_table(
2675 &bytes.slice(footer.global_buff_offsets_start as usize..),
2676 &footer,
2677 file_version,
2678 )?;
2679 if gbo_table.is_empty() {
2680 return Err(Error::internal(
2681 "File did not contain any global buffers, schema expected".to_string(),
2682 ));
2683 }
2684 let schema_start = gbo_table[0].position as usize;
2685 let schema_size = gbo_table[0].size as usize;
2686
2687 let schema_bytes = bytes.slice(schema_start..(schema_start + schema_size));
2688 let (_, schema) = FileReader::decode_schema(schema_bytes)?;
2689 let projection = ReaderProjection::from_whole_schema(&schema, file_version);
2690
2691 let column_metadata_start = footer.column_meta_start as usize;
2694 let column_metadata_end = footer.global_buff_offsets_start as usize;
2695 let column_metadata_bytes = bytes.slice(column_metadata_start..column_metadata_end);
2696 let column_metadatas =
2697 FileReader::read_all_column_metadata(column_metadata_bytes, &footer)?;
2698
2699 let page_table = FileReader::meta_to_col_infos(&column_metadatas, file_version)?;
2700
2701 Ok(Self {
2702 data: bytes,
2703 num_rows: page_table
2704 .first()
2705 .map(|col| col.page_infos.iter().map(|page| page.num_rows).sum::<u64>())
2706 .unwrap_or(0),
2707 page_table,
2708 top_level_columns: projection.column_indices,
2709 schema: Arc::new(schema),
2710 })
2711 }
2712}
2713
2714#[cfg(test)]
2715mod tests {
2716 use std::{
2717 collections::{BTreeMap, HashMap},
2718 pin::Pin,
2719 sync::Arc,
2720 };
2721
2722 use arrow_array::{
2723 Int32Array, ListArray, RecordBatch, RecordBatchIterator, UInt32Array,
2724 types::{Float64Type, Int32Type},
2725 };
2726 use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
2727 use arrow_schema::{DataType, Field, Fields, Schema as ArrowSchema};
2728 use bytes::Bytes;
2729 use futures::{StreamExt, prelude::stream::TryStreamExt};
2730 use lance_arrow::{BLOB_META_KEY, RecordBatchExt};
2731 use lance_core::{ArrowResult, datatypes::Schema};
2732 use lance_datagen::{ArrayGeneratorExt, BatchCount, ByteCount, RowCount, array, gen_batch};
2733 use lance_encoding::{
2734 constants::{STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_SPARSE},
2735 decoder::{
2736 DecodeBatchScheduler, DecoderPlugins, EncodedBatchLayout, FilterExpression,
2737 PageEncoding, ReadBatchTask, decode_batch,
2738 },
2739 encoder::{EncodedBatch, EncodingOptions, encode_batch},
2740 format::pb21,
2741 version::LanceFileVersion,
2742 };
2743 use lance_io::{stream::RecordBatchStream, utils::CachedFileSize};
2744 use log::debug;
2745 use rstest::rstest;
2746 use tokio::sync::mpsc;
2747
2748 use crate::reader::{
2749 EncodedBatchReaderExt, FileReader, FileReaderOptions, ProjectedFileReader,
2750 ReaderProjection, validate_field_length, verify_uniform_lengths,
2751 };
2752 use crate::testing::{FsFixture, WrittenFile, test_cache, write_lance_file};
2753 use crate::writer::{EncodedBatchWriteExt, FileWriter, FileWriterOptions};
2754 use lance_encoding::decoder::DecoderConfig;
2755
2756 fn footer_version(bytes: &[u8]) -> (u16, u16) {
2757 let version_start = bytes.len() - 8;
2758 (
2759 u16::from_le_bytes([bytes[version_start], bytes[version_start + 1]]),
2760 u16::from_le_bytes([bytes[version_start + 2], bytes[version_start + 3]]),
2761 )
2762 }
2763
2764 #[tokio::test]
2765 async fn sparse_file_writer_reader_scan_range_and_take_roundtrip() {
2766 let fs = FsFixture::default();
2767 let sparse_metadata = HashMap::from([(
2768 STRUCTURAL_ENCODING_META_KEY.to_string(),
2769 STRUCTURAL_ENCODING_SPARSE.to_string(),
2770 )]);
2771 let value_field =
2772 Field::new("values", DataType::Int32, true).with_metadata(sparse_metadata.clone());
2773 let item_field = Arc::new(Field::new("item", DataType::Int32, true));
2774 let list_field = Field::new("items", DataType::List(item_field.clone()), true)
2775 .with_metadata(sparse_metadata);
2776 let arrow_schema = Arc::new(ArrowSchema::new(vec![value_field, list_field]));
2777 let list = ListArray::try_new(
2778 item_field,
2779 OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 2, 2, 2, 3, 3, 5])),
2780 Arc::new(Int32Array::from(vec![
2781 Some(1),
2782 None,
2783 Some(3),
2784 Some(4),
2785 Some(5),
2786 ])),
2787 Some(NullBuffer::from(vec![true, false, true, true, true, true])),
2788 )
2789 .unwrap();
2790 let batch = RecordBatch::try_new(
2791 arrow_schema.clone(),
2792 vec![
2793 Arc::new(Int32Array::from(vec![
2794 Some(10),
2795 None,
2796 Some(30),
2797 Some(40),
2798 None,
2799 Some(60),
2800 ])),
2801 Arc::new(list),
2802 ],
2803 )
2804 .unwrap();
2805 let input = RecordBatchIterator::new(vec![Ok(batch.clone())], arrow_schema);
2806 write_lance_file(
2807 input,
2808 &fs,
2809 FileWriterOptions {
2810 format_version: Some(LanceFileVersion::V2_3),
2811 ..Default::default()
2812 },
2813 )
2814 .await;
2815
2816 let file_scheduler = fs
2817 .scheduler
2818 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
2819 .await
2820 .unwrap();
2821 let file_reader = FileReader::try_open(
2822 file_scheduler,
2823 None,
2824 Arc::<DecoderPlugins>::default(),
2825 &test_cache(),
2826 FileReaderOptions::default(),
2827 )
2828 .await
2829 .unwrap();
2830 assert_eq!(file_reader.metadata.column_infos.len(), 2);
2831 assert!(
2832 file_reader
2833 .metadata
2834 .column_infos
2835 .iter()
2836 .flat_map(|column| column.page_infos.iter())
2837 .all(|page| {
2838 matches!(
2839 &page.encoding,
2840 PageEncoding::Structural(layout)
2841 if matches!(
2842 layout.layout,
2843 Some(pb21::page_layout::Layout::SparseLayout(_))
2844 )
2845 )
2846 })
2847 );
2848
2849 let scan = file_reader
2850 .read_stream(
2851 lance_io::ReadBatchParams::RangeFull,
2852 1024,
2853 1,
2854 FilterExpression::no_filter(),
2855 )
2856 .await
2857 .unwrap()
2858 .try_collect::<Vec<_>>()
2859 .await
2860 .unwrap();
2861 assert_eq!(scan, vec![batch.clone()]);
2862
2863 let range = file_reader
2864 .read_stream(
2865 lance_io::ReadBatchParams::Range(1..5),
2866 1024,
2867 1,
2868 FilterExpression::no_filter(),
2869 )
2870 .await
2871 .unwrap()
2872 .try_collect::<Vec<_>>()
2873 .await
2874 .unwrap();
2875 assert_eq!(range, vec![batch.slice(1, 4)]);
2876
2877 let indices = UInt32Array::from(vec![0, 3, 5]);
2878 let take = file_reader
2879 .read_stream(
2880 lance_io::ReadBatchParams::Indices(indices.clone()),
2881 1024,
2882 1,
2883 FilterExpression::no_filter(),
2884 )
2885 .await
2886 .unwrap()
2887 .try_collect::<Vec<_>>()
2888 .await
2889 .unwrap();
2890 assert_eq!(take, vec![batch.take(&indices).unwrap()]);
2891 }
2892
2893 async fn create_some_file(fs: &FsFixture, version: LanceFileVersion) -> WrittenFile {
2894 let location_type = DataType::Struct(Fields::from(vec![
2895 Field::new("x", DataType::Float64, true),
2896 Field::new("y", DataType::Float64, true),
2897 ]));
2898 let categories_type = DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)));
2899
2900 let mut reader = gen_batch()
2901 .col("score", array::rand::<Float64Type>())
2902 .col("location", array::rand_type(&location_type))
2903 .col("categories", array::rand_type(&categories_type))
2904 .col("binary", array::rand_type(&DataType::Binary));
2905 if version <= LanceFileVersion::V2_0 {
2906 reader = reader.col("large_bin", array::rand_type(&DataType::LargeBinary));
2907 }
2908 let reader = reader.into_reader_rows(RowCount::from(1000), BatchCount::from(100));
2909
2910 write_lance_file(
2911 reader,
2912 fs,
2913 FileWriterOptions {
2914 format_version: Some(version),
2915 ..Default::default()
2916 },
2917 )
2918 .await
2919 }
2920
2921 async fn create_wide_direct_file(fs: &FsFixture, num_columns: usize) -> WrittenFile {
2922 let mut reader = gen_batch();
2923 for column_idx in 0..num_columns {
2924 reader = reader.col(format!("c{column_idx}"), array::step::<Int32Type>());
2925 }
2926 let reader = reader.into_reader_rows(RowCount::from(1000), BatchCount::from(100));
2927
2928 write_lance_file(
2929 reader,
2930 fs,
2931 FileWriterOptions {
2932 format_version: Some(LanceFileVersion::V2_1),
2933 ..Default::default()
2934 },
2935 )
2936 .await
2937 }
2938
2939 async fn create_wide_fixed_size_list_file(fs: &FsFixture, num_columns: usize) -> WrittenFile {
2940 let data_type =
2941 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4);
2942 let mut reader = gen_batch();
2943 for column_idx in 0..num_columns {
2944 reader = reader.col(
2945 format!("c{column_idx}"),
2946 array::rand_type(&data_type).with_random_nulls(0.1),
2947 );
2948 }
2949 let reader = reader.into_reader_rows(RowCount::from(64), BatchCount::from(4));
2950
2951 write_lance_file(
2952 reader,
2953 fs,
2954 FileWriterOptions {
2955 format_version: Some(LanceFileVersion::V2_1),
2956 ..Default::default()
2957 },
2958 )
2959 .await
2960 }
2961
2962 async fn create_wide_structural_file(fs: &FsFixture, num_groups: usize) -> WrittenFile {
2963 let struct_type = DataType::Struct(Fields::from(vec![
2964 Field::new("x", DataType::Int32, true),
2965 Field::new("y", DataType::Int32, true),
2966 ]));
2967 let list_type = DataType::List(Arc::new(Field::new("item", DataType::Int32, true)));
2968 let mut reader = gen_batch();
2969 for group_idx in 0..num_groups {
2970 reader = reader
2971 .col(
2972 format!("s{group_idx}"),
2973 array::rand_type(&struct_type).with_random_nulls(0.5),
2974 )
2975 .col(
2976 format!("l{group_idx}"),
2977 array::rand_type(&list_type).with_random_nulls(0.5),
2978 );
2979 }
2980 let reader = reader.into_reader_rows(RowCount::from(64), BatchCount::from(4));
2981
2982 write_lance_file(
2983 reader,
2984 fs,
2985 FileWriterOptions {
2986 format_version: Some(LanceFileVersion::V2_1),
2987 ..Default::default()
2988 },
2989 )
2990 .await
2991 }
2992
2993 type Transformer = Box<dyn Fn(&RecordBatch) -> RecordBatch>;
2994
2995 async fn verify_expected(
2996 expected: &[RecordBatch],
2997 mut actual: Pin<Box<dyn RecordBatchStream>>,
2998 read_size: u32,
2999 transform: Option<Transformer>,
3000 ) {
3001 let mut remaining = expected.iter().map(|batch| batch.num_rows()).sum::<usize>() as u32;
3002 let mut expected_iter = expected.iter().map(|batch| {
3003 if let Some(transform) = &transform {
3004 transform(batch)
3005 } else {
3006 batch.clone()
3007 }
3008 });
3009 let mut next_expected = expected_iter.next().unwrap().clone();
3010 while let Some(actual) = actual.next().await {
3011 let mut actual = actual.unwrap();
3012 let mut rows_to_verify = actual.num_rows() as u32;
3013 let expected_length = remaining.min(read_size);
3014 assert_eq!(expected_length, rows_to_verify);
3015
3016 while rows_to_verify > 0 {
3017 let next_slice_len = (next_expected.num_rows() as u32).min(rows_to_verify);
3018 assert_eq!(
3019 next_expected.slice(0, next_slice_len as usize),
3020 actual.slice(0, next_slice_len as usize)
3021 );
3022 remaining -= next_slice_len;
3023 rows_to_verify -= next_slice_len;
3024 if remaining > 0 {
3025 if next_slice_len == next_expected.num_rows() as u32 {
3026 next_expected = expected_iter.next().unwrap().clone();
3027 } else {
3028 next_expected = next_expected.slice(
3029 next_slice_len as usize,
3030 next_expected.num_rows() - next_slice_len as usize,
3031 );
3032 }
3033 }
3034 if rows_to_verify > 0 {
3035 actual = actual.slice(
3036 next_slice_len as usize,
3037 actual.num_rows() - next_slice_len as usize,
3038 );
3039 }
3040 }
3041 }
3042 assert_eq!(remaining, 0);
3043 }
3044
3045 async fn collect_read_tasks(
3046 tasks: Pin<Box<dyn futures::Stream<Item = ReadBatchTask> + Send>>,
3047 readahead: usize,
3048 ) -> Vec<RecordBatch> {
3049 tasks
3050 .map(|task| task.task)
3051 .buffered(readahead)
3052 .try_collect::<Vec<_>>()
3053 .await
3054 .unwrap()
3055 }
3056
3057 async fn read_file_with_mutated_bytes(
3061 version: LanceFileVersion,
3062 batch: RecordBatch,
3063 pattern: &[u8],
3064 patch_offset: usize,
3065 patch: &[u8],
3066 ) -> lance_core::Result<Vec<RecordBatch>> {
3067 let fs = FsFixture::default();
3068 let schema = batch.schema();
3069 write_lance_file(
3070 RecordBatchIterator::new(vec![Ok(batch)], schema),
3071 &fs,
3072 FileWriterOptions {
3073 format_version: Some(version),
3074 ..Default::default()
3075 },
3076 )
3077 .await;
3078
3079 let mut bytes = fs
3080 .object_store
3081 .read_one_all(&fs.tmp_path)
3082 .await
3083 .unwrap()
3084 .to_vec();
3085 let matches = bytes
3086 .windows(pattern.len())
3087 .enumerate()
3088 .filter_map(|(position, window)| (window == pattern).then_some(position))
3089 .collect::<Vec<_>>();
3090 assert_eq!(
3091 matches.len(),
3092 1,
3093 "expected the byte pattern to appear exactly once in the file"
3094 );
3095 let patch_start = matches[0] + patch_offset;
3096 bytes[patch_start..patch_start + patch.len()].copy_from_slice(patch);
3097 fs.object_store.put(&fs.tmp_path, &bytes).await.unwrap();
3098
3099 let file_scheduler = fs
3100 .scheduler
3101 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3102 .await
3103 .unwrap();
3104 let file_reader = FileReader::try_open(
3105 file_scheduler,
3106 None,
3107 Arc::<DecoderPlugins>::default(),
3108 &test_cache(),
3109 FileReaderOptions::default(),
3110 )
3111 .await
3112 .unwrap();
3113 file_reader
3114 .read_stream(
3115 lance_io::ReadBatchParams::RangeFull,
3116 1024,
3117 16,
3118 FilterExpression::no_filter(),
3119 )
3120 .await?
3121 .try_collect::<Vec<_>>()
3122 .await
3123 }
3124
3125 #[rstest]
3135 #[tokio::test]
3136 async fn test_default_reader_rejects_out_of_bounds_variable_width_offsets(
3137 #[values(LanceFileVersion::V2_1, LanceFileVersion::V2_2, LanceFileVersion::V2_3)]
3138 version: LanceFileVersion,
3139 ) {
3140 use arrow_array::{Array, DictionaryArray, Int32Array, StringArray};
3141
3142 let values = StringArray::from(vec!["alpha", "beta", "gamma"]);
3143 let indices = Int32Array::from((0..300).map(|i| i % 3).collect::<Vec<i32>>());
3144 let dictionary = DictionaryArray::new(indices, Arc::new(values));
3145 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
3146 "category",
3147 dictionary.data_type().clone(),
3148 false,
3149 )]));
3150 let batch = RecordBatch::try_new(arrow_schema, vec![Arc::new(dictionary)]).unwrap();
3151
3152 let offsets_tail_pattern = [5_i32, 9, 14]
3159 .iter()
3160 .flat_map(|value| value.to_le_bytes())
3161 .collect::<Vec<u8>>();
3162 let error = read_file_with_mutated_bytes(
3163 version,
3164 batch,
3165 &offsets_tail_pattern,
3166 8,
3167 &100_000_i32.to_le_bytes(),
3168 )
3169 .await
3170 .expect_err("out-of-bounds offsets must fail the read");
3171 assert!(
3172 matches!(error, lance_core::Error::CorruptFile { .. }),
3173 "expected CorruptFile, got: {error}"
3174 );
3175 assert!(
3176 error.to_string().contains("out of bounds"),
3177 "unexpected message: {error}"
3178 );
3179 }
3180
3181 #[rstest]
3186 #[tokio::test]
3187 async fn test_default_reader_rejects_out_of_bounds_miniblock_offsets(
3188 #[values(LanceFileVersion::V2_1, LanceFileVersion::V2_2, LanceFileVersion::V2_3)]
3189 version: LanceFileVersion,
3190 ) {
3191 use arrow_array::StringArray;
3192
3193 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
3194 "strings",
3195 DataType::Utf8,
3196 false,
3197 )]));
3198 let batch = RecordBatch::try_new(
3199 arrow_schema,
3200 vec![Arc::new(StringArray::from(vec!["alpha", "beta", "gamma"]))],
3201 )
3202 .unwrap();
3203
3204 let chunk_offsets_pattern = [16_i32, 21, 25, 30]
3209 .iter()
3210 .flat_map(|value| value.to_le_bytes())
3211 .collect::<Vec<u8>>();
3212 let error = read_file_with_mutated_bytes(
3213 version,
3214 batch,
3215 &chunk_offsets_pattern,
3216 12,
3217 &100_000_i32.to_le_bytes(),
3218 )
3219 .await
3220 .expect_err("an out-of-bounds chunk offset must fail the read");
3221 assert!(
3222 matches!(error, lance_core::Error::CorruptFile { .. }),
3223 "expected CorruptFile, got: {error}"
3224 );
3225 assert!(
3226 error.to_string().contains("out of bounds"),
3227 "unexpected message: {error}"
3228 );
3229 }
3230
3231 #[tokio::test]
3232 async fn test_round_trip() {
3233 let fs = FsFixture::default();
3234
3235 let WrittenFile { data, .. } = create_some_file(&fs, LanceFileVersion::V2_0).await;
3236
3237 let file_size = fs.object_store.size(&fs.tmp_path).await.unwrap() as usize;
3238 let footer = fs
3239 .object_store
3240 .open(&fs.tmp_path)
3241 .await
3242 .unwrap()
3243 .get_range(file_size - 8..file_size)
3244 .await
3245 .unwrap();
3246 assert_eq!(footer_version(&footer), (0, 3));
3247 assert_eq!(
3248 crate::determine_file_version(&fs.object_store, &fs.tmp_path, Some(file_size))
3249 .await
3250 .unwrap(),
3251 LanceFileVersion::V2_0
3252 );
3253
3254 for read_size in [32, 1024, 1024 * 1024] {
3255 let file_scheduler = fs
3256 .scheduler
3257 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3258 .await
3259 .unwrap();
3260 let file_reader = FileReader::try_open(
3261 file_scheduler,
3262 None,
3263 Arc::<DecoderPlugins>::default(),
3264 &test_cache(),
3265 FileReaderOptions::default(),
3266 )
3267 .await
3268 .unwrap();
3269
3270 let schema = file_reader.schema();
3271 assert_eq!(schema.metadata.get("foo").unwrap(), "bar");
3272
3273 let batch_stream = file_reader
3274 .read_stream(
3275 lance_io::ReadBatchParams::RangeFull,
3276 read_size,
3277 16,
3278 FilterExpression::no_filter(),
3279 )
3280 .await
3281 .unwrap();
3282
3283 verify_expected(&data, batch_stream, read_size, None).await;
3284 }
3285 }
3286
3287 #[rstest]
3288 #[test_log::test(tokio::test)]
3289 async fn test_encoded_batch_round_trip(
3290 #[values(LanceFileVersion::V2_0)] version: LanceFileVersion,
3292 ) {
3293 let data = gen_batch()
3294 .col("x", array::rand::<Int32Type>())
3295 .col("y", array::rand_utf8(ByteCount::from(16), false))
3296 .into_batch_rows(RowCount::from(10000))
3297 .unwrap();
3298
3299 let lance_schema = Arc::new(Schema::try_from(data.schema().as_ref()).unwrap());
3300
3301 let encoding_options = EncodingOptions {
3302 cache_bytes_per_column: 4096,
3303 max_page_bytes: 32 * 1024 * 1024,
3304 keep_original_array: true,
3305 buffer_alignment: 64,
3306 };
3307
3308 let encoding_strategy = crate::versions::v2_0::encoding_strategy();
3309
3310 let encoded_batch = encode_batch(
3311 &data,
3312 lance_schema.clone(),
3313 encoding_strategy.as_ref(),
3314 &encoding_options,
3315 )
3316 .await
3317 .unwrap();
3318
3319 let bytes = encoded_batch.try_to_self_described_lance(version).unwrap();
3321 assert_eq!(footer_version(&bytes), (2, 0));
3322
3323 let decoded_batch = EncodedBatch::try_from_self_described_lance(bytes).unwrap();
3324
3325 let decoded = decode_batch(
3326 &decoded_batch,
3327 &FilterExpression::no_filter(),
3328 Arc::<DecoderPlugins>::default(),
3329 false,
3330 EncodedBatchLayout::Array,
3331 None,
3332 )
3333 .await
3334 .unwrap();
3335
3336 assert_eq!(data, decoded);
3337
3338 let bytes = encoded_batch.try_to_mini_lance(version).unwrap();
3340 assert_eq!(footer_version(&bytes), (2, 0));
3341 let decoded_batch =
3342 EncodedBatch::try_from_mini_lance(bytes, lance_schema.as_ref(), LanceFileVersion::V2_0)
3343 .unwrap();
3344 let decoded = decode_batch(
3345 &decoded_batch,
3346 &FilterExpression::no_filter(),
3347 Arc::<DecoderPlugins>::default(),
3348 false,
3349 EncodedBatchLayout::Array,
3350 None,
3351 )
3352 .await
3353 .unwrap();
3354
3355 assert_eq!(data, decoded);
3356 }
3357
3358 #[rstest]
3359 #[test_log::test(tokio::test)]
3360 async fn test_projection(
3361 #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1, LanceFileVersion::V2_2)]
3362 version: LanceFileVersion,
3363 ) {
3364 let fs = FsFixture::default();
3365
3366 let written_file = create_some_file(&fs, version).await;
3367 let file_scheduler = fs
3368 .scheduler
3369 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3370 .await
3371 .unwrap();
3372
3373 let field_id_mapping = written_file
3374 .field_id_mapping
3375 .iter()
3376 .copied()
3377 .collect::<BTreeMap<_, _>>();
3378
3379 let empty_projection = ReaderProjection {
3380 column_indices: Vec::default(),
3381 schema: Arc::new(Schema::default()),
3382 };
3383
3384 for columns in [
3385 vec!["score"],
3386 vec!["location"],
3387 vec!["categories"],
3388 vec!["score.x"],
3389 vec!["score", "categories"],
3390 vec!["score", "location"],
3391 vec!["location", "categories"],
3392 vec!["score.y", "location", "categories"],
3393 ] {
3394 debug!("Testing round trip with projection {:?}", columns);
3395 for use_field_ids in [true, false] {
3396 let file_reader = FileReader::try_open(
3398 file_scheduler.clone(),
3399 None,
3400 Arc::<DecoderPlugins>::default(),
3401 &test_cache(),
3402 FileReaderOptions::default(),
3403 )
3404 .await
3405 .unwrap();
3406
3407 let projected_schema = written_file.schema.project(&columns).unwrap();
3408 let projection = if use_field_ids {
3409 ReaderProjection::from_field_ids(
3410 file_reader.metadata.version(),
3411 &projected_schema,
3412 &field_id_mapping,
3413 )
3414 .unwrap()
3415 } else {
3416 ReaderProjection::from_column_names(
3417 file_reader.metadata.version(),
3418 &written_file.schema,
3419 &columns,
3420 )
3421 .unwrap()
3422 };
3423
3424 let batch_stream = file_reader
3425 .read_stream_projected(
3426 lance_io::ReadBatchParams::RangeFull,
3427 1024,
3428 16,
3429 projection.clone(),
3430 FilterExpression::no_filter(),
3431 )
3432 .await
3433 .unwrap();
3434
3435 let projection_arrow = ArrowSchema::from(projection.schema.as_ref());
3436 verify_expected(
3437 &written_file.data,
3438 batch_stream,
3439 1024,
3440 Some(Box::new(move |batch: &RecordBatch| {
3441 batch.project_by_schema(&projection_arrow).unwrap()
3442 })),
3443 )
3444 .await;
3445
3446 let file_reader = FileReader::try_open(
3448 file_scheduler.clone(),
3449 Some(projection.clone()),
3450 Arc::<DecoderPlugins>::default(),
3451 &test_cache(),
3452 FileReaderOptions::default(),
3453 )
3454 .await
3455 .unwrap();
3456
3457 let batch_stream = file_reader
3458 .read_stream(
3459 lance_io::ReadBatchParams::RangeFull,
3460 1024,
3461 16,
3462 FilterExpression::no_filter(),
3463 )
3464 .await
3465 .unwrap();
3466
3467 let projection_arrow = ArrowSchema::from(projection.schema.as_ref());
3468 verify_expected(
3469 &written_file.data,
3470 batch_stream,
3471 1024,
3472 Some(Box::new(move |batch: &RecordBatch| {
3473 batch.project_by_schema(&projection_arrow).unwrap()
3474 })),
3475 )
3476 .await;
3477
3478 assert!(
3479 file_reader
3480 .read_stream_projected(
3481 lance_io::ReadBatchParams::RangeFull,
3482 1024,
3483 16,
3484 empty_projection.clone(),
3485 FilterExpression::no_filter(),
3486 )
3487 .await
3488 .is_err()
3489 );
3490 }
3491 }
3492
3493 assert!(
3494 FileReader::try_open(
3495 file_scheduler.clone(),
3496 Some(empty_projection),
3497 Arc::<DecoderPlugins>::default(),
3498 &test_cache(),
3499 FileReaderOptions::default(),
3500 )
3501 .await
3502 .is_err()
3503 );
3504
3505 let arrow_schema = ArrowSchema::new(vec![
3506 Field::new("x", DataType::Int32, true),
3507 Field::new("y", DataType::Int32, true),
3508 ]);
3509 let schema = Schema::try_from(&arrow_schema).unwrap();
3510
3511 let projection_with_dupes = ReaderProjection {
3512 column_indices: vec![0, 0],
3513 schema: Arc::new(schema),
3514 };
3515
3516 assert!(
3517 FileReader::try_open(
3518 file_scheduler.clone(),
3519 Some(projection_with_dupes),
3520 Arc::<DecoderPlugins>::default(),
3521 &test_cache(),
3522 FileReaderOptions::default(),
3523 )
3524 .await
3525 .is_err()
3526 );
3527 }
3528
3529 #[tokio::test]
3530 async fn test_lazy_reader_direct_projection_matches_eager_reader() {
3531 let fs = FsFixture::default();
3532 let written_file = create_wide_direct_file(&fs, 16).await;
3533
3534 let file_scheduler = fs
3535 .scheduler
3536 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3537 .await
3538 .unwrap();
3539 let projection = ReaderProjection::from_column_names(
3540 LanceFileVersion::V2_1,
3541 &written_file.schema,
3542 &["c10"],
3543 )
3544 .unwrap();
3545
3546 let eager_reader = FileReader::try_open(
3547 file_scheduler.clone(),
3548 None,
3549 Arc::<DecoderPlugins>::default(),
3550 &test_cache(),
3551 FileReaderOptions::default(),
3552 )
3553 .await
3554 .unwrap();
3555 let expected = eager_reader
3556 .read_stream_projected(
3557 lance_io::ReadBatchParams::RangeFull,
3558 127,
3559 16,
3560 projection.clone(),
3561 FilterExpression::no_filter(),
3562 )
3563 .await
3564 .unwrap()
3565 .try_collect::<Vec<_>>()
3566 .await
3567 .unwrap();
3568
3569 let cache = test_cache();
3570 let lazy_reader = ProjectedFileReader::try_open(
3571 file_scheduler,
3572 Some(projection.clone()),
3573 Arc::<DecoderPlugins>::default(),
3574 &cache,
3575 FileReaderOptions::default(),
3576 )
3577 .await
3578 .unwrap();
3579 let tasks = lazy_reader
3580 .read_tasks(
3581 lance_io::ReadBatchParams::RangeFull,
3582 127,
3583 None,
3584 FilterExpression::no_filter(),
3585 )
3586 .await
3587 .unwrap();
3588 let actual = collect_read_tasks(tasks, 16).await;
3589
3590 assert_eq!(expected, actual);
3591 }
3592
3593 #[tokio::test]
3594 async fn test_lazy_reader_loads_only_requested_column_metadata() {
3595 let fs = FsFixture::default();
3596 let written_file = create_wide_direct_file(&fs, 512).await;
3597
3598 let projection = ReaderProjection::from_column_names(
3599 LanceFileVersion::V2_1,
3600 &written_file.schema,
3601 &["c0"],
3602 )
3603 .unwrap();
3604 let file_scheduler = fs
3605 .scheduler
3606 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3607 .await
3608 .unwrap();
3609 let lazy_reader = ProjectedFileReader::try_open(
3610 file_scheduler,
3611 Some(projection.clone()),
3612 Arc::<DecoderPlugins>::default(),
3613 &test_cache(),
3614 FileReaderOptions::default(),
3615 )
3616 .await
3617 .unwrap();
3618 let selected_column = projection.column_indices[0] as usize;
3619 let requested_metadata_bytes = lazy_reader
3620 .metadata_index()
3621 .unwrap()
3622 .column_metadata_offsets[selected_column]
3623 .1;
3624 let total_metadata_bytes = lazy_reader
3625 .metadata_index()
3626 .unwrap()
3627 .column_metadata_offsets
3628 .iter()
3629 .map(|(_, length)| *length)
3630 .sum::<u64>();
3631 assert!(
3632 total_metadata_bytes > 8 * fs.object_store.block_size() as u64,
3633 "test file metadata is too small to prove lazy loading: {total_metadata_bytes} bytes"
3634 );
3635
3636 fs.object_store.io_stats_incremental();
3637 let tasks = lazy_reader
3638 .read_tasks(
3639 lance_io::ReadBatchParams::Range(0..0),
3640 1024,
3641 Some(projection.clone()),
3642 FilterExpression::no_filter(),
3643 )
3644 .await
3645 .unwrap();
3646 let batches = collect_read_tasks(tasks, 1).await;
3647 assert!(batches.is_empty());
3648
3649 let stats = fs.object_store.io_stats_incremental();
3650 assert!(
3651 stats.read_bytes < total_metadata_bytes / 2,
3652 "lazy read fetched too much metadata: read {} bytes, requested column metadata is {} bytes, total column metadata is {} bytes",
3653 stats.read_bytes,
3654 requested_metadata_bytes,
3655 total_metadata_bytes
3656 );
3657
3658 fs.object_store.io_stats_incremental();
3659 let tasks = lazy_reader
3660 .read_tasks(
3661 lance_io::ReadBatchParams::Range(0..0),
3662 1024,
3663 Some(projection),
3664 FilterExpression::no_filter(),
3665 )
3666 .await
3667 .unwrap();
3668 let batches = collect_read_tasks(tasks, 1).await;
3669 assert!(batches.is_empty());
3670
3671 let stats = fs.object_store.io_stats_incremental();
3672 assert_eq!(
3673 stats.read_iops, 0,
3674 "cached column metadata should avoid repeat metadata I/O"
3675 );
3676 assert_eq!(
3677 stats.read_bytes, 0,
3678 "cached column metadata should avoid repeat metadata reads"
3679 );
3680 }
3681
3682 async fn assert_lazy_projection_matches_eager_and_reads_metadata_subset(
3683 fs: &FsFixture,
3684 projection: ReaderProjection,
3685 shape: &str,
3686 ) -> Vec<RecordBatch> {
3687 let file_scheduler = fs
3688 .scheduler
3689 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3690 .await
3691 .unwrap();
3692 let eager_reader = FileReader::try_open(
3693 file_scheduler.clone(),
3694 None,
3695 Arc::<DecoderPlugins>::default(),
3696 &test_cache(),
3697 FileReaderOptions::default(),
3698 )
3699 .await
3700 .unwrap();
3701 let expected = eager_reader
3702 .read_stream_projected(
3703 lance_io::ReadBatchParams::RangeFull,
3704 127,
3705 16,
3706 projection.clone(),
3707 FilterExpression::no_filter(),
3708 )
3709 .await
3710 .unwrap()
3711 .try_collect::<Vec<_>>()
3712 .await
3713 .unwrap();
3714
3715 let cache = test_cache();
3716 let lazy_reader = ProjectedFileReader::try_open(
3717 file_scheduler,
3718 Some(projection.clone()),
3719 Arc::<DecoderPlugins>::default(),
3720 &cache,
3721 FileReaderOptions::default(),
3722 )
3723 .await
3724 .unwrap();
3725 let metadata_index = lazy_reader.metadata_index().unwrap();
3726 let requested_metadata_bytes = projection
3727 .column_indices
3728 .iter()
3729 .map(|column_index| metadata_index.column_metadata_offsets[*column_index as usize].1)
3730 .sum::<u64>();
3731 let total_metadata_bytes = metadata_index
3732 .column_metadata_offsets
3733 .iter()
3734 .map(|(_, length)| *length)
3735 .sum::<u64>();
3736 assert!(total_metadata_bytes > requested_metadata_bytes * 8);
3737
3738 fs.object_store.io_stats_incremental();
3739 let tasks = lazy_reader
3740 .read_tasks(
3741 lance_io::ReadBatchParams::Range(0..0),
3742 127,
3743 None,
3744 FilterExpression::no_filter(),
3745 )
3746 .await
3747 .unwrap();
3748 assert!(collect_read_tasks(tasks, 1).await.is_empty());
3749 let metadata_stats = fs.object_store.io_stats_incremental();
3750 assert!(
3751 metadata_stats.read_bytes < total_metadata_bytes / 2,
3752 "lazy {shape} read fetched too much metadata: read {} bytes, requested column metadata is {} bytes, total column metadata is {} bytes",
3753 metadata_stats.read_bytes,
3754 requested_metadata_bytes,
3755 total_metadata_bytes
3756 );
3757
3758 let tasks = lazy_reader
3759 .read_tasks(
3760 lance_io::ReadBatchParams::RangeFull,
3761 127,
3762 None,
3763 FilterExpression::no_filter(),
3764 )
3765 .await
3766 .unwrap();
3767 let actual = collect_read_tasks(tasks, 16).await;
3768 assert_eq!(expected, actual);
3769 actual
3770 }
3771
3772 #[tokio::test]
3773 async fn test_lazy_reader_fixed_size_list_projection_matches_eager_reader() {
3774 let fs = FsFixture::default();
3775 let written_file = create_wide_fixed_size_list_file(&fs, 512).await;
3776 let projection = ReaderProjection::from_column_names(
3777 LanceFileVersion::V2_1,
3778 &written_file.schema,
3779 &["c17", "c509"],
3780 )
3781 .unwrap();
3782 assert!(ProjectedFileReader::supports_projection(
3783 &projection,
3784 LanceFileVersion::V2_1
3785 ));
3786 assert!(!ProjectedFileReader::supports_projection(
3787 &projection,
3788 LanceFileVersion::V2_0
3789 ));
3790 assert_lazy_projection_matches_eager_and_reads_metadata_subset(
3791 &fs,
3792 projection,
3793 "fixed-size-list",
3794 )
3795 .await;
3796 }
3797
3798 #[tokio::test]
3799 async fn test_lazy_reader_nested_projection_compacts_physical_columns() {
3800 let fs = FsFixture::default();
3801 let written_file = create_wide_structural_file(&fs, 128).await;
3802 let projection = ReaderProjection::from_column_names(
3803 LanceFileVersion::V2_1,
3804 &written_file.schema,
3805 &["s97.y", "l4", "s3"],
3806 )
3807 .unwrap();
3808
3809 assert_eq!(
3810 projection
3811 .schema
3812 .fields
3813 .iter()
3814 .map(|field| field.name.as_str())
3815 .collect::<Vec<_>>(),
3816 vec!["s97", "l4", "s3"]
3817 );
3818 assert_eq!(projection.schema.fields[0].children.len(), 1);
3819 assert_eq!(projection.schema.fields[0].children[0].name, "y");
3820 assert_eq!(projection.schema.fields[2].children.len(), 2);
3821 assert_eq!(projection.column_indices.len(), 4);
3822 assert!(
3823 projection
3824 .column_indices
3825 .windows(2)
3826 .any(|indices| indices[0] > indices[1]),
3827 "the projection must reorder physical columns to exercise compact remapping"
3828 );
3829 assert!(ProjectedFileReader::supports_projection(
3830 &projection,
3831 LanceFileVersion::V2_1
3832 ));
3833 let actual = assert_lazy_projection_matches_eager_and_reads_metadata_subset(
3834 &fs, projection, "nested",
3835 )
3836 .await;
3837 assert!(
3838 actual
3839 .iter()
3840 .flat_map(|batch| batch.columns())
3841 .any(|column| column.null_count() > 0),
3842 "the structural projection must exercise nullable arrays"
3843 );
3844 }
3845
3846 #[rstest]
3847 #[case::before_metadata_region(90, 5)]
3848 #[case::after_metadata_region(190, 20)]
3849 fn test_decode_cmo_table_rejects_out_of_range_offsets(
3850 #[case] position: u64,
3851 #[case] length: u64,
3852 ) {
3853 let mut cmo_table = [0; 16];
3854 cmo_table[0..8].copy_from_slice(&position.to_le_bytes());
3855 cmo_table[8..16].copy_from_slice(&length.to_le_bytes());
3856 let footer = super::Footer {
3857 column_meta_start: 100,
3858 column_meta_offsets_start: 200,
3859 global_buff_offsets_start: 200,
3860 num_global_buffers: 0,
3861 num_columns: 1,
3862 major_version: 2,
3863 minor_version: 1,
3864 };
3865
3866 let err = FileReader::decode_cmo_table(Bytes::copy_from_slice(&cmo_table), &footer)
3867 .expect_err("out-of-range CMO entries must be rejected");
3868 assert!(
3869 matches!(err, lance_core::Error::InvalidInput { .. }),
3870 "expected InvalidInput, got {err:?}"
3871 );
3872 }
3873
3874 #[rstest]
3875 #[case::blob(BLOB_META_KEY)]
3876 #[case::packed_struct("lance-encoding:packed")]
3877 #[tokio::test]
3878 async fn test_lazy_reader_rejects_opaque_projection(#[case] metadata_key: &str) {
3879 let fs = FsFixture::default();
3880 let written_file = create_some_file(&fs, LanceFileVersion::V2_1).await;
3881
3882 let ordinary_projection = ReaderProjection::from_column_names(
3883 LanceFileVersion::V2_1,
3884 &written_file.schema,
3885 &["location.x"],
3886 )
3887 .unwrap();
3888 assert_eq!(ordinary_projection.schema.fields[0].children.len(), 1);
3889 assert!(ProjectedFileReader::supports_projection(
3890 &ordinary_projection,
3891 LanceFileVersion::V2_1
3892 ));
3893
3894 let file_scheduler = fs
3895 .scheduler
3896 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3897 .await
3898 .unwrap();
3899 let err = ProjectedFileReader::try_open(
3900 file_scheduler.clone(),
3901 None,
3902 Arc::<DecoderPlugins>::default(),
3903 &test_cache(),
3904 FileReaderOptions::default(),
3905 )
3906 .await
3907 .unwrap_err();
3908 assert!(
3909 matches!(err, lance_core::Error::InvalidInput { .. }),
3910 "expected InvalidInput, got {err:?}"
3911 );
3912
3913 let mut projection = ordinary_projection;
3914 Arc::make_mut(&mut projection.schema).fields[0]
3915 .metadata
3916 .insert(metadata_key.to_string(), "true".to_string());
3917 assert!(!ProjectedFileReader::supports_projection(
3918 &projection,
3919 LanceFileVersion::V2_1
3920 ));
3921
3922 let err = ProjectedFileReader::try_open(
3923 file_scheduler,
3924 Some(projection),
3925 Arc::<DecoderPlugins>::default(),
3926 &test_cache(),
3927 FileReaderOptions::default(),
3928 )
3929 .await
3930 .unwrap_err();
3931 assert!(
3932 matches!(err, lance_core::Error::NotSupported { .. }),
3933 "expected NotSupported for {metadata_key}, got {err:?}"
3934 );
3935 }
3936
3937 #[tokio::test]
3944 async fn test_lazy_reader_validates_unequal_length_projection() {
3945 use arrow_array::Int32Array;
3946 use lance_io::ReadBatchParams;
3947
3948 let arrow_schema = Arc::new(ArrowSchema::new(vec![
3949 Field::new("a", DataType::Int32, true),
3950 Field::new("c", DataType::Int32, true),
3951 ]));
3952 let lance_schema = Schema::try_from(arrow_schema.as_ref()).unwrap();
3953
3954 let fs = FsFixture::default();
3955 let options = FileWriterOptions {
3956 format_version: Some(LanceFileVersion::V2_1),
3957 ..Default::default()
3958 };
3959 let mut writer = FileWriter::try_new(
3960 fs.object_store.create(&fs.tmp_path).await.unwrap(),
3961 lance_schema.clone(),
3962 options,
3963 )
3964 .unwrap();
3965 writer
3967 .write_column(0, Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])))
3968 .await
3969 .unwrap();
3970 writer
3971 .write_column(1, Arc::new(Int32Array::from(vec![100])))
3972 .await
3973 .unwrap();
3974 writer.finish().await.unwrap();
3975
3976 let file_scheduler = fs
3977 .scheduler
3978 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3979 .await
3980 .unwrap();
3981 let cache = test_cache();
3982 let open_indexed = |names: &[&str]| {
3983 let projection =
3984 ReaderProjection::from_column_names(LanceFileVersion::V2_1, &lance_schema, names)
3985 .unwrap();
3986 ProjectedFileReader::try_open(
3987 file_scheduler.clone(),
3988 Some(projection),
3989 Arc::<DecoderPlugins>::default(),
3990 &cache,
3991 FileReaderOptions::default(),
3992 )
3993 };
3994
3995 let lazy = open_indexed(&["a", "c"]).await.unwrap();
3998 let err = match lazy
4001 .read_tasks(
4002 ReadBatchParams::RangeFull,
4003 1024,
4004 None,
4005 FilterExpression::no_filter(),
4006 )
4007 .await
4008 {
4009 Ok(_) => panic!("expected the mismatched-length projection to be rejected"),
4010 Err(e) => e.to_string(),
4011 };
4012 assert!(
4013 err.contains("a=5") && err.contains("c=1"),
4014 "error should name each column's length, got: {err}"
4015 );
4016
4017 let lazy = open_indexed(&["c"]).await.unwrap();
4020 let tasks = lazy
4021 .read_tasks(
4022 ReadBatchParams::RangeFull,
4023 1024,
4024 None,
4025 FilterExpression::no_filter(),
4026 )
4027 .await
4028 .unwrap();
4029 let batches = collect_read_tasks(tasks, 16).await;
4030 let values: Vec<Option<i32>> = batches
4031 .iter()
4032 .flat_map(|b| {
4033 b.column(0)
4034 .as_any()
4035 .downcast_ref::<Int32Array>()
4036 .unwrap()
4037 .iter()
4038 .collect::<Vec<_>>()
4039 })
4040 .collect();
4041 assert_eq!(values, vec![Some(100)]);
4042 }
4043
4044 #[test_log::test(tokio::test)]
4045 async fn test_compressing_buffer() {
4046 let fs = FsFixture::default();
4047
4048 let written_file = create_some_file(&fs, LanceFileVersion::V2_0).await;
4049 let file_scheduler = fs
4050 .scheduler
4051 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4052 .await
4053 .unwrap();
4054
4055 let file_reader = FileReader::try_open(
4057 file_scheduler.clone(),
4058 None,
4059 Arc::<DecoderPlugins>::default(),
4060 &test_cache(),
4061 FileReaderOptions::default(),
4062 )
4063 .await
4064 .unwrap();
4065
4066 let mut projection = written_file.schema.project(&["score"]).unwrap();
4067 for field in projection.fields.iter_mut() {
4068 field
4069 .metadata
4070 .insert("lance:compression".to_string(), "zstd".to_string());
4071 }
4072 let projection = ReaderProjection {
4073 column_indices: projection.fields.iter().map(|f| f.id as u32).collect(),
4074 schema: Arc::new(projection),
4075 };
4076
4077 let batch_stream = file_reader
4078 .read_stream_projected(
4079 lance_io::ReadBatchParams::RangeFull,
4080 1024,
4081 16,
4082 projection.clone(),
4083 FilterExpression::no_filter(),
4084 )
4085 .await
4086 .unwrap();
4087
4088 let projection_arrow = Arc::new(ArrowSchema::from(projection.schema.as_ref()));
4089 verify_expected(
4090 &written_file.data,
4091 batch_stream,
4092 1024,
4093 Some(Box::new(move |batch: &RecordBatch| {
4094 batch.project_by_schema(&projection_arrow).unwrap()
4095 })),
4096 )
4097 .await;
4098 }
4099
4100 #[tokio::test]
4101 async fn test_read_all() {
4102 let fs = FsFixture::default();
4103 let WrittenFile { data, .. } = create_some_file(&fs, LanceFileVersion::V2_0).await;
4104 let total_rows = data.iter().map(|batch| batch.num_rows()).sum::<usize>();
4105
4106 let file_scheduler = fs
4107 .scheduler
4108 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4109 .await
4110 .unwrap();
4111 let file_reader = FileReader::try_open(
4112 file_scheduler.clone(),
4113 None,
4114 Arc::<DecoderPlugins>::default(),
4115 &test_cache(),
4116 FileReaderOptions::default(),
4117 )
4118 .await
4119 .unwrap();
4120
4121 let batches = file_reader
4122 .read_stream(
4123 lance_io::ReadBatchParams::RangeFull,
4124 total_rows as u32,
4125 16,
4126 FilterExpression::no_filter(),
4127 )
4128 .await
4129 .unwrap()
4130 .try_collect::<Vec<_>>()
4131 .await
4132 .unwrap();
4133 assert_eq!(batches.len(), 1);
4134 assert_eq!(batches[0].num_rows(), total_rows);
4135 }
4136
4137 #[rstest]
4138 #[tokio::test]
4139 async fn test_blocking_take(
4140 #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1, LanceFileVersion::V2_2)]
4141 version: LanceFileVersion,
4142 ) {
4143 let fs = FsFixture::default();
4144 let WrittenFile { data, schema, .. } = create_some_file(&fs, version).await;
4145 let total_rows = data.iter().map(|batch| batch.num_rows()).sum::<usize>();
4146
4147 let file_scheduler = fs
4148 .scheduler
4149 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4150 .await
4151 .unwrap();
4152 let file_reader = FileReader::try_open(
4153 file_scheduler.clone(),
4154 Some(ReaderProjection::from_column_names(version, &schema, &["score"]).unwrap()),
4155 Arc::<DecoderPlugins>::default(),
4156 &test_cache(),
4157 FileReaderOptions::default(),
4158 )
4159 .await
4160 .unwrap();
4161
4162 let batches = tokio::task::spawn_blocking(move || {
4163 file_reader
4164 .read_stream_projected_blocking(
4165 lance_io::ReadBatchParams::Indices(UInt32Array::from(vec![0, 1, 2, 3, 4])),
4166 total_rows as u32,
4167 None,
4168 FilterExpression::no_filter(),
4169 )
4170 .unwrap()
4171 .collect::<ArrowResult<Vec<_>>>()
4172 .unwrap()
4173 })
4174 .await
4175 .unwrap();
4176
4177 assert_eq!(batches.len(), 1);
4178 assert_eq!(batches[0].num_rows(), 5);
4179 assert_eq!(batches[0].num_columns(), 1);
4180 }
4181
4182 #[tokio::test(flavor = "multi_thread")]
4183 async fn test_drop_in_progress() {
4184 let fs = FsFixture::default();
4185 let WrittenFile { data, .. } = create_some_file(&fs, LanceFileVersion::V2_0).await;
4186 let total_rows = data.iter().map(|batch| batch.num_rows()).sum::<usize>();
4187
4188 let file_scheduler = fs
4189 .scheduler
4190 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4191 .await
4192 .unwrap();
4193 let file_reader = FileReader::try_open(
4194 file_scheduler.clone(),
4195 None,
4196 Arc::<DecoderPlugins>::default(),
4197 &test_cache(),
4198 FileReaderOptions::default(),
4199 )
4200 .await
4201 .unwrap();
4202
4203 let mut batches = file_reader
4204 .read_stream(
4205 lance_io::ReadBatchParams::RangeFull,
4206 (total_rows / 10) as u32,
4207 16,
4208 FilterExpression::no_filter(),
4209 )
4210 .await
4211 .unwrap();
4212
4213 drop(file_reader);
4214
4215 let batch = batches.next().await.unwrap().unwrap();
4216 assert!(batch.num_rows() > 0);
4217
4218 drop(batches);
4220 }
4221
4222 #[tokio::test]
4223 async fn drop_while_scheduling() {
4224 let fs = FsFixture::default();
4234 let written_file = create_some_file(&fs, LanceFileVersion::V2_0).await;
4235 let total_rows = written_file
4236 .data
4237 .iter()
4238 .map(|batch| batch.num_rows())
4239 .sum::<usize>();
4240
4241 let file_scheduler = fs
4242 .scheduler
4243 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4244 .await
4245 .unwrap();
4246 let file_reader = FileReader::try_open(
4247 file_scheduler.clone(),
4248 None,
4249 Arc::<DecoderPlugins>::default(),
4250 &test_cache(),
4251 FileReaderOptions::default(),
4252 )
4253 .await
4254 .unwrap();
4255
4256 let projection =
4257 ReaderProjection::from_whole_schema(&written_file.schema, LanceFileVersion::V2_0);
4258 let column_infos = file_reader
4259 .collect_columns_from_projection(&projection)
4260 .unwrap();
4261 let mut decode_scheduler = DecodeBatchScheduler::try_new(
4262 &projection.schema,
4263 &projection.column_indices,
4264 &column_infos,
4265 &vec![],
4266 total_rows as u64,
4267 Arc::<DecoderPlugins>::default(),
4268 file_reader.core.scheduler.clone(),
4269 test_cache(),
4270 &FilterExpression::no_filter(),
4271 &DecoderConfig::default(),
4272 )
4273 .await
4274 .unwrap();
4275
4276 let range = 0..total_rows as u64;
4277
4278 let (tx, rx) = mpsc::unbounded_channel();
4279
4280 drop(rx);
4282
4283 decode_scheduler.schedule_range(
4285 range,
4286 &FilterExpression::no_filter(),
4287 tx,
4288 file_reader.core.scheduler.clone(),
4289 )
4290 }
4291
4292 #[tokio::test]
4293 async fn test_read_empty_range() {
4294 let fs = FsFixture::default();
4295 create_some_file(&fs, LanceFileVersion::V2_0).await;
4296
4297 let file_scheduler = fs
4298 .scheduler
4299 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4300 .await
4301 .unwrap();
4302 let file_reader = FileReader::try_open(
4303 file_scheduler.clone(),
4304 None,
4305 Arc::<DecoderPlugins>::default(),
4306 &test_cache(),
4307 FileReaderOptions::default(),
4308 )
4309 .await
4310 .unwrap();
4311
4312 let batches = file_reader
4314 .read_stream(
4315 lance_io::ReadBatchParams::Range(0..0),
4316 1024,
4317 16,
4318 FilterExpression::no_filter(),
4319 )
4320 .await
4321 .unwrap()
4322 .try_collect::<Vec<_>>()
4323 .await
4324 .unwrap();
4325
4326 assert_eq!(batches.len(), 0);
4327
4328 let batches = file_reader
4330 .read_stream(
4331 lance_io::ReadBatchParams::Ranges(Arc::new([0..1, 2..2])),
4332 1024,
4333 16,
4334 FilterExpression::no_filter(),
4335 )
4336 .await
4337 .unwrap()
4338 .try_collect::<Vec<_>>()
4339 .await
4340 .unwrap();
4341 assert_eq!(batches.len(), 1);
4342 }
4343
4344 async fn write_file_with_global_buffer(fs: &FsFixture, buffer: Bytes) {
4345 let lance_schema =
4346 lance_core::datatypes::Schema::try_from(&ArrowSchema::new(vec![Field::new(
4347 "foo",
4348 DataType::Int32,
4349 true,
4350 )]))
4351 .unwrap();
4352
4353 let mut file_writer = FileWriter::try_new(
4354 fs.object_store.create(&fs.tmp_path).await.unwrap(),
4355 lance_schema,
4356 FileWriterOptions::default(),
4357 )
4358 .unwrap();
4359
4360 let buf_index = file_writer.add_global_buffer(buffer).await.unwrap();
4361 assert_eq!(buf_index, 1);
4362
4363 file_writer.finish().await.unwrap();
4364 }
4365
4366 #[rstest]
4370 #[case::within_tail_window(true)]
4371 #[case::outside_tail_window(false)]
4372 #[tokio::test]
4373 async fn test_read_global_buffer(#[case] within_window: bool) {
4374 let fs = FsFixture::default();
4375
4376 let block_size = fs.object_store.block_size();
4377 let buffer = if within_window {
4378 Bytes::from_static(b"hello")
4379 } else {
4380 Bytes::from(vec![7u8; 2 * block_size])
4381 };
4382 let expected_read_iops = if within_window { 0 } else { 1 };
4383
4384 write_file_with_global_buffer(&fs, buffer.clone()).await;
4385
4386 let file_scheduler = fs
4387 .scheduler
4388 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4389 .await
4390 .unwrap();
4391 let file_reader = FileReader::try_open(
4392 file_scheduler,
4393 None,
4394 Arc::<DecoderPlugins>::default(),
4395 &test_cache(),
4396 FileReaderOptions::default(),
4397 )
4398 .await
4399 .unwrap();
4400
4401 let retained = &file_reader.metadata().retained_global_buffers;
4404 assert!(!retained.contains_key(&0), "schema must not be retained");
4405 assert_eq!(retained.contains_key(&1), within_window);
4406
4407 fs.object_store.io_stats_incremental();
4409
4410 let buf = file_reader.read_global_buffer(1).await.unwrap();
4411 assert_eq!(buf, buffer);
4412
4413 let stats = fs.object_store.io_stats_incremental();
4414 assert_eq!(stats.read_iops, expected_read_iops);
4415 }
4416
4417 #[tokio::test]
4420 async fn test_read_global_buffer_no_user_buffers() {
4421 let fs = FsFixture::default();
4422 create_some_file(&fs, LanceFileVersion::V2_1).await;
4423
4424 let file_scheduler = fs
4425 .scheduler
4426 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4427 .await
4428 .unwrap();
4429 let file_reader = FileReader::try_open(
4430 file_scheduler,
4431 None,
4432 Arc::<DecoderPlugins>::default(),
4433 &test_cache(),
4434 FileReaderOptions::default(),
4435 )
4436 .await
4437 .unwrap();
4438
4439 let metadata = file_reader.metadata();
4440 assert_eq!(metadata.file_buffers.len(), 1, "expected only the schema");
4441 assert!(
4442 metadata.retained_global_buffers.is_empty(),
4443 "a file with no user global buffers must retain nothing"
4444 );
4445 }
4446
4447 #[rstest]
4448 #[tokio::test]
4449 async fn test_deep_size_of_includes_column_metadata(
4450 #[values(
4451 LanceFileVersion::V2_0,
4452 LanceFileVersion::V2_1,
4453 LanceFileVersion::V2_2,
4454 LanceFileVersion::V2_3
4455 )]
4456 version: LanceFileVersion,
4457 ) {
4458 use lance_core::deepsize::DeepSizeOf;
4463
4464 let fs = FsFixture::default();
4465 let _written = create_some_file(&fs, version).await;
4466 let cache = test_cache();
4467 let file_scheduler = fs
4468 .scheduler
4469 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4470 .await
4471 .unwrap();
4472 let file_reader = FileReader::try_open(
4473 file_scheduler,
4474 None,
4475 Arc::<DecoderPlugins>::default(),
4476 &cache,
4477 FileReaderOptions::default(),
4478 )
4479 .await
4480 .unwrap();
4481
4482 let metadata = file_reader.metadata();
4483 let deep_size = metadata.deep_size_of();
4484
4485 assert!(
4491 deep_size > 1024,
4492 "deep_size_of ({deep_size}) is suspiciously small — \
4493 column_metadatas and column_infos may not be accounted for"
4494 );
4495
4496 assert!(
4498 !metadata.column_metadatas.is_empty(),
4499 "Expected non-empty column_metadatas"
4500 );
4501
4502 let num_columns = metadata.column_metadatas.len();
4505 assert!(
4506 deep_size > num_columns * 50,
4507 "deep_size_of ({deep_size}) should scale with column count ({num_columns})"
4508 );
4509 }
4510
4511 #[tokio::test]
4512 async fn test_read_global_buffer_out_of_range() {
4513 let fs = FsFixture::default();
4514
4515 write_file_with_global_buffer(&fs, Bytes::from_static(b"hello")).await;
4516
4517 let file_scheduler = fs
4518 .scheduler
4519 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4520 .await
4521 .unwrap();
4522 let file_reader = FileReader::try_open(
4523 file_scheduler,
4524 None,
4525 Arc::<DecoderPlugins>::default(),
4526 &test_cache(),
4527 FileReaderOptions::default(),
4528 )
4529 .await
4530 .unwrap();
4531
4532 let err = file_reader.read_global_buffer(2).await.unwrap_err();
4535 assert!(
4536 matches!(err, lance_core::Error::InvalidInput { .. }),
4537 "expected InvalidInput, got: {err:?}"
4538 );
4539 let msg = err.to_string();
4540 assert!(msg.contains('2'), "error should mention the index: {msg}");
4541 }
4542
4543 #[rstest]
4548 fn test_validate_struct_child_lengths(#[values(false, true)] is_structural: bool) {
4549 let run = |dt: DataType, indices: &[u32], lengths: Vec<u64>| -> lance_core::Result<u64> {
4550 let arrow = ArrowSchema::new(vec![Field::new("s", dt, true)]);
4551 let schema = Schema::try_from(&arrow).unwrap();
4552 let column_len = |c: usize| Ok(lengths[c]);
4553 let mut cursor = 0usize;
4554 let mut field_lengths = Vec::new();
4555 for field in &schema.fields {
4556 let rows = validate_field_length(
4557 field,
4558 is_structural,
4559 true,
4560 indices,
4561 &mut cursor,
4562 &column_len,
4563 )?;
4564 field_lengths.push((field.name.as_str(), rows));
4565 }
4566 verify_uniform_lengths(&field_lengths)
4567 };
4568
4569 let struct_ty = || {
4570 DataType::Struct(Fields::from(vec![
4571 Field::new("a", DataType::Int32, true),
4572 Field::new("b", DataType::Int32, true),
4573 ]))
4574 };
4575
4576 let (indices, equal, unequal): (&[u32], Vec<u64>, Vec<u64>) = if is_structural {
4579 (&[0, 1], vec![5, 5], vec![5, 3])
4580 } else {
4581 (&[0, 1, 2], vec![5, 5, 5], vec![5, 5, 3])
4582 };
4583
4584 assert_eq!(run(struct_ty(), indices, equal).unwrap(), 5);
4585
4586 let err = run(struct_ty(), indices, unequal).unwrap_err();
4587 let msg = err.to_string();
4588 assert!(
4589 msg.contains("differing lengths") && msg.contains('b'),
4590 "expected a child-length error naming 'b', got: {msg}"
4591 );
4592 }
4593
4594 #[test]
4595 fn test_validate_v2_0_unloaded_blob_projection_is_opaque() {
4596 let metadata = HashMap::from([(BLOB_META_KEY.to_string(), "true".to_string())]);
4597 let arrow = ArrowSchema::new(vec![
4598 Field::new("blob", DataType::LargeBinary, true).with_metadata(metadata),
4599 ]);
4600 let mut schema = Schema::try_from(&arrow).unwrap();
4601 schema.fields[0].unloaded_mut();
4602 let projection = ReaderProjection {
4603 schema: Arc::new(schema),
4604 column_indices: vec![0],
4605 };
4606 let column_len = |column: usize| {
4607 assert_eq!(column, 0);
4608 Ok(3)
4609 };
4610 let mut cursor = 0usize;
4611
4612 let rows = validate_field_length(
4613 &projection.schema.fields[0],
4614 false,
4615 true,
4616 &projection.column_indices,
4617 &mut cursor,
4618 &column_len,
4619 )
4620 .unwrap();
4621
4622 assert_eq!(rows, 3);
4623 assert_eq!(cursor, 1);
4624 }
4625
4626 #[test]
4627 fn test_validate_length_list_and_empty_struct() {
4628 let validate = |dt: DataType,
4629 is_structural: bool,
4630 indices: &[u32],
4631 lengths: Vec<u64>|
4632 -> lance_core::Result<u64> {
4633 let arrow = ArrowSchema::new(vec![Field::new("f", dt, true)]);
4634 let schema = Schema::try_from(&arrow).unwrap();
4635 let column_len = |c: usize| Ok(lengths[c]);
4636 let mut cursor = 0usize;
4637 validate_field_length(
4638 &schema.fields[0],
4639 is_structural,
4640 true,
4641 indices,
4642 &mut cursor,
4643 &column_len,
4644 )
4645 };
4646
4647 let list_ty = DataType::List(Arc::new(Field::new("item", DataType::Int32, true)));
4652 assert_eq!(
4653 validate(list_ty.clone(), false, &[0, 1], vec![5, 17]).unwrap(),
4654 5
4655 );
4656 assert_eq!(validate(list_ty, true, &[0], vec![5]).unwrap(), 5);
4657
4658 let list_of_struct = DataType::List(Arc::new(Field::new(
4665 "item",
4666 DataType::Struct(Fields::from(vec![
4667 Field::new("a", DataType::Int32, true),
4668 Field::new("b", DataType::Int32, true),
4669 ])),
4670 true,
4671 )));
4672 assert_eq!(
4673 validate(
4674 list_of_struct.clone(),
4675 false,
4676 &[0, 1, 2, 3],
4677 vec![6, 6, 29, 29]
4678 )
4679 .unwrap(),
4680 6
4681 );
4682 assert_eq!(
4684 validate(list_of_struct, true, &[0, 1], vec![29, 29]).unwrap(),
4685 29
4686 );
4687
4688 let empty_struct = DataType::Struct(Fields::empty());
4690 assert_eq!(
4691 validate(empty_struct.clone(), false, &[0], vec![9]).unwrap(),
4692 9
4693 );
4694 assert_eq!(validate(empty_struct, true, &[0], vec![9]).unwrap(), 9);
4695 }
4696}