diff --git a/rust/lance-encoding/src/array_encoding/logical/binary.rs b/rust/lance-encoding/src/array_encoding/logical/binary.rs index 697c1503e0a..1bb900a6785 100644 --- a/rust/lance-encoding/src/array_encoding/logical/binary.rs +++ b/rust/lance-encoding/src/array_encoding/logical/binary.rs @@ -11,7 +11,7 @@ use arrow_array::{ use arrow_schema::DataType; use futures::{FutureExt, future::BoxFuture}; -use lance_core::Result; +use lance_core::{Error, Result}; use log::trace; use crate::{ @@ -152,7 +152,7 @@ pub struct BinaryArrayDecoder { } impl BinaryArrayDecoder { - fn from_list_array(array: &GenericListArray) -> ArrayRef { + fn from_list_array(array: &GenericListArray) -> Result { let values = array .values() .as_primitive::() @@ -160,11 +160,8 @@ impl BinaryArrayDecoder { .inner() .clone(); let offsets = array.offsets().clone(); - Arc::new(GenericByteArray::::new( - offsets, - values, - array.nulls().cloned(), - )) + let array = GenericByteArray::::try_new(offsets, values, array.nulls().cloned())?; + Ok(Arc::new(array)) } } @@ -173,14 +170,89 @@ impl DecodeArrayTask for BinaryArrayDecoder { let data_type = self.data_type; let (arr, _) = self.inner.decode()?; let result = match data_type { - DataType::Binary => Self::from_list_array::(arr.as_list::()), - DataType::LargeBinary => Self::from_list_array::(arr.as_list::()), - DataType::Utf8 => Self::from_list_array::(arr.as_list::()), - DataType::LargeUtf8 => Self::from_list_array::(arr.as_list::()), - _ => panic!("Binary decoder does not support this data type"), + DataType::Binary => Self::from_list_array::(arr.as_list::())?, + DataType::LargeBinary => { + Self::from_list_array::(arr.as_list::())? + } + DataType::Utf8 => Self::from_list_array::(arr.as_list::())?, + DataType::LargeUtf8 => Self::from_list_array::(arr.as_list::())?, + other => { + return Err(Error::internal(format!( + "Binary decoder does not support data type {other}" + ))); + } }; // data_size is only tracked in the v2.1 structural decode path; the v2.0 array // v2.0 path does not need it so we return 0. Ok((result, 0)) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::{ListArray, UInt8Array}; + use arrow_buffer::OffsetBuffer; + use arrow_schema::Field; + + use super::*; + use crate::decoder::DecodeArrayTask; + + struct StubDecodeTask { + array: ArrayRef, + } + + impl DecodeArrayTask for StubDecodeTask { + fn decode(self: Box) -> Result<(ArrayRef, u64)> { + Ok((self.array, 0)) + } + } + + fn make_single_byte_list(value: u8) -> ListArray { + let offsets = OffsetBuffer::from_lengths([1_usize]); + let values: ArrayRef = Arc::new(UInt8Array::from(vec![value])); + ListArray::try_new( + Arc::new(Field::new("item", DataType::UInt8, false)), + offsets, + values, + None, + ) + .unwrap() + } + + #[test] + fn logical_utf8_decode_preserves_non_overflow_arrow_error() { + let list = make_single_byte_list(0xFF_u8); + let decoder = BinaryArrayDecoder { + inner: Box::new(StubDecodeTask { + array: Arc::new(list), + }), + data_type: DataType::Utf8, + }; + + let error = Box::new(decoder).decode().unwrap_err(); + let message = error.to_string(); + assert!(!message.contains("more than 2GiB of string/binary data")); + assert!(message.to_lowercase().contains("utf")); + } + + #[test] + fn logical_binary_decode_returns_internal_error_for_unsupported_type() { + let list = make_single_byte_list(b'x'); + let decoder = BinaryArrayDecoder { + inner: Box::new(StubDecodeTask { + array: Arc::new(list), + }), + data_type: DataType::Int32, + }; + + let error = Box::new(decoder).decode().unwrap_err(); + assert!(matches!(error, Error::Internal { .. })); + assert!( + error + .to_string() + .contains("Binary decoder does not support data type Int32") + ); + } +} diff --git a/rust/lance-encoding/src/array_encoding/logical/list.rs b/rust/lance-encoding/src/array_encoding/logical/list.rs index d76532ab3bb..60e32301fdc 100644 --- a/rust/lance-encoding/src/array_encoding/logical/list.rs +++ b/rust/lance-encoding/src/array_encoding/logical/list.rs @@ -637,6 +637,51 @@ struct ListDecodeTask { offset_type: DataType, } +fn oversized_batch_error( + items_field: &Field, + requested_range: Range, + decodable_prefix_end: u64, + num_items: u64, +) -> Error { + let prefix_detail = if decodable_prefix_end > requested_range.start { + format!( + "rows {}..{} fit, but requesting rows {}..{} would require {num_items} {}", + requested_range.start, + decodable_prefix_end, + requested_range.start, + requested_range.end, + if items_field.data_type() == &DataType::UInt8 { + "bytes" + } else { + "items" + } + ) + } else { + format!( + "requesting rows {}..{} would require {num_items} {}", + requested_range.start, + requested_range.end, + if items_field.data_type() == &DataType::UInt8 { + "bytes" + } else { + "items" + } + ) + }; + if items_field.data_type() == &DataType::UInt8 { + Error::not_supported(format!( + "Could not create array with more than 2GiB of string/binary data in a single batch \ + ({prefix_detail}). Please reduce the batch_size, set LANCE_DEFAULT_BATCH_SIZE to a \ + smaller value, or convert the column to large_string/large_binary." + )) + } else { + Error::not_supported(format!( + "Could not create a list array with more than i32::MAX items in a single batch \ + ({prefix_detail}). Please reduce the batch_size." + )) + } +} + impl DecodeArrayTask for ListDecodeTask { fn decode(self: Box) -> Result<(ArrayRef, u64)> { let items = self @@ -782,11 +827,14 @@ impl LogicalPageDecoder for ListPageDecoder { } } if actual_num_rows < num_rows { - // TODO: We should be able to automatically - // shrink the read batch size if we detect the batches are going to be huge (maybe - // even achieve this with a read_batch_bytes parameter, though some estimation may - // still be required) - return Err(Error::not_supported_source(format!("loading a batch of {} lists would require creating an array with over i32::MAX items and we don't yet support returning smaller than requested batches", num_rows).into())); + let requested_range = self.rows_drained..self.rows_drained + num_rows; + let num_items = self.offsets[requested_range.end as usize] - item_start; + return Err(oversized_batch_error( + self.items_field.as_ref(), + requested_range, + self.rows_drained + actual_num_rows, + num_items, + )); } let offsets = self.offsets [self.rows_drained as usize..(self.rows_drained + actual_num_rows + 1) as usize] @@ -1281,3 +1329,74 @@ impl FieldEncoder for ListFieldEncoder { .boxed() } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_buffer::BooleanBuffer; + use arrow_schema::{DataType, Field}; + + use super::{ListPageDecoder, oversized_batch_error}; + use crate::decoder::LogicalPageDecoder; + + #[test] + fn oversized_binary_batch_error_is_actionable() { + let error = oversized_batch_error( + &Field::new("item", DataType::UInt8, false), + 32..160, + 96, + i32::MAX as u64 + 1, + ); + assert!( + error + .to_string() + .contains("more than 2GiB of string/binary data") + ); + assert!(error.to_string().contains("rows 32..96 fit")); + assert!(error.to_string().contains("requesting rows 32..160")); + assert!(error.to_string().contains("batch_size")); + assert!(error.to_string().contains("LANCE_DEFAULT_BATCH_SIZE")); + assert!(error.to_string().contains("large_string/large_binary")); + } + + #[test] + fn oversized_list_batch_error_is_actionable() { + let error = oversized_batch_error( + &Field::new("item", DataType::Int32, false), + 8..32, + 16, + i32::MAX as u64 + 1, + ); + let message = error.to_string(); + assert!(message.contains("list array")); + assert!(message.contains("more than i32::MAX items")); + assert!(message.contains("rows 8..16 fit")); + assert!(message.contains("requesting rows 8..32")); + } + + #[test] + fn list_decoder_overflow_reports_row_span() { + let mut decoder = ListPageDecoder { + unloaded: None, + offsets: Arc::<[u64]>::from(vec![0_u64, 1, 2, 3, 3 + i32::MAX as u64 + 1]), + validity: BooleanBuffer::from_iter([true, true, true, true]), + item_decoder: None, + num_rows: 4, + rows_drained: 2, + rows_loaded: 4, + items_field: Arc::new(Field::new("item", DataType::UInt8, false)), + offset_type: DataType::Int32, + data_type: DataType::List(Arc::new(Field::new("item", DataType::UInt8, false))), + }; + + let Err(error) = decoder.drain(2) else { + panic!("expected overflow error"); + }; + let message = error.to_string(); + assert!(message.contains("rows 2..3 fit")); + assert!(message.contains("requesting rows 2..4")); + assert!(message.contains(&(i32::MAX as u64 + 2).to_string())); + assert!(message.contains("batch_size")); + } +} diff --git a/rust/lance-encoding/src/array_encoding/physical/binary.rs b/rust/lance-encoding/src/array_encoding/physical/binary.rs index 294b295def3..8f413e2bdf3 100644 --- a/rust/lance-encoding/src/array_encoding/physical/binary.rs +++ b/rust/lance-encoding/src/array_encoding/physical/binary.rs @@ -30,7 +30,26 @@ use crate::{ use arrow_array::{PrimitiveArray, UInt64Array}; use arrow_schema::DataType; -use lance_core::Result; +use lance_core::{Error, Result}; + +fn oversized_binary_batch_error(num_rows: u64, num_bytes: u64) -> Error { + Error::not_supported(format!( + "Could not create array with more than 2GiB of string/binary data in a single batch \ + ({} rows would require {} bytes). Please reduce the batch_size, set \ + LANCE_DEFAULT_BATCH_SIZE to a smaller value, or convert the column to \ + large_string/large_binary.", + num_rows, num_bytes + )) +} + +fn oversized_large_binary_batch_error(num_rows: u64, num_bytes: u64) -> Error { + Error::not_supported(format!( + "Could not create large_string/large_binary array in a single batch because {} rows \ + would require {} bytes, which exceeds i64::MAX. Please reduce the batch_size or set \ + LANCE_DEFAULT_BATCH_SIZE to a smaller value.", + num_rows, num_bytes + )) +} struct IndicesNormalizer { indices: Vec, @@ -315,14 +334,25 @@ impl PrimitivePageDecoder for BinaryPageDecoder { // Normalize and cast (TODO: could fuse these into one pass for micro-optimization) let target_vec = target_offsets.values(); let start = target_vec[0]; - let offsets_buffer = - match bytes_per_offset { - 4 => ScalarBuffer::from_iter(target_vec.iter().map(|x| (x - start) as i32)) - .into_inner(), - 8 => ScalarBuffer::from_iter(target_vec.iter().map(|x| (x - start) as i64)) - .into_inner(), - _ => panic!("Unsupported offsets type"), - }; + let end = *target_vec.last().unwrap(); + let num_bytes = end - start; + let offsets_buffer = match bytes_per_offset { + 4 => { + if num_bytes > i32::MAX as u64 { + return Err(oversized_binary_batch_error(num_rows, num_bytes)); + } + ScalarBuffer::from_iter(target_vec.iter().map(|&offset| (offset - start) as i32)) + .into_inner() + } + 8 => { + if num_bytes > i64::MAX as u64 { + return Err(oversized_large_binary_batch_error(num_rows, num_bytes)); + } + ScalarBuffer::from_iter(target_vec.iter().map(|&offset| (offset - start) as i64)) + .into_inner() + } + _ => panic!("Unsupported offsets type"), + }; let bytes_to_skip = self.decoded_indices.value(rows_to_skip as usize); let num_bytes = self @@ -538,6 +568,38 @@ mod tests { use super::*; + #[derive(Debug)] + struct EmptyBytesDecoder; + + impl PrimitivePageDecoder for EmptyBytesDecoder { + fn decode(&self, _rows_to_skip: u64, _num_rows: u64) -> Result { + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 8, + data: LanceBuffer::empty(), + num_values: 0, + block_info: BlockInfo::new(), + })) + } + } + + #[derive(Debug)] + struct BytesDecoder { + bytes: Vec, + } + + impl PrimitivePageDecoder for BytesDecoder { + fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result { + let start = rows_to_skip as usize; + let end = start + num_rows as usize; + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 8, + data: LanceBuffer::from(self.bytes[start..end].to_vec()), + num_values: num_rows, + block_info: BlockInfo::new(), + })) + } + } + #[test] fn test_encode_indices_adjusts_nulls() { // Null entries in string arrays should be adjusted @@ -568,4 +630,70 @@ mod tests { ); assert_eq!(null_adjustment, 7); } + + #[test] + fn test_binary_overflow_error_is_actionable() { + let num_rows = 1; + let start = 100_u64; + let end = start + i32::MAX as u64 + 1; + let decoded_indices = UInt64Array::from(vec![start, end]); + let decoder = BinaryPageDecoder { + decoded_indices, + validity: BooleanBuffer::from_iter([true]), + offsets_type: DataType::Int32, + bytes_decoder: Box::new(EmptyBytesDecoder), + }; + + let error = decoder.decode(0, num_rows).unwrap_err(); + assert!(matches!(error, Error::NotSupported { .. })); + let message = error.to_string(); + assert!(message.contains("more than 2GiB of string/binary data")); + assert!(message.contains("batch_size")); + assert!(message.contains("LANCE_DEFAULT_BATCH_SIZE")); + assert!(message.contains("large_string/large_binary")); + } + + #[test] + fn test_large_binary_overflow_error_is_actionable() { + let num_rows = 1; + let start = 100_u64; + let end = start + i64::MAX as u64 + 1; + let decoded_indices = UInt64Array::from(vec![start, end]); + let decoder = BinaryPageDecoder { + decoded_indices, + validity: BooleanBuffer::from_iter([true]), + offsets_type: DataType::Int64, + bytes_decoder: Box::new(EmptyBytesDecoder), + }; + + let error = decoder.decode(0, num_rows).unwrap_err(); + assert!(matches!(error, Error::NotSupported { .. })); + let message = error.to_string(); + assert!(message.contains("large_string/large_binary")); + assert!(message.contains("exceeds i64::MAX")); + assert!(message.contains("batch_size")); + assert!(message.contains("LANCE_DEFAULT_BATCH_SIZE")); + } + + #[test] + fn test_large_binary_decode_success_path() { + let decoded_indices = UInt64Array::from(vec![100_u64, 102, 105]); + let mut bytes = vec![0_u8; 100]; + bytes.extend_from_slice(b"abcde"); + let decoder = BinaryPageDecoder { + decoded_indices, + validity: BooleanBuffer::from_iter([true, true]), + offsets_type: DataType::Int64, + bytes_decoder: Box::new(BytesDecoder { bytes }), + }; + + let data = decoder.decode(0, 2).unwrap(); + let variable = data.as_variable_width().unwrap(); + assert_eq!(variable.bits_per_offset, 64); + assert_eq!(variable.data.as_ref(), b"abcde"); + assert_eq!( + variable.offsets.borrow_to_typed_slice::().as_ref(), + &[0_i64, 2, 5] + ); + } }