diff --git a/rust/lance-core/src/datatypes.rs b/rust/lance-core/src/datatypes.rs index 8837037c308..2f5c7b0680e 100644 --- a/rust/lance-core/src/datatypes.rs +++ b/rust/lance-core/src/datatypes.rs @@ -117,7 +117,7 @@ impl LogicalType { self.0.starts_with("fixed_size_list:struct:") } - fn is_struct(&self) -> bool { + pub fn is_struct(&self) -> bool { self.0 == "struct" } diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index d5355c089b2..0fc0baf0ee0 100644 --- a/rust/lance-file/src/reader.rs +++ b/rust/lance-file/src/reader.rs @@ -321,18 +321,15 @@ impl ReaderProjection { ) -> Result<()> { for field in fields { let is_structural = file_version >= LanceFileVersion::V2_1; + let (contributes, recurse) = field_column_shape(field, is_structural); // In the 2.0 system we needed ids for intermediate fields. In 2.1+ // we only need ids for leaf fields. - if (!is_structural - || field.children.is_empty() - || field.is_blob() - || field.is_packed_struct()) + if contributes && let Some(column_idx) = field_id_to_column_index.get(&(field.id as u32)).copied() { column_indices.push(column_idx); } - // Don't recurse into children if the field is a blob or packed struct in 2.1 - if !is_structural || (!field.is_blob() && !field.is_packed_struct()) { + if recurse { Self::from_field_ids_helper( file_version, field.children.iter(), @@ -520,6 +517,119 @@ struct Footer { const FOOTER_LEN: usize = 40; +// How a field maps onto physical columns, shared by the projection-building and +// projection-validation walks so they stay in lockstep. In the 2.0 layout every +// field (including structs and lists) has its own column; in 2.1 only leaves do, +// and blob/packed-struct fields are opaque (a single column with no descent). +// Returns `(contributes, recurse)`: whether the field has its own column and +// whether to walk into its children. The DFS order is the field's own column (if +// any) followed by its children, so a field's root (first) column is always the +// first entry of its sub-slice. +fn field_column_shape(field: &Field, is_structural: bool) -> (bool, bool) { + let contributes = + !is_structural || field.children.is_empty() || field.is_blob() || field.is_packed_struct(); + let recurse = !is_structural || (!field.is_blob() && !field.is_packed_struct()); + (contributes, recurse) +} + +// Whether a field's children each cover the same rows as the field itself. Struct +// children do (one value per parent row), so they must share its length. List, +// map, and fixed-size-list items have an independent cardinality (item count, not +// row count) and are validated only against themselves. +fn children_share_parent_length(field: &Field) -> bool { + field.logical_type.is_struct() +} + +// Validate one field's slice of a projection's flat `column_indices`, returning +// the field's top-level row count (the page-row sum of its root column). Walks the +// same DFS order as `from_field_ids_helper`, advancing `cursor` past every column +// the field contributes. +// +// `comparable` tracks whether the field's row count shares the read's top-level +// cardinality. A struct's children must all match that count -- the decoders +// combine them assuming equal lengths and would otherwise panic or read past a +// shorter child -- so the equality check runs only while `comparable` holds. Once +// the walk descends through a list/map/fixed-size-list its items have an +// independent cardinality (item count, not row count), so `comparable` turns off +// for that whole subtree and a nested struct's children are no longer compared. +fn validate_field_length Result>( + field: &Field, + is_structural: bool, + comparable: bool, + column_indices: &[u32], + cursor: &mut usize, + column_len: &F, +) -> Result { + let (contributes, recurse) = field_column_shape(field, is_structural); + let mut field_rows: Option = None; + if contributes { + let column = *column_indices.get(*cursor).ok_or_else(|| { + Error::invalid_input(format!( + "projection supplied fewer column indices than its fields require \ + (ran out at field '{}')", + field.name + )) + })?; + *cursor += 1; + field_rows = Some(column_len(column as usize)?); + } + if recurse { + // Only enforce equal-length children for a struct whose own count is still + // at the top-level cardinality; below a list/map/fixed-size-list the items + // have an independent cardinality, so neither this field nor its + // descendants are comparable. + let enforce_children = comparable && children_share_parent_length(field); + for child in &field.children { + let child_rows = validate_field_length( + child, + is_structural, + enforce_children, + column_indices, + cursor, + column_len, + )?; + // A struct that contributes no column of its own (the 2.1 layout) + // takes its row count from its first child. + let expected = *field_rows.get_or_insert(child_rows); + if enforce_children && child_rows != expected { + return Err(Error::invalid_input(format!( + "cannot read field '{}': its children have differing lengths \ + (child '{}' has {} rows, but the field has {}); a struct's \ + children must all have the same length", + field.name, child.name, child_rows, expected + ))); + } + } + } + field_rows.ok_or_else(|| { + Error::invalid_input(format!( + "projected field '{}' maps to no columns", + field.name + )) + }) +} + +// The reader combines a projection's columns into rectangular batches, so they +// must all have the same length. Returns that common length, or a descriptive +// error (naming each column's length) when they differ. Ordinary files always +// pass; only files written with `FileWriter::write_column` whose columns ended up +// unequal can fail, and those must be read separately. +fn verify_uniform_lengths(field_lengths: &[(&str, u64)]) -> Result { + let first = field_lengths.first().map_or(0, |&(_, len)| len); + if field_lengths.iter().all(|&(_, len)| len == first) { + return Ok(first); + } + let columns = field_lengths + .iter() + .map(|(name, len)| format!("{name}={len}")) + .collect::>() + .join(", "); + Err(Error::invalid_input(format!( + "cannot read columns of differing lengths together ({columns}); \ + read each column (or equal-length group) separately" + ))) +} + impl FileReader { pub fn with_scheduler(&self, scheduler: Arc) -> Self { Self { @@ -549,6 +659,29 @@ impl FileReader { self.core.num_rows() } + /// The number of rows stored in a single physical column. + /// + /// For ordinary (rectangular) files every column has the same length, equal + /// to [`num_rows`](Self::num_rows). Files written with + /// [`FileWriter::write_column`](crate::writer::FileWriter::write_column) + /// may have columns of differing lengths; this returns the length of one + /// such column, derived by summing its pages' row counts. Errors if + /// `column_index` is out of bounds. + pub fn column_num_rows(&self, column_index: usize) -> Result { + let column = self + .metadata + .column_metadatas + .get(column_index) + .ok_or_else(|| { + Error::invalid_input(format!( + "column index {} is out of bounds (file has {} columns)", + column_index, + self.metadata.column_metadatas.len() + )) + })?; + Ok(column.pages.iter().map(|page| page.length).sum()) + } + pub fn metadata(&self) -> &Arc { &self.metadata } @@ -1546,12 +1679,20 @@ impl FileReader { ) -> Result> { let projection = projection.unwrap_or_else(|| self.core.base_projection.clone()); Self::validate_projection(&projection, &self.metadata)?; + // Apply the same projection-length validation as the async path. This + // reader is always backed by full metadata, so we can build the prepared + // projection synchronously (no column-metadata I/O) and reuse the shared + // check. `read_len` is the projection's common column length, which + // `RangeFull`/`RangeFrom` resolve against rather than `num_rows`. + let prepared = PreparedProjection { + column_infos: self.metadata.column_infos.clone(), + decoder_projection: projection.clone(), + }; + let read_len = self.core.prepared_read_length(&prepared)?; let verify_bound = |params: &ReadBatchParams, bound: u64, inclusive: bool| { - if bound > self.core.num_rows() || bound == self.core.num_rows() && inclusive { + if bound > read_len || (bound == read_len && inclusive) { Err(Error::invalid_input(format!( - "cannot read {:?} from file with {} rows", - params, - self.core.num_rows() + "cannot read {params:?} from columns with {read_len} rows" ))) } else { Ok(()) @@ -1592,7 +1733,7 @@ impl FileReader { ReadBatchParams::RangeFrom(range) => { verify_bound(¶ms, range.start as u64, true)?; self.read_range_blocking( - range.start as u64..self.core.num_rows(), + range.start as u64..read_len, batch_size, projection, filter, @@ -1603,7 +1744,7 @@ impl FileReader { self.read_range_blocking(0..range.end as u64, batch_size, projection, filter) } ReadBatchParams::RangeFull => { - self.read_range_blocking(0..self.core.num_rows(), batch_size, projection, filter) + self.read_range_blocking(0..read_len, batch_size, projection, filter) } } } @@ -1939,17 +2080,64 @@ impl FileReadCore { }) } + // The common length to read across a prepared projection, after validating + // its columns can be combined into rectangular batches. Each top-level field + // is checked for internal consistency (see `validate_field_length`); the + // top-level fields must then share a length, since one read combines them. + // Ordinary files always pass (every column has `num_rows` rows); files + // written with `FileWriter::write_column` whose columns ended up unequal are + // rejected here and must be read separately. + // + // `column_infos` and `decoder_projection.column_indices` line up for both + // metadata providers: the full provider keeps absolute indices into the whole + // file, while the indexed (lazy) provider loads only the projected columns and + // renumbers them 0..N -- in either case `column_infos[column_index]` is the + // requested column. + fn prepared_read_length(&self, prepared: &PreparedProjection) -> Result { + let is_structural = self.version() >= LanceFileVersion::V2_1; + let column_infos = &prepared.column_infos; + let column_len = |column: usize| -> Result { + let info = column_infos.get(column).ok_or_else(|| { + Error::invalid_input(format!( + "projection references column index {} but only {} columns are available", + column, + column_infos.len() + )) + })?; + Ok(info.page_infos.iter().map(|page| page.num_rows).sum()) + }; + let column_indices = &prepared.decoder_projection.column_indices; + let fields = &prepared.decoder_projection.schema.fields; + let mut cursor = 0usize; + let mut field_lengths = Vec::with_capacity(fields.len()); + for field in fields { + let rows = validate_field_length( + field, + is_structural, + true, + column_indices, + &mut cursor, + &column_len, + )?; + field_lengths.push((field.name.as_str(), rows)); + } + if cursor != column_indices.len() { + return Err(Error::invalid_input(format!( + "projection supplied {} column indices but its fields require {}", + column_indices.len(), + cursor + ))); + } + verify_uniform_lengths(&field_lengths) + } + async fn read_range( &self, range: Range, batch_size: u32, - projection: ReaderProjection, + prepared: PreparedProjection, filter: FilterExpression, ) -> Result> { - let prepared = self - .metadata_provider - .prepare_projection(&projection, &self.scheduler, &self.cache) - .await?; FileReader::do_read_range( prepared.column_infos, self.scheduler.clone(), @@ -1970,12 +2158,8 @@ impl FileReadCore { &self, indices: Vec, batch_size: u32, - projection: ReaderProjection, + prepared: PreparedProjection, ) -> Result> { - let prepared = self - .metadata_provider - .prepare_projection(&projection, &self.scheduler, &self.cache) - .await?; FileReader::do_take_rows( prepared.column_infos, self.scheduler.clone(), @@ -1995,13 +2179,9 @@ impl FileReadCore { &self, ranges: Vec>, batch_size: u32, - projection: ReaderProjection, + prepared: PreparedProjection, filter: FilterExpression, ) -> Result> { - let prepared = self - .metadata_provider - .prepare_projection(&projection, &self.scheduler, &self.cache) - .await?; FileReader::do_read_ranges( prepared.column_infos, self.scheduler.clone(), @@ -2025,12 +2205,21 @@ impl FileReadCore { filter: FilterExpression, ) -> Result + Send>>> { let projection = projection.unwrap_or_else(|| self.base_projection.clone()); + let prepared = self + .metadata_provider + .prepare_projection(&projection, &self.scheduler, &self.cache) + .await?; + // All projected columns must share a length: the reader combines them + // into rectangular batches. Ordinary files satisfy this (every column + // has `num_rows` rows); files written with `FileWriter::write_column` + // may not, and such columns must be read separately. `read_len` is that + // common length, which `RangeFull`/`RangeFrom` resolve against (rather + // than `num_rows`, the file's longest column). + let read_len = self.prepared_read_length(&prepared)?; let verify_bound = |params: &ReadBatchParams, bound: u64, inclusive: bool| { - if bound > self.num_rows() || bound == self.num_rows() && inclusive { + if bound > read_len || (bound == read_len && inclusive) { Err(Error::invalid_input(format!( - "cannot read {:?} from file with {} rows", - params, - self.num_rows() + "cannot read {params:?} from columns with {read_len} rows" ))) } else { Ok(()) @@ -2049,14 +2238,14 @@ impl FileReadCore { } } let indices = indices.iter().map(|idx| idx.unwrap() as u64).collect(); - self.take_rows(indices, batch_size, projection).await + self.take_rows(indices, batch_size, prepared).await } ReadBatchParams::Range(range) => { verify_bound(¶ms, range.end as u64, false)?; self.read_range( range.start as u64..range.end as u64, batch_size, - projection, + prepared, filter, ) .await @@ -2067,26 +2256,21 @@ impl FileReadCore { verify_bound(¶ms, range.end, false)?; ranges_u64.push(range.start..range.end); } - self.read_ranges(ranges_u64, batch_size, projection, filter) + self.read_ranges(ranges_u64, batch_size, prepared, filter) .await } ReadBatchParams::RangeFrom(range) => { verify_bound(¶ms, range.start as u64, true)?; - self.read_range( - range.start as u64..self.num_rows(), - batch_size, - projection, - filter, - ) - .await + self.read_range(range.start as u64..read_len, batch_size, prepared, filter) + .await } ReadBatchParams::RangeTo(range) => { verify_bound(¶ms, range.end as u64, false)?; - self.read_range(0..range.end as u64, batch_size, projection, filter) + self.read_range(0..range.end as u64, batch_size, prepared, filter) .await } ReadBatchParams::RangeFull => { - self.read_range(0..self.num_rows(), batch_size, projection, filter) + self.read_range(0..read_len, batch_size, prepared, filter) .await } } @@ -2415,7 +2599,8 @@ mod tests { use tokio::sync::mpsc; use crate::reader::{ - EncodedBatchReaderExt, FileReader, FileReaderOptions, ProjectedFileReader, ReaderProjection, + EncodedBatchReaderExt, FileReader, FileReaderOptions, ProjectedFileReader, + ReaderProjection, validate_field_length, verify_uniform_lengths, }; use crate::testing::{FsFixture, WrittenFile, test_cache, write_lance_file}; use crate::writer::{EncodedBatchWriteExt, FileWriter, FileWriterOptions}; @@ -3042,6 +3227,113 @@ mod tests { ); } + // The projection-length validation lives in `FileReadCore`, shared by the + // eager and the lazy (indexed) metadata providers. The indexed provider loads + // only the projected columns and renumbers them 0..N, so this checks that the + // renumbered `column_infos`/`column_indices` still line up for the length + // check: a mismatched-length projection is rejected through + // `ProjectedFileReader`, and a single short column resolves to its own length. + #[tokio::test] + async fn test_lazy_reader_validates_unequal_length_projection() { + use arrow_array::Int32Array; + use lance_io::ReadBatchParams; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ])); + let lance_schema = Schema::try_from(arrow_schema.as_ref()).unwrap(); + + let fs = FsFixture::default(); + let options = FileWriterOptions { + format_version: Some(LanceFileVersion::V2_1), + ..Default::default() + }; + let mut writer = FileWriter::try_new( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema.clone(), + options, + ) + .unwrap(); + // "a" has 5 rows, "c" has 1 -- an unequal-length file. + writer + .write_column(0, Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))) + .await + .unwrap(); + writer + .write_column(1, Arc::new(Int32Array::from(vec![100]))) + .await + .unwrap(); + writer.finish().await.unwrap(); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let cache = test_cache(); + let open_indexed = |names: &[&str]| { + let projection = + ReaderProjection::from_column_names(LanceFileVersion::V2_1, &lance_schema, names) + .unwrap(); + ProjectedFileReader::try_open( + file_scheduler.clone(), + Some(projection), + Arc::::default(), + &cache, + FileReaderOptions::default(), + ) + }; + + // A mismatched-length projection [a, c] (5 vs 1) is rejected at read time, + // through the indexed provider's renumbered column infos. + let lazy = open_indexed(&["a", "c"]).await.unwrap(); + // (`read_tasks` yields a stream, which is not `Debug`, so match rather + // than `unwrap_err`.) + let err = match lazy + .read_tasks( + ReadBatchParams::RangeFull, + 1024, + None, + FilterExpression::no_filter(), + ) + .await + { + Ok(_) => panic!("expected the mismatched-length projection to be rejected"), + Err(e) => e.to_string(), + }; + assert!( + err.contains("a=5") && err.contains("c=1"), + "error should name each column's length, got: {err}" + ); + + // A single short column resolves to its own length (1), not the file's + // longest column. + let lazy = open_indexed(&["c"]).await.unwrap(); + let tasks = lazy + .read_tasks( + ReadBatchParams::RangeFull, + 1024, + None, + FilterExpression::no_filter(), + ) + .await + .unwrap(); + let batches = collect_read_tasks(tasks, 16).await; + let values: Vec> = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .collect::>() + }) + .collect(); + assert_eq!(values, vec![Some(100)]); + } + #[test_log::test(tokio::test)] async fn test_compressing_buffer() { let fs = FsFixture::default(); @@ -3540,4 +3832,126 @@ mod tests { let msg = err.to_string(); assert!(msg.contains('2'), "error should mention the index: {msg}"); } + + // Exercises the projection length-validation walk in isolation, feeding + // synthetic per-column lengths so we can reach cases no current writer can + // actually produce -- in particular a struct whose children diverge in + // length, which the decoders would otherwise panic on or misread. + #[rstest] + fn test_validate_struct_child_lengths(#[values(false, true)] is_structural: bool) { + let run = |dt: DataType, indices: &[u32], lengths: Vec| -> lance_core::Result { + let arrow = ArrowSchema::new(vec![Field::new("s", dt, true)]); + let schema = Schema::try_from(&arrow).unwrap(); + let column_len = |c: usize| Ok(lengths[c]); + let mut cursor = 0usize; + let mut field_lengths = Vec::new(); + for field in &schema.fields { + let rows = validate_field_length( + field, + is_structural, + true, + indices, + &mut cursor, + &column_len, + )?; + field_lengths.push((field.name.as_str(), rows)); + } + verify_uniform_lengths(&field_lengths) + }; + + let struct_ty = || { + DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])) + }; + + // In 2.1 a struct contributes no column of its own (just its two leaves); + // in 2.0 it also has its own column first. + let (indices, equal, unequal): (&[u32], Vec, Vec) = if is_structural { + (&[0, 1], vec![5, 5], vec![5, 3]) + } else { + (&[0, 1, 2], vec![5, 5, 5], vec![5, 5, 3]) + }; + + assert_eq!(run(struct_ty(), indices, equal).unwrap(), 5); + + let err = run(struct_ty(), indices, unequal).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("differing lengths") && msg.contains('b'), + "expected a child-length error naming 'b', got: {msg}" + ); + } + + #[test] + fn test_validate_length_list_and_empty_struct() { + let validate = |dt: DataType, + is_structural: bool, + indices: &[u32], + lengths: Vec| + -> lance_core::Result { + let arrow = ArrowSchema::new(vec![Field::new("f", dt, true)]); + let schema = Schema::try_from(&arrow).unwrap(); + let column_len = |c: usize| Ok(lengths[c]); + let mut cursor = 0usize; + validate_field_length( + &schema.fields[0], + is_structural, + true, + indices, + &mut cursor, + &column_len, + ) + }; + + // A list's items have a different cardinality than its rows; that gap + // must not be flagged as a mismatch. In 2.0 the list is an offsets column + // (rows) plus an items column (item count); in 2.1 it is a single column + // whose page rows are the top-level row count. + let list_ty = DataType::List(Arc::new(Field::new("item", DataType::Int32, true))); + assert_eq!( + validate(list_ty.clone(), false, &[0, 1], vec![5, 17]).unwrap(), + 5 + ); + assert_eq!(validate(list_ty, true, &[0], vec![5]).unwrap(), 5); + + // A list of structs: the struct's children sit below the list boundary, so + // their (item-count) lengths must NOT be compared against the list's row + // count. In 2.0 each field has a column [list, struct, a, b] = [6, 6, 29, + // 29]; the list resolves to its own row count (6) and the struct's longer + // children are not flagged. (Regression: this previously errored because + // the nested struct's children were checked against the struct's count.) + let list_of_struct = DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])), + true, + ))); + assert_eq!( + validate( + list_of_struct.clone(), + false, + &[0, 1, 2, 3], + vec![6, 6, 29, 29] + ) + .unwrap(), + 6 + ); + // In 2.1 only the two leaves carry columns; still no false mismatch. + assert_eq!( + validate(list_of_struct, true, &[0, 1], vec![29, 29]).unwrap(), + 29 + ); + + // An empty struct contributes a single column and validates to its length. + let empty_struct = DataType::Struct(Fields::empty()); + assert_eq!( + validate(empty_struct.clone(), false, &[0], vec![9]).unwrap(), + 9 + ); + assert_eq!(validate(empty_struct, true, &[0], vec![9]).unwrap(), 9); + } } diff --git a/rust/lance-file/src/writer.rs b/rust/lance-file/src/writer.rs index 12bd50df6fe..f7042b6b99c 100644 --- a/rust/lance-file/src/writer.rs +++ b/rust/lance-file/src/writer.rs @@ -6,7 +6,7 @@ use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::AtomicBool; -use arrow_array::RecordBatch; +use arrow_array::{ArrayRef, RecordBatch}; use arrow_data::ArrayData; use bytes::{Buf, BufMut, Bytes, BytesMut}; @@ -221,6 +221,11 @@ pub struct FileWriter { field_id_to_column_indices: Vec<(u32, u32)>, num_columns: u32, rows_written: u64, + // The number of rows written for each top-level field (i.e. each entry in + // `column_writers`). With `write_batch` every field advances together and + // these are all equal, but `write_column` advances one field at a time, so + // a single file may end up with columns of differing item counts. + field_rows_written: Vec, global_buffers: Vec<(u64, u64)>, schema_metadata: HashMap, options: FileWriterOptions, @@ -277,6 +282,7 @@ impl FileWriter { column_metadata: Vec::new(), num_columns: 0, rows_written: 0, + field_rows_written: Vec::new(), field_id_to_column_indices: Vec::new(), global_buffers: Vec::new(), schema_metadata: HashMap::new(), @@ -467,6 +473,7 @@ impl FileWriter { BatchEncoder::try_new(&schema, encoding_strategy.as_ref(), &encoding_options)?; self.num_columns = encoder.num_columns(); + self.field_rows_written = vec![0; encoder.field_encoders.len()]; self.column_writers = encoder.field_encoders; self.column_metadata = vec![initial_column_metadata(); self.num_columns as usize]; self.field_id_to_column_indices = encoder.field_id_to_column_index; @@ -490,13 +497,14 @@ impl FileWriter { batch: &RecordBatch, external_buffers: &mut OutOfLineBuffers, ) -> Result>> { - self.schema + let field_arrays = self + .schema .as_ref() .unwrap() .fields .iter() - .zip(self.column_writers.iter_mut()) - .map(|(field, column_writer)| { + .enumerate() + .map(|(field_idx, field)| { let array = batch .column_by_name(&field.name) @@ -507,19 +515,57 @@ impl FileWriter { ) .into(), ))?; + Ok((field_idx, array.clone())) + }) + .collect::>>()?; + self.encode_columns(&field_arrays, external_buffers) + } + + // Encode a set of `(field index, array)` pairs, each advancing only its own + // column. Each task captures its field's current row offset at encode time, + // so `advance_columns` must run after this call (never before); the order of + // the returned tasks relative to `write_pages` does not matter. + fn encode_columns( + &mut self, + field_arrays: &[(usize, ArrayRef)], + external_buffers: &mut OutOfLineBuffers, + ) -> Result>> { + // Snapshot the starting row number of each field before borrowing the + // column writers mutably below. + let row_numbers = field_arrays + .iter() + .map(|(field_idx, _)| self.field_rows_written[*field_idx]) + .collect::>(); + field_arrays + .iter() + .zip(row_numbers) + .map(|((field_idx, array), row_number)| { let repdef = RepDefBuilder::default(); let num_rows = array.len() as u64; - column_writer.maybe_encode( + self.column_writers[*field_idx].maybe_encode( array.clone(), external_buffers, repdef, - self.rows_written, + row_number, num_rows, ) }) .collect::>>() } + // Advance the per-field row counters after a set of columns has been + // written, keeping `rows_written` (the file's logical length) in sync as the + // longest column. Only the written fields move, so their new totals fold into + // `rows_written` directly without rescanning every field. (`write_batch` + // advances every field uniformly and tracks this inline instead.) + fn advance_columns(&mut self, field_arrays: &[(usize, ArrayRef)]) { + for (field_idx, array) in field_arrays { + let new_total = self.field_rows_written[*field_idx] + array.len() as u64; + self.field_rows_written[*field_idx] = new_total; + self.rows_written = self.rows_written.max(new_total); + } + } + /// Schedule a batch of data to be written to the file /// /// Note: the future returned by this method may complete before the data has been fully @@ -557,18 +603,97 @@ impl FileWriter { .flatten() .collect::>(); - self.rows_written = match self.rows_written.checked_add(batch.num_rows() as u64) { - Some(rows_written) => rows_written, - None => { - return Err(Error::invalid_input_source(format!("cannot write batch with {} rows because {} rows have already been written and Lance files cannot contain more than 2^64 rows", num_rows, self.rows_written).into())); - } - }; + // `write_batch` advances every field by the same amount, so the longest + // column simply grows by `num_rows`. Guard against overflowing the row + // counter. + if self.rows_written.checked_add(num_rows).is_none() { + return Err(Error::invalid_input_source(format!("cannot write batch with {} rows because {} rows have already been written and Lance files cannot contain more than 2^64 rows", num_rows, self.rows_written).into())); + } + for field_rows in self.field_rows_written.iter_mut() { + *field_rows += num_rows; + } + self.rows_written += num_rows; self.write_pages(encoding_tasks).await?; Ok(()) } + /// Write a single column, advancing only that column's row counter. + /// + /// Unlike [`write_batch`](Self::write_batch), which advances every column + /// from a single shared row counter, this method advances one column + /// independently. Used across calls it produces a single file whose columns + /// may have different item counts. + /// + /// `column_index` refers to a top-level field in the writer's schema (the + /// same order as the schema's fields); a nested child cannot be targeted on + /// its own. Because each call writes the whole field from a single array, the + /// children of a struct field always advance together and stay equal-length; + /// only different top-level fields can diverge in length. A column may be + /// written across multiple calls; its values are appended. A field that is + /// never written ends up as a zero-length column. The writer must have been + /// created with an explicit schema (via [`try_new`](Self::try_new)); a lazy + /// schema cannot be inferred here because individual calls need not cover + /// every field. + /// + /// ``` + /// # use arrow_array::{ArrayRef, Int32Array}; + /// # use std::sync::Arc; + /// # use lance_file::writer::FileWriter; + /// # async fn example(writer: &mut FileWriter) -> lance_core::Result<()> { + /// // Field 0 gets three values, field 1 gets one — a non-rectangular file. + /// writer.write_column(0, Arc::new(Int32Array::from(vec![1, 2, 3]))).await?; + /// writer.write_column(1, Arc::new(Int32Array::from(vec![10]))).await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn write_column(&mut self, column_index: usize, array: ArrayRef) -> Result<()> { + let schema = self.schema.as_ref().ok_or_else(|| { + Error::invalid_input_source( + "write_column requires the writer to be created with an explicit schema".into(), + ) + })?; + let field = schema.fields.get(column_index).ok_or_else(|| { + Error::invalid_input_source( + format!( + "write_column: field index {} is out of bounds (schema has {} fields)", + column_index, + schema.fields.len() + ) + .into(), + ) + })?; + if array.len() as u64 > u32::MAX as u64 { + return Err(Error::invalid_input_source( + "cannot write Lance files with more than 2^32 rows".into(), + )); + } + Self::verify_field_nullability(&array.to_data(), field)?; + + // A never-advanced field simply remains a zero-length column, which the + // encoders handle at `finish` time. + if array.is_empty() { + return Ok(()); + } + + let columns = [(column_index, array)]; + let mut external_buffers = + OutOfLineBuffers::new(self.tell().await?, PAGE_BUFFER_ALIGNMENT as u64); + let encoding_tasks = self.encode_columns(&columns, &mut external_buffers)?; + for external_buffer in external_buffers.take_buffers() { + Self::do_write_buffer(&mut self.writer, &external_buffer).await?; + } + let encoding_tasks = encoding_tasks + .into_iter() + .flatten() + .collect::>(); + + self.advance_columns(&columns); + self.write_pages(encoding_tasks).await?; + Ok(()) + } + async fn write_column_metadata( &mut self, metadata: pbfile::ColumnMetadata, @@ -974,11 +1099,11 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; - use crate::reader::{FileReader, FileReaderOptions, describe_encoding}; + use crate::reader::{FileReader, FileReaderOptions, ReaderProjection, describe_encoding}; use crate::testing::FsFixture; use crate::writer::{ENV_LANCE_FILE_WRITER_MAX_PAGE_BYTES, FileWriter, FileWriterOptions}; use arrow_array::builder::{Float32Builder, Int32Builder}; - use arrow_array::{Int32Array, RecordBatch, UInt64Array}; + use arrow_array::{ArrayRef, Int32Array, RecordBatch, UInt64Array}; use arrow_array::{RecordBatchReader, StringArray, types::Float64Type}; use arrow_schema::{DataType, Field, Field as ArrowField, Schema, Schema as ArrowSchema}; use lance_core::cache::LanceCache; @@ -990,6 +1115,7 @@ mod tests { use lance_encoding::version::LanceFileVersion; use lance_io::object_store::ObjectStore; use lance_io::utils::CachedFileSize; + use rstest::rstest; #[tokio::test] async fn test_basic_write() { @@ -1040,6 +1166,611 @@ mod tests { file_writer.finish().await.unwrap(); } + // Read a single column back at an explicit range/index set, returning its + // `Int32` values. Reading one column (or an equal-length group) at a time is + // how unequal-length files are consumed: a full scan across columns of + // differing lengths cannot form a single rectangular batch. + async fn read_int32_column( + reader: &FileReader, + schema: &LanceSchema, + version: LanceFileVersion, + name: &str, + params: lance_io::ReadBatchParams, + ) -> Vec> { + use futures::TryStreamExt; + use lance_encoding::decoder::FilterExpression; + + let projection = ReaderProjection::from_column_names(version, schema, &[name]).unwrap(); + let batches: Vec = reader + .read_stream_projected(params, 1024, 16, projection, FilterExpression::no_filter()) + .await + .unwrap() + .try_collect() + .await + .unwrap(); + batches + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .collect::>() + }) + .collect() + } + + /// A single file may hold columns of differing item counts, written by + /// advancing each column's row counter independently (no shared global + /// counter). + #[rstest] + #[tokio::test] + async fn test_write_columns_unequal_lengths( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + use lance_io::ReadBatchParams; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ArrowField::new("c", DataType::Int32, true), + ])); + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + let fs = FsFixture::default(); + let options = FileWriterOptions { + format_version: Some(version), + ..Default::default() + }; + let mut writer = FileWriter::try_new( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema.clone(), + options, + ) + .unwrap(); + + // Field "a" gets 5 values across two calls (appending), field "b" gets a + // single value, and field "c" is never written (a zero-length column). + let a1: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let b: ArrayRef = Arc::new(Int32Array::from(vec![10])); + writer.write_column(0, a1).await.unwrap(); + writer.write_column(1, b).await.unwrap(); + let a2: ArrayRef = Arc::new(Int32Array::from(vec![4, 5])); + writer.write_column(0, a2).await.unwrap(); + // An empty array is a no-op whether or not the field already has rows: + // field "a" keeps its 5 rows, field "c" stays a zero-length column. + let empty: ArrayRef = Arc::new(Int32Array::from(Vec::::new())); + writer.write_column(0, empty.clone()).await.unwrap(); + writer.write_column(2, empty).await.unwrap(); + + let summary = writer.finish().await.unwrap(); + // The file's logical length is the longest column. + assert_eq!(summary.num_rows, 5); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + + // Per-column row counts are recorded in / derivable from file metadata. + assert_eq!(reader.num_rows(), 5); + assert_eq!(reader.column_num_rows(0).unwrap(), 5); + assert_eq!(reader.column_num_rows(1).unwrap(), 1); + assert_eq!(reader.column_num_rows(2).unwrap(), 0); + assert!(reader.column_num_rows(3).is_err()); + + // Each column reads back independently at its own length. + assert_eq!( + read_int32_column( + &reader, + &lance_schema, + version, + "a", + ReadBatchParams::Range(0..5) + ) + .await, + vec![Some(1), Some(2), Some(3), Some(4), Some(5)], + ); + assert_eq!( + read_int32_column( + &reader, + &lance_schema, + version, + "b", + ReadBatchParams::Range(0..1) + ) + .await, + vec![Some(10)], + ); + + // Random access by position within the longer column returns the right + // value even though other columns are shorter. (The take path requires + // strictly increasing indices.) + assert_eq!( + read_int32_column( + &reader, + &lance_schema, + version, + "a", + ReadBatchParams::Indices(arrow_array::UInt32Array::from(vec![0, 2, 4])), + ) + .await, + vec![Some(1), Some(3), Some(5)], + ); + } + + /// Reading an unequal-length file: + /// - a projection whose columns are equal length full-scans normally; + /// - a full scan across columns of differing length is rejected up front, + /// before any batch is produced (even though a prefix would be rectangular); + /// - a bounded read is valid as long as every projected column covers it; + /// - a single-column `RangeFull` resolves to that column's own length, not + /// the file's (maximum) length. + #[rstest] + #[tokio::test] + async fn test_read_unequal_length_projection( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + use futures::TryStreamExt; + use lance_encoding::decoder::FilterExpression; + use lance_io::ReadBatchParams; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ArrowField::new("c", DataType::Int32, true), + ])); + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + let fs = FsFixture::default(); + let options = FileWriterOptions { + format_version: Some(version), + ..Default::default() + }; + let mut writer = FileWriter::try_new( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema.clone(), + options, + ) + .unwrap(); + // "a" and "b" are equal length (5); "c" is shorter (1). + writer + .write_column(0, Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))) + .await + .unwrap(); + writer + .write_column(1, Arc::new(Int32Array::from(vec![6, 7, 8, 9, 10]))) + .await + .unwrap(); + writer + .write_column(2, Arc::new(Int32Array::from(vec![100]))) + .await + .unwrap(); + writer.finish().await.unwrap(); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + + let read = |names: &'static [&'static str], params: ReadBatchParams| { + let projection = + ReaderProjection::from_column_names(version, &lance_schema, names).unwrap(); + async { + match reader + .read_stream_projected( + params, + 1024, + 16, + projection, + FilterExpression::no_filter(), + ) + .await + { + Ok(stream) => stream.try_collect::>().await, + Err(e) => Err(e), + } + } + }; + let col_values = |batches: &[RecordBatch], idx: usize| -> Vec> { + batches + .iter() + .flat_map(|b| { + b.column(idx) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .collect::>() + }) + .collect() + }; + + // Equal-length projection [a, b] full-scans into rectangular batches. + let batches = read(&["a", "b"], ReadBatchParams::RangeFull).await.unwrap(); + assert_eq!( + col_values(&batches, 0), + vec![Some(1), Some(2), Some(3), Some(4), Some(5)] + ); + assert_eq!( + col_values(&batches, 1), + vec![Some(6), Some(7), Some(8), Some(9), Some(10)] + ); + + // A mismatched-length projection [a, c] (5 vs 1) is rejected before any + // batch is yielded, regardless of the read params — its columns cannot + // be combined into rectangular batches. The error names each column's + // length so the caller can see which column is the odd one out. + let err = read(&["a", "c"], ReadBatchParams::RangeFull) + .await + .unwrap_err() + .to_string(); + assert!( + err.contains("a=5") && err.contains("c=1"), + "error should name each column's length, got: {err}" + ); + assert!( + read(&["a", "c"], ReadBatchParams::Range(0..1)) + .await + .is_err(), + "even a common-prefix read of unequal-length columns must error" + ); + + // A single-column RangeFull resolves to that column's own length. + let batches = read(&["c"], ReadBatchParams::RangeFull).await.unwrap(); + assert_eq!(col_values(&batches, 0), vec![Some(100)]); + let batches = read(&["a"], ReadBatchParams::RangeFull).await.unwrap(); + assert_eq!( + col_values(&batches, 0), + vec![Some(1), Some(2), Some(3), Some(4), Some(5)] + ); + + // RangeFrom/RangeTo likewise resolve against the projected column's own + // length rather than the file's longest column. + let batches = read(&["a"], ReadBatchParams::RangeFrom(2..)).await.unwrap(); + assert_eq!(col_values(&batches, 0), vec![Some(3), Some(4), Some(5)]); + // RangeFrom on the short column "c" resolves to length 1, not 5. + let batches = read(&["c"], ReadBatchParams::RangeFrom(0..)).await.unwrap(); + assert_eq!(col_values(&batches, 0), vec![Some(100)]); + let batches = read(&["a"], ReadBatchParams::RangeTo(..3)).await.unwrap(); + assert_eq!(col_values(&batches, 0), vec![Some(1), Some(2), Some(3)]); + // A bound past the projected column's length errors. + assert!( + read(&["a"], ReadBatchParams::RangeTo(..6)).await.is_err(), + "RangeTo past the column length must error" + ); + assert!( + read(&["c"], ReadBatchParams::RangeFrom(2..)).await.is_err(), + "RangeFrom past the column length must error" + ); + } + + /// A struct and a list column each map to multiple physical columns, and a + /// list's item column is longer than its top-level row count. The + /// projection-length check must partition `column_indices` by top-level + /// field and use each field's root column, so an ordinary (rectangular) file + /// with nested columns still reads under the new validation path. + #[rstest] + #[tokio::test] + async fn test_read_nested_columns_under_validation( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + use arrow_array::types::Int32Type; + use arrow_array::{ListArray, StructArray}; + use futures::TryStreamExt; + use lance_encoding::decoder::FilterExpression; + use lance_io::ReadBatchParams; + + let struct_type = DataType::Struct( + vec![ + ArrowField::new("x", DataType::Int32, true), + ArrowField::new("y", DataType::Int32, true), + ] + .into(), + ); + let list_type = DataType::List(Arc::new(ArrowField::new("item", DataType::Int32, true))); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("s", struct_type, true), + ArrowField::new("lst", list_type, true), + ])); + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let s: ArrayRef = Arc::new(StructArray::from(vec![ + ( + Arc::new(ArrowField::new("x", DataType::Int32, true)), + Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("y", DataType::Int32, true)), + Arc::new(Int32Array::from(vec![11, 21, 31])) as ArrayRef, + ), + ])); + // 3 lists, 6 items: the item column is longer than the top-level rows. + let lst: ArrayRef = Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2)]), + Some(vec![Some(3)]), + Some(vec![Some(4), Some(5), Some(6)]), + ])); + let batch = RecordBatch::try_new(arrow_schema.clone(), vec![a, s, lst]).unwrap(); + + let fs = FsFixture::default(); + let options = FileWriterOptions { + format_version: Some(version), + ..Default::default() + }; + let mut writer = FileWriter::try_new( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema.clone(), + options, + ) + .unwrap(); + writer.write_batch(&batch).await.unwrap(); + writer.finish().await.unwrap(); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + + // If `validate_field_length` mispartitioned the physical columns, the + // length check would read the wrong root column (e.g. the list's item + // column, length 6) and spuriously reject this rectangular file. + for names in [&["a", "s", "lst"][..], &["a", "lst"][..], &["a", "s"][..]] { + let projection = + ReaderProjection::from_column_names(version, &lance_schema, names).unwrap(); + let batches: Vec = reader + .read_stream_projected( + ReadBatchParams::RangeFull, + 1024, + 16, + projection, + FilterExpression::no_filter(), + ) + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total_rows, 3, + "projection {names:?} should read 3 top-level rows" + ); + } + } + + /// `write_column` rejects invalid inputs at the API boundary with + /// descriptive errors: a writer without an explicit schema, an + /// out-of-bounds field index, and a null written into a non-nullable field. + #[tokio::test] + async fn test_write_column_validation_errors() { + // A lazy-schema writer cannot infer the schema from a single column. + let fs = FsFixture::default(); + let mut lazy_writer = FileWriter::new_lazy( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + FileWriterOptions::default(), + ); + let err = lazy_writer + .write_column(0, Arc::new(Int32Array::from(vec![1, 2, 3]))) + .await + .unwrap_err() + .to_string(); + assert!( + err.contains("explicit schema"), + "expected explicit-schema error, got: {err}" + ); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, false), + ArrowField::new("b", DataType::Int32, true), + ])); + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + // An out-of-bounds field index is rejected, naming the index and count. + let fs = FsFixture::default(); + let mut writer = FileWriter::try_new( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema.clone(), + FileWriterOptions::default(), + ) + .unwrap(); + let err = writer + .write_column(5, Arc::new(Int32Array::from(vec![1]))) + .await + .unwrap_err() + .to_string(); + assert!( + err.contains('5') && err.contains('2'), + "expected out-of-bounds error naming index 5 and 2 fields, got: {err}" + ); + + // A null in a non-nullable field ("a") is rejected. + let err = writer + .write_column(0, Arc::new(Int32Array::from(vec![Some(1), None, Some(3)]))) + .await + .unwrap_err() + .to_string(); + assert!( + err.contains("non-null"), + "expected nullability error, got: {err}" + ); + } + + /// The blocking read path applies the same projection-length validation as + /// the async path: a short single column resolves to its own length, and a + /// mismatched-length projection errors up front. + #[rstest] + #[tokio::test] + async fn test_blocking_read_unequal_length( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + use lance_encoding::decoder::FilterExpression; + use lance_io::ReadBatchParams; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("c", DataType::Int32, true), + ])); + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + let fs = FsFixture::default(); + let options = FileWriterOptions { + format_version: Some(version), + ..Default::default() + }; + let mut writer = FileWriter::try_new( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema.clone(), + options, + ) + .unwrap(); + writer + .write_column(0, Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))) + .await + .unwrap(); + writer + .write_column(1, Arc::new(Int32Array::from(vec![100]))) + .await + .unwrap(); + writer.finish().await.unwrap(); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let reader = Arc::new( + FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(), + ); + + // Single short column: RangeFull resolves to its own length (1). + let proj_c = ReaderProjection::from_column_names(version, &lance_schema, &["c"]).unwrap(); + let reader_c = reader.clone(); + let batches = tokio::task::spawn_blocking(move || { + reader_c + .read_stream_projected_blocking( + ReadBatchParams::RangeFull, + 1024, + Some(proj_c), + FilterExpression::no_filter(), + ) + .unwrap() + .collect::, _>>() + .unwrap() + }) + .await + .unwrap(); + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 1); + + // A mismatched projection [a, c] errors on the blocking path too. + let proj_ac = + ReaderProjection::from_column_names(version, &lance_schema, &["a", "c"]).unwrap(); + let reader_ac = reader.clone(); + let is_err = tokio::task::spawn_blocking(move || { + reader_ac + .read_stream_projected_blocking( + ReadBatchParams::RangeFull, + 1024, + Some(proj_ac), + FilterExpression::no_filter(), + ) + .is_err() + }) + .await + .unwrap(); + assert!( + is_err, + "blocking full scan across unequal-length columns must error" + ); + } + + /// Files written the ordinary (rectangular) way keep equal column lengths, + /// so the unequal-length support is backwards compatible. + #[tokio::test] + async fn test_write_batch_keeps_equal_lengths() { + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ])); + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + let fs = FsFixture::default(); + let mut writer = FileWriter::try_new( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema, + FileWriterOptions::default(), + ) + .unwrap(); + let batch = RecordBatch::try_new( + arrow_schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![4, 5, 6])), + ], + ) + .unwrap(); + writer.write_batch(&batch).await.unwrap(); + let summary = writer.finish().await.unwrap(); + assert_eq!(summary.num_rows, 3); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + assert_eq!(reader.column_num_rows(0).unwrap(), 3); + assert_eq!(reader.column_num_rows(1).unwrap(), 3); + } + #[tokio::test] async fn test_max_page_bytes_enforced() { let arrow_field = Field::new("data", DataType::UInt64, false);