fix: return actionable errors for binary/string offset overflow

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,19 +152,16 @@ pub struct BinaryArrayDecoder {
 }
 
 impl BinaryArrayDecoder {
-    fn from_list_array<T: ByteArrayType>(array: &GenericListArray<T::Offset>) -> ArrayRef {
+    fn from_list_array<T: ByteArrayType>(array: &GenericListArray<T::Offset>) -> Result<ArrayRef> {
         let values = array
             .values()
             .as_primitive::<UInt8Type>()
             .values()
             .inner()
             .clone();
         let offsets = array.offsets().clone();
-        Arc::new(GenericByteArray::<T>::new(
-            offsets,
-            values,
-            array.nulls().cloned(),
-        ))
+        let array = GenericByteArray::<T>::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::<BinaryType>(arr.as_list::<i32>()),
-            DataType::LargeBinary => Self::from_list_array::<LargeBinaryType>(arr.as_list::<i64>()),
-            DataType::Utf8 => Self::from_list_array::<Utf8Type>(arr.as_list::<i32>()),
-            DataType::LargeUtf8 => Self::from_list_array::<LargeUtf8Type>(arr.as_list::<i64>()),
-            _ => panic!("Binary decoder does not support this data type"),
+            DataType::Binary => Self::from_list_array::<BinaryType>(arr.as_list::<i32>())?,
+            DataType::LargeBinary => {
+                Self::from_list_array::<LargeBinaryType>(arr.as_list::<i64>())?
+            }
+            DataType::Utf8 => Self::from_list_array::<Utf8Type>(arr.as_list::<i32>())?,
+            DataType::LargeUtf8 => Self::from_list_array::<LargeUtf8Type>(arr.as_list::<i64>())?,
+            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<Self>) -> 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")
+        );
+    }
+}

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<u64>,
+    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<Self>) -> 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"));
+    }
+}