This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new cbd0e02f96 Return errors instead of panicking in fallible functions
(#10755)
cbd0e02f96 is described below
commit cbd0e02f9653c17836ec58bdc889947e4d550984
Author: Emil Ernerfeldt <[email protected]>
AuthorDate: Tue Sep 1 08:55:14 2026 +0200
Return errors instead of panicking in fallible functions (#10755)
# Which issue does this PR close?
- Part of https://github.com/apache/arrow-rs/issues/10553
# Rationale for this change
While clearing `clippy::missing_panics_doc`, a pattern turned up:
functions that
already return `Result`, but report some failures by panicking. Several
are
reachable from untrusted input, so the panic is a denial of service
rather than a
bug report.
Documenting those panics would normalise them, so this returns errors
instead.
This is one of four PRs splitting up the work.
1. **this PR** - return errors from fallible functions
2. #10759 - remove unreachable panics
3. #10760 - document the panics that genuinely remain
4. #10761 - `#[expect]` the unreachable ones, so the lint can be turned
on
# What changes are included in this PR?
Functions that could report the failure and did not.
# Are these changes tested?
Covered by the existing tests, plus new tests for the `b64_decode` and
`ArrayData::validate_values` error paths.
# Are there any user-facing changes?
No API changes. Some calls that used to panic now return an `Err`, which
is the
point of the PR.
---------
Co-authored-by: Claude Opus 5 (1M context) <[email protected]>
Co-authored-by: Jeffrey Vo <[email protected]>
---
.../fixed_size_binary_dictionary_builder.rs | 10 +-
.../builder/generic_bytes_dictionary_builder.rs | 10 +-
.../src/builder/primitive_dictionary_builder.rs | 10 +-
arrow-cast/src/base64.rs | 19 +-
arrow-data/src/data.rs | 198 +++++++++++++++++---
arrow-flight/src/sql/client.rs | 31 +++-
arrow-integration-test/src/datatype.rs | 30 +++-
arrow-integration-test/src/field.rs | 6 +-
arrow-integration-testing/src/lib.rs | 35 ++--
arrow-ipc/src/reader.rs | 200 ++++++++++++++++-----
arrow-json/src/reader/mod.rs | 8 +-
arrow-schema/src/ffi.rs | 47 +++--
arrow-select/src/dictionary.rs | 4 +-
arrow-string/src/concat_elements.rs | 20 ++-
parquet/src/arrow/arrow_reader/mod.rs | 13 +-
parquet/src/arrow/arrow_writer/mod.rs | 13 +-
parquet/src/arrow/async_reader/mod.rs | 13 +-
parquet/src/column/reader.rs | 24 ++-
parquet/src/file/writer.rs | 2 +-
19 files changed, 542 insertions(+), 151 deletions(-)
diff --git a/arrow-array/src/builder/fixed_size_binary_dictionary_builder.rs
b/arrow-array/src/builder/fixed_size_binary_dictionary_builder.rs
index a2a10ac4c4..98b7fed0a6 100644
--- a/arrow-array/src/builder/fixed_size_binary_dictionary_builder.rs
+++ b/arrow-array/src/builder/fixed_size_binary_dictionary_builder.rs
@@ -159,9 +159,13 @@ where
Ok(Self {
state,
dedup,
- keys_builder: new_keys
- .into_builder()
- .expect("underlying buffer has no references"),
+ keys_builder: new_keys.into_builder().map_err(|_| {
+ ArrowError::ComputeError(
+ "Internal Error: the keys just derived from the source
builder are \
+ unexpectedly shared, so they cannot be reused as a
builder"
+ .to_string(),
+ )
+ })?,
values_builder,
byte_width,
})
diff --git a/arrow-array/src/builder/generic_bytes_dictionary_builder.rs
b/arrow-array/src/builder/generic_bytes_dictionary_builder.rs
index aed195a033..2a7e0e9173 100644
--- a/arrow-array/src/builder/generic_bytes_dictionary_builder.rs
+++ b/arrow-array/src/builder/generic_bytes_dictionary_builder.rs
@@ -214,9 +214,13 @@ where
Ok(Self {
state,
dedup,
- keys_builder: new_keys
- .into_builder()
- .expect("underlying buffer has no references"),
+ keys_builder: new_keys.into_builder().map_err(|_| {
+ ArrowError::ComputeError(
+ "Internal Error: the keys just derived from the source
builder are \
+ unexpectedly shared, so they cannot be reused as a
builder"
+ .to_string(),
+ )
+ })?,
values_builder,
})
}
diff --git a/arrow-array/src/builder/primitive_dictionary_builder.rs
b/arrow-array/src/builder/primitive_dictionary_builder.rs
index 90b0c14a75..bc5d878500 100644
--- a/arrow-array/src/builder/primitive_dictionary_builder.rs
+++ b/arrow-array/src/builder/primitive_dictionary_builder.rs
@@ -227,9 +227,13 @@ where
Ok(Self {
map,
- keys_builder: new_keys
- .into_builder()
- .expect("underlying buffer has no references"),
+ keys_builder: new_keys.into_builder().map_err(|_| {
+ ArrowError::ComputeError(
+ "Internal Error: the keys just derived from the source
builder are \
+ unexpectedly shared, so they cannot be reused as a
builder"
+ .to_string(),
+ )
+ })?,
values_builder,
})
}
diff --git a/arrow-cast/src/base64.rs b/arrow-cast/src/base64.rs
index 0cd724df7c..e3cc03ff09 100644
--- a/arrow-cast/src/base64.rs
+++ b/arrow-cast/src/base64.rs
@@ -61,6 +61,10 @@ pub fn b64_encode<E: Engine, O: OffsetSizeTrait>(
}
/// Base64 decode each element of `array` with the provided [`Engine`]
+///
+/// # Errors
+///
+/// Returns an error if a value is not valid base64 for `engine`.
pub fn b64_decode<E: Engine, O: OffsetSizeTrait>(
engine: &E,
array: &GenericBinaryArray<O>,
@@ -74,7 +78,11 @@ pub fn b64_decode<E: Engine, O: OffsetSizeTrait>(
for v in array {
if let Some(v) = v {
- let len = engine.decode_slice(v, &mut buffer[offset..]).unwrap();
+ let len = engine
+ .decode_slice(v, &mut buffer[offset..])
+ .map_err(|err| {
+ ArrowError::InvalidArgumentError(format!("Failed to decode
base64: {err}"))
+ })?;
// This cannot overflow as `len` is less than `v.len()` and `a` is
valid
offset += len;
}
@@ -120,6 +128,15 @@ mod tests {
test_engine(&BASE64_STANDARD_NO_PAD, &data);
}
+ #[test]
+ fn test_b64_decode_invalid_input() {
+ let data: BinaryArray = vec![Some(b"!!!not base64!!!".to_vec())]
+ .into_iter()
+ .collect();
+ let err = b64_decode(&BASE64_STANDARD, &data).unwrap_err().to_string();
+ assert!(err.contains("Failed to decode base64"), "{err}");
+ }
+
/// Safe-Rust `Engine` that writes invalid UTF-8 into the encode buffer
/// (#10284). `b64_encode` must reject it rather than build an unsound
/// `StringArray`.
diff --git a/arrow-data/src/data.rs b/arrow-data/src/data.rs
index 7bf9520785..ed73bf6273 100644
--- a/arrow-data/src/data.rs
+++ b/arrow-data/src/data.rs
@@ -1090,7 +1090,7 @@ impl ArrayData {
/// For an empty array, the `buffer` can also be empty.
fn typed_offsets<T: ArrowNativeType + num_traits::Num>(&self) ->
Result<&[T], ArrowError> {
// An empty list-like array can have 0 offsets
- if self.len == 0 && self.buffers[0].is_empty() {
+ if self.len == 0 && self.buffer_at(0)?.is_empty() {
return Ok(&[]);
}
@@ -1105,7 +1105,7 @@ impl ArrayData {
idx: usize,
len: usize,
) -> Result<&[T], ArrowError> {
- let buffer = &self.buffers[idx];
+ let buffer = self.buffer_at(idx)?;
let required_elements = checked_len_plus_offset(&self.data_type, len,
self.offset)?;
let byte_width = mem::size_of::<T>();
@@ -1376,6 +1376,36 @@ impl ArrayData {
self.get_valid_child_data(0, expected_type)
}
+ /// Returns `buffers[idx]`, or an error if there is no such buffer.
+ ///
+ /// [`Self::validate_values`] can be called on its own, without the buffer
counts
+ /// having been checked by [`Self::validate`] first, so the index may be
missing.
+ fn buffer_at(&self, idx: usize) -> Result<&Buffer, ArrowError> {
+ self.buffers.get(idx).ok_or_else(|| {
+ ArrowError::InvalidArgumentError(format!(
+ "{} should contain at least {} buffer(s), had {}",
+ self.data_type,
+ idx + 1,
+ self.buffers.len()
+ ))
+ })
+ }
+
+ /// Returns `child_data[idx]`, or an error if there is no such child.
+ ///
+ /// [`Self::validate_values`] can be called on its own, without the child
counts
+ /// having been checked by [`Self::validate`] first, so the index may be
missing.
+ fn child_at(&self, idx: usize) -> Result<&ArrayData, ArrowError> {
+ self.child_data.get(idx).ok_or_else(|| {
+ ArrowError::InvalidArgumentError(format!(
+ "{} should contain at least {} child data array(s), had {}",
+ self.data_type,
+ idx + 1,
+ self.child_data.len()
+ ))
+ })
+ }
+
/// Returns `Err` if self.child_data does not have exactly `expected_len`
elements
fn validate_num_child_data(&self, expected_len: usize) -> Result<(),
ArrowError> {
if self.child_data.len() != expected_len {
@@ -1552,8 +1582,8 @@ impl ArrayData {
match &self.data_type {
DataType::Utf8 => self.validate_utf8::<i32>(),
DataType::LargeUtf8 => self.validate_utf8::<i64>(),
- DataType::Binary =>
self.validate_offsets_full::<i32>(self.buffers[1].len()),
- DataType::LargeBinary =>
self.validate_offsets_full::<i64>(self.buffers[1].len()),
+ DataType::Binary =>
self.validate_offsets_full::<i32>(self.buffer_at(1)?.len()),
+ DataType::LargeBinary =>
self.validate_offsets_full::<i64>(self.buffer_at(1)?.len()),
DataType::BinaryView => {
let views = self.typed_buffer::<u128>(0, self.len)?;
validate_binary_view(views, &self.buffers[1..])
@@ -1563,11 +1593,11 @@ impl ArrayData {
validate_string_view(views, &self.buffers[1..])
}
DataType::List(_) | DataType::Map(_, _) => {
- let child = &self.child_data[0];
+ let child = self.child_at(0)?;
self.validate_offsets_full::<i32>(child.len)
}
DataType::LargeList(_) => {
- let child = &self.child_data[0];
+ let child = self.child_at(0)?;
self.validate_offsets_full::<i64>(child.len)
}
DataType::Union(_, _) => {
@@ -1579,7 +1609,12 @@ impl ArrayData {
Ok(())
}
DataType::Dictionary(key_type, _value_type) => {
- let dictionary_length: i64 =
self.child_data[0].len.try_into().unwrap();
+ let dictionary_length = self.child_at(0)?.len;
+ let dictionary_length =
i64::try_from(dictionary_length).map_err(|_| {
+ ArrowError::InvalidArgumentError(format!(
+ "Dictionary of {dictionary_length} values is too long
for an i64"
+ ))
+ })?;
let max_value = dictionary_length - 1;
match key_type.as_ref() {
DataType::UInt8 => self.check_bounds::<u8>(max_value),
@@ -1590,16 +1625,20 @@ impl ArrayData {
DataType::Int16 => self.check_bounds::<i16>(max_value),
DataType::Int32 => self.check_bounds::<i32>(max_value),
DataType::Int64 => self.check_bounds::<i64>(max_value),
- _ => unreachable!(),
+ _ => Err(ArrowError::InvalidArgumentError(format!(
+ "Dictionary key type must be an integer, got
{key_type}"
+ ))),
}
}
DataType::RunEndEncoded(run_ends, _values) => {
- let run_ends_data = self.child_data()[0].clone();
+ let run_ends_data = self.child_at(0)?;
match run_ends.data_type() {
DataType::Int16 => run_ends_data.check_run_ends::<i16>(),
DataType::Int32 => run_ends_data.check_run_ends::<i32>(),
DataType::Int64 => run_ends_data.check_run_ends::<i64>(),
- _ => unreachable!(),
+ data_type => Err(ArrowError::InvalidArgumentError(format!(
+ "Run end type must be Int16, Int32 or Int64, got
{data_type}"
+ ))),
}
}
_ => {
@@ -1670,7 +1709,7 @@ impl ArrayData {
where
T: ArrowNativeType + TryInto<usize> + num_traits::Num +
std::fmt::Display,
{
- let values_buffer = &self.buffers[1].as_slice();
+ let values_buffer = &self.buffer_at(1)?.as_slice();
if let Ok(values_str) = std::str::from_utf8(values_buffer) {
// Validate Offsets are correct
self.validate_each_offset::<T, _>(values_buffer.len(),
|string_index, range| {
@@ -1715,15 +1754,9 @@ impl ArrayData {
where
T: ArrowNativeType + TryInto<i64> + num_traits::Num +
std::fmt::Display,
{
- let required_len = checked_len_plus_offset(&self.data_type, self.len,
self.offset)?;
- let buffer = &self.buffers[0];
-
- // This should have been checked as part of `validate()` prior
- // to calling `validate_full()` but double check to be sure
- assert!(buffer.len() / mem::size_of::<T>() >= required_len);
-
- // Justification: buffer size was validated above
- let indexes: &[T] =
&buffer.typed_data::<T>()[self.offset..required_len];
+ // `validate()` checks the buffer size too, but `validate_values()`
can be called
+ // on its own, so do not assume it has run.
+ let indexes: &[T] = self.typed_buffer::<T>(0, self.len)?;
indexes.iter().enumerate().try_for_each(|(i, &dict_index)| {
// Do not check the value is null (value can be arbitrary)
@@ -2958,6 +2991,131 @@ mod tests {
);
}
+ /// Without `force_validate`, `build_unchecked` skips validation, so these
can
+ /// reach `validate_values` with a data type that has an invalid child
type.
+ #[test]
+ #[cfg(not(feature = "force_validate"))]
+ fn test_validate_values_rejects_a_non_integer_dictionary_key() {
+ let values = valid_non_nullable_int32_array_data(2);
+ let data_type = DataType::Dictionary(Box::new(DataType::Utf8),
Box::new(DataType::Int32));
+ let dictionary = unsafe {
+ ArrayData::builder(data_type)
+ .len(1)
+ .add_child_data(values)
+ .build_unchecked()
+ };
+
+ let err = dictionary.validate_values().expect_err("should get error");
+ assert_eq!(
+ err.to_string(),
+ "Invalid argument error: Dictionary key type must be an integer,
got Utf8"
+ );
+ }
+
+ #[test]
+ #[cfg(not(feature = "force_validate"))]
+ fn test_validate_values_rejects_a_non_integer_run_end() {
+ let data_type = DataType::RunEndEncoded(
+ Arc::new(Field::new("run_ends", DataType::Utf8, false)),
+ Arc::new(Field::new("values", DataType::Int32, true)),
+ );
+ let run_end_encoded = unsafe {
+ ArrayData::builder(data_type)
+ .len(1)
+ .add_child_data(valid_non_nullable_int32_array_data(1))
+ .add_child_data(valid_non_nullable_int32_array_data(1))
+ .build_unchecked()
+ };
+
+ let err = run_end_encoded
+ .validate_values()
+ .expect_err("should get error");
+ assert_eq!(
+ err.to_string(),
+ "Invalid argument error: Run end type must be Int16, Int32 or
Int64, got Utf8"
+ );
+ }
+
+ /// `validate_values` must report missing children rather than index out
of bounds.
+ #[test]
+ #[cfg(not(feature = "force_validate"))]
+ fn test_validate_values_rejects_missing_child_data() {
+ let int32 = Box::new(DataType::Int32);
+ let field = || Arc::new(Field::new("f", DataType::Int32, true));
+ let data_types = [
+ DataType::Dictionary(int32.clone(), int32.clone()),
+ DataType::List(field()),
+ DataType::LargeList(field()),
+ DataType::RunEndEncoded(field(), field()),
+ ];
+
+ for data_type in data_types {
+ let data = unsafe {
+ ArrayData::builder(data_type.clone())
+ .len(1)
+ .build_unchecked()
+ };
+ let err = data.validate_values().expect_err("should get error");
+ assert_eq!(
+ err.to_string(),
+ format!(
+ "Invalid argument error: {data_type} should contain at
least 1 child data array(s), had 0"
+ )
+ );
+ }
+ }
+
+ /// `validate_values` must report missing buffers rather than index out of
bounds.
+ #[test]
+ #[cfg(not(feature = "force_validate"))]
+ fn test_validate_values_rejects_missing_buffers() {
+ // (data type, index of the first missing buffer)
+ let cases = [
+ (DataType::Utf8, 1),
+ (DataType::LargeUtf8, 1),
+ (DataType::Binary, 1),
+ (DataType::LargeBinary, 1),
+ (DataType::BinaryView, 0),
+ (DataType::Utf8View, 0),
+ ];
+
+ for (data_type, missing) in cases {
+ let data = unsafe {
+ ArrayData::builder(data_type.clone())
+ .len(1)
+ .build_unchecked()
+ };
+ let err = data.validate_values().expect_err("should get error");
+ assert_eq!(
+ err.to_string(),
+ format!(
+ "Invalid argument error: {data_type} should contain at
least {} buffer(s), had 0",
+ missing + 1
+ )
+ );
+ }
+ }
+
+ /// A dictionary whose keys buffer is too small must be reported, not
asserted on.
+ #[test]
+ #[cfg(not(feature = "force_validate"))]
+ fn test_validate_values_rejects_a_short_dictionary_keys_buffer() {
+ let data_type = DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Int32));
+ let dictionary = unsafe {
+ ArrayData::builder(data_type)
+ .len(4)
+ .add_buffer(Buffer::from_slice_ref([1_i32, 0]))
+ .add_child_data(valid_non_nullable_int32_array_data(2))
+ .build_unchecked()
+ };
+
+ let err = dictionary.validate_values().expect_err("should get error");
+ assert_eq!(
+ err.to_string(),
+ "Invalid argument error: Buffer 0 of Dictionary(Int32, Int32)
isn't large enough. Expected 16 bytes got 8"
+ );
+ }
+
#[test]
fn should_fail_validation_when_having_map_field_type_is_not_struct() {
let map_field = Field::new("key", DataType::Int32, false);
diff --git a/arrow-flight/src/sql/client.rs b/arrow-flight/src/sql/client.rs
index 6ea1a04fc1..e00f7254a3 100644
--- a/arrow-flight/src/sql/client.rs
+++ b/arrow-flight/src/sql/client.rs
@@ -222,7 +222,9 @@ where
.into_request(),
)?;
let mut result = self.flight_client.do_put(req).await?.into_inner();
- let result = result.message().await?.unwrap();
+ let result = result.message().await?.ok_or_else(|| {
+ FlightError::protocol("Server closed the stream without sending a
result")
+ })?;
let result: DoPutUpdateResult =
Message::decode(&*result.app_metadata)?;
Ok(result.record_count)
}
@@ -258,7 +260,9 @@ where
return Err(FlightError::ExternalError(Box::new(msg)));
}
- let result = result.message().await?.unwrap();
+ let result = result.message().await?.ok_or_else(|| {
+ FlightError::protocol("Server closed the stream without sending a
result")
+ })?;
let result: DoPutUpdateResult =
Message::decode(&*result.app_metadata)?;
Ok(result.record_count)
}
@@ -387,9 +391,16 @@ where
};
let req = self.set_request_headers(action.into_request())?;
let mut result = self.flight_client.do_action(req).await?.into_inner();
- let result = result.message().await?.unwrap();
+ let result = result.message().await?.ok_or_else(|| {
+ FlightError::protocol("Server closed the stream without sending a
result")
+ })?;
let any = Any::decode(&*result.body)?;
- let prepared_result: ActionCreatePreparedStatementResult =
any.unpack()?.unwrap();
+ let prepared_result: ActionCreatePreparedStatementResult =
+ any.unpack()?.ok_or_else(|| {
+ FlightError::protocol(
+ "Server did not return an
ActionCreatePreparedStatementResult",
+ )
+ })?;
let dataset_schema = match prepared_result.dataset_schema.len() {
0 => Schema::empty(),
_ => Schema::try_from(IpcMessage(prepared_result.dataset_schema))?,
@@ -415,9 +426,13 @@ where
};
let req = self.set_request_headers(action.into_request())?;
let mut result = self.flight_client.do_action(req).await?.into_inner();
- let result = result.message().await?.unwrap();
+ let result = result.message().await?.ok_or_else(|| {
+ FlightError::protocol("Server closed the stream without sending a
result")
+ })?;
let any = Any::decode(&*result.body)?;
- let begin_result: ActionBeginTransactionResult =
any.unpack()?.unwrap();
+ let begin_result: ActionBeginTransactionResult =
any.unpack()?.ok_or_else(|| {
+ FlightError::protocol("Server did not return an
ActionBeginTransactionResult")
+ })?;
Ok(begin_result.transaction_id)
}
@@ -542,7 +557,9 @@ where
..Default::default()
}]))
.await?;
- let result = result.message().await?.unwrap();
+ let result = result.message().await?.ok_or_else(|| {
+ FlightError::protocol("Server closed the stream without sending a
result")
+ })?;
let result: DoPutUpdateResult =
Message::decode(&*result.app_metadata)?;
Ok(result.record_count)
}
diff --git a/arrow-integration-test/src/datatype.rs
b/arrow-integration-test/src/datatype.rs
index 69174a1c22..9ff0077639 100644
--- a/arrow-integration-test/src/datatype.rs
+++ b/arrow-integration-test/src/datatype.rs
@@ -19,6 +19,15 @@ use arrow::datatypes::{DataType, Field, Fields,
IntervalUnit, TimeUnit, UnionMod
use arrow::error::{ArrowError, Result};
use std::sync::Arc;
+/// Read a JSON number as an integer of type `T`.
+fn json_int<T: TryFrom<i64>>(what: &str, value: &serde_json::Value) ->
Result<T> {
+ let int = value
+ .as_i64()
+ .ok_or_else(|| ArrowError::ParseError(format!("Expecting {what} to be
an integer")))?;
+ T::try_from(int)
+ .map_err(|_| ArrowError::ParseError(format!("{what} is out of range
for its type: {int}")))
+}
+
/// Parse a data type from a JSON representation.
pub fn data_type_from_json(json: &serde_json::Value) -> Result<DataType> {
use serde_json::Value;
@@ -36,7 +45,10 @@ pub fn data_type_from_json(json: &serde_json::Value) ->
Result<DataType> {
Some(s) if s == "fixedsizebinary" => {
// return a list with any type as its child isn't defined in
the map
if let Some(Value::Number(size)) = map.get("byteWidth") {
- Ok(DataType::FixedSizeBinary(size.as_i64().unwrap() as
i32))
+ Ok(DataType::FixedSizeBinary(json_int(
+ "byteWidth",
+ &Value::Number(size.clone()),
+ )?))
} else {
Err(ArrowError::ParseError(
"Expecting a byteWidth for
fixedsizebinary".to_string(),
@@ -46,19 +58,19 @@ pub fn data_type_from_json(json: &serde_json::Value) ->
Result<DataType> {
Some(s) if s == "decimal" => {
// return a list with any type as its child isn't defined in
the map
let precision = match map.get("precision") {
- Some(p) => Ok(p.as_u64().unwrap().try_into().unwrap()),
+ Some(p) => json_int("precision", p),
None => Err(ArrowError::ParseError(
"Expecting a precision for decimal".to_string(),
)),
}?;
let scale = match map.get("scale") {
- Some(s) => Ok(s.as_u64().unwrap().try_into().unwrap()),
+ Some(s) => json_int("scale", s),
_ => Err(ArrowError::ParseError(
"Expecting a scale for decimal".to_string(),
)),
}?;
let bit_width: usize = match map.get("bitWidth") {
- Some(b) => b.as_u64().unwrap() as usize,
+ Some(b) => json_int("bitWidth", b)?,
_ => 128, // Default bit width
};
@@ -197,7 +209,7 @@ pub fn data_type_from_json(json: &serde_json::Value) ->
Result<DataType> {
if let Some(Value::Number(size)) = map.get("listSize") {
Ok(DataType::FixedSizeList(
default_field,
- size.as_i64().unwrap() as i32,
+ json_int("listSize", &Value::Number(size.clone()))?,
))
} else {
Err(ArrowError::ParseError(
@@ -238,10 +250,14 @@ pub fn data_type_from_json(json: &serde_json::Value) ->
Result<DataType> {
)));
};
if let Some(values) = map.get("typeIds") {
- let values = values.as_array().unwrap();
+ let values = values.as_array().ok_or_else(|| {
+ ArrowError::ParseError("Expecting typeIds to be an
array".to_string())
+ })?;
let fields = values
.iter()
- .map(|t| (t.as_i64().unwrap() as i8,
default_field.clone()))
+ .map(|t| Ok((json_int::<i8>("a type id", t)?,
default_field.clone())))
+ .collect::<Result<Vec<_>>>()?
+ .into_iter()
.collect();
Ok(DataType::Union(fields, union_mode))
diff --git a/arrow-integration-test/src/field.rs
b/arrow-integration-test/src/field.rs
index da30217431..68019cc537 100644
--- a/arrow-integration-test/src/field.rs
+++ b/arrow-integration-test/src/field.rs
@@ -259,7 +259,11 @@ pub fn field_from_json(json: &serde_json::Value) ->
Result<Field> {
}
};
dict_id = match dictionary.get("id") {
- Some(Value::Number(n)) => n.as_i64().unwrap(),
+ Some(Value::Number(n)) => n.as_i64().ok_or_else(|| {
+ ArrowError::ParseError(
+ "Field 'id' attribute is not an
integer".to_string(),
+ )
+ })?,
_ => {
return Err(ArrowError::ParseError(
"Field missing 'id' attribute".to_string(),
diff --git a/arrow-integration-testing/src/lib.rs
b/arrow-integration-testing/src/lib.rs
index 613408ae59..a1f7d2a70a 100644
--- a/arrow-integration-testing/src/lib.rs
+++ b/arrow-integration-testing/src/lib.rs
@@ -58,23 +58,27 @@ pub struct ArrowFile {
impl ArrowFile {
/// Read a single [RecordBatch] from the file
pub fn read_batch(&self, batch_num: usize) -> Result<RecordBatch> {
- let b = self.arrow_json["batches"].get(batch_num).unwrap();
- let json_batch: ArrowJsonBatch =
serde_json::from_value(b.clone()).unwrap();
- record_batch_from_json(&self.schema, json_batch,
Some(&self.dictionaries))
+ let b = self.arrow_json["batches"].get(batch_num).ok_or_else(|| {
+ ArrowError::ParseError(format!("Arrow JSON has no batch
{batch_num}"))
+ })?;
+ self.batch_from_json(b)
}
/// Read all [RecordBatch]es from the file
pub fn read_batches(&self) -> Result<Vec<RecordBatch>> {
self.arrow_json["batches"]
.as_array()
- .unwrap()
+ .ok_or_else(|| ArrowError::ParseError("Arrow JSON has no 'batches'
array".to_string()))?
.iter()
- .map(|b| {
- let json_batch: ArrowJsonBatch =
serde_json::from_value(b.clone()).unwrap();
- record_batch_from_json(&self.schema, json_batch,
Some(&self.dictionaries))
- })
+ .map(|b| self.batch_from_json(b))
.collect()
}
+
+ fn batch_from_json(&self, batch: &Value) -> Result<RecordBatch> {
+ let json_batch: ArrowJsonBatch = serde_json::from_value(batch.clone())
+ .map_err(|err| ArrowError::ParseError(format!("Invalid Arrow JSON
batch: {err}")))?;
+ record_batch_from_json(&self.schema, json_batch,
Some(&self.dictionaries))
+ }
}
/// Canonicalize the names of map fields in a schema
@@ -122,17 +126,20 @@ pub fn canonicalize_schema(schema: &Schema) -> Schema {
pub fn open_json_file(json_name: &str) -> Result<ArrowFile> {
let json_file = File::open(json_name)?;
let reader = BufReader::new(json_file);
- let arrow_json: Value = serde_json::from_reader(reader).unwrap();
+ let arrow_json: Value = serde_json::from_reader(reader)
+ .map_err(|err| ArrowError::ParseError(format!("Invalid Arrow JSON:
{err}")))?;
let schema = schema_from_json(&arrow_json["schema"])?;
// read dictionaries
let mut dictionaries = HashMap::new();
if let Some(dicts) = arrow_json.get("dictionaries") {
- for d in dicts
- .as_array()
- .expect("Unable to get dictionaries as array")
- {
+ let dicts = dicts.as_array().ok_or_else(|| {
+ ArrowError::ParseError("Arrow JSON 'dictionaries' is not an
array".to_string())
+ })?;
+ for d in dicts {
let json_dict: ArrowJsonDictionaryBatch =
- serde_json::from_value(d.clone()).expect("Unable to get
dictionary from JSON");
+ serde_json::from_value(d.clone()).map_err(|err| {
+ ArrowError::ParseError(format!("Invalid Arrow JSON
dictionary: {err}"))
+ })?;
// TODO: convert to a concrete Arrow type
dictionaries.insert(json_dict.id, json_dict);
}
diff --git a/arrow-ipc/src/reader.rs b/arrow-ipc/src/reader.rs
index c8eb45515b..3b409d2a2f 100644
--- a/arrow-ipc/src/reader.rs
+++ b/arrow-ipc/src/reader.rs
@@ -102,12 +102,7 @@ impl RecordBatchDecoder<'_> {
self.create_primitive_array(field_node, data_type, &buffers)
}
BinaryView | Utf8View => {
- let count = variadic_counts
- .pop_front()
- .ok_or(ArrowError::IpcError(format!(
- "Missing variadic count for {data_type} column"
- )))?;
- let count = count + 2; // view and null buffer.
+ let count = self.next_variadic_buffer_count(variadic_counts,
data_type)?;
let buffers = (0..count)
.map(|_| self.next_buffer())
.collect::<Result<Vec<_>, _>>()?;
@@ -546,6 +541,11 @@ impl<'a> RecordBatchDecoder<'a> {
}
/// Read the record batch, consuming the reader
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if the message does not describe a batch matching the
schema,
+ /// for example if it declares more variadic buffer counts than the schema
uses.
pub fn read_record_batch(mut self) -> Result<RecordBatch, ArrowError> {
let mut variadic_counts: VecDeque<i64> = self
.batch
@@ -598,7 +598,7 @@ impl<'a> RecordBatchDecoder<'a> {
))
}
} else {
- assert!(variadic_counts.is_empty());
+ check_variadic_counts_consumed(&variadic_counts)?;
RecordBatch::try_new_with_options(schema, columns, &options)
}
} else {
@@ -619,7 +619,7 @@ impl<'a> RecordBatchDecoder<'a> {
))
}
} else {
- assert!(variadic_counts.is_empty());
+ check_variadic_counts_consumed(&variadic_counts)?;
RecordBatch::try_new_with_options(schema, children, &options)
}
}
@@ -637,8 +637,11 @@ impl<'a> RecordBatchDecoder<'a> {
)
}
- fn skip_buffer(&mut self) {
- self.buffers.next().unwrap();
+ fn skip_buffer(&mut self) -> Result<(), ArrowError> {
+ self.buffers.next().ok_or_else(|| {
+ ArrowError::IpcError("Buffer count mismatched with
metadata".to_string())
+ })?;
+ Ok(())
}
fn next_node(&mut self, field: &Field) -> Result<&'a FieldNode,
ArrowError> {
@@ -659,42 +662,36 @@ impl<'a> RecordBatchDecoder<'a> {
match field.data_type() {
Utf8 | Binary | LargeBinary | LargeUtf8 => {
for _ in 0..3 {
- self.skip_buffer()
+ self.skip_buffer()?;
}
}
Utf8View | BinaryView => {
- let count = variadic_count
- .pop_front()
- .ok_or(ArrowError::IpcError(format!(
- "Missing variadic count for {} column",
- field.data_type()
- )))?;
- let count = count + 2; // view and null buffer.
- for _i in 0..count {
- self.skip_buffer()
+ let count = self.next_variadic_buffer_count(variadic_count,
field.data_type())?;
+ for _ in 0..count {
+ self.skip_buffer()?;
}
}
FixedSizeBinary(_) => {
- self.skip_buffer();
- self.skip_buffer();
+ self.skip_buffer()?;
+ self.skip_buffer()?;
}
List(list_field) | LargeList(list_field) | Map(list_field, _) => {
- self.skip_buffer();
- self.skip_buffer();
+ self.skip_buffer()?;
+ self.skip_buffer()?;
self.skip_field(list_field, variadic_count)?;
}
ListView(list_field) | LargeListView(list_field) => {
- self.skip_buffer(); // Null buffer
- self.skip_buffer(); // Offsets
- self.skip_buffer(); // Sizes
+ self.skip_buffer()?; // Null buffer
+ self.skip_buffer()?; // Offsets
+ self.skip_buffer()?; // Sizes
self.skip_field(list_field, variadic_count)?;
}
FixedSizeList(list_field, _) => {
- self.skip_buffer();
+ self.skip_buffer()?;
self.skip_field(list_field, variadic_count)?;
}
Struct(struct_fields) => {
- self.skip_buffer();
+ self.skip_buffer()?;
// skip for each field
for struct_field in struct_fields {
@@ -706,17 +703,17 @@ impl<'a> RecordBatchDecoder<'a> {
self.skip_field(values_field, variadic_count)?;
}
Dictionary(_, _) => {
- self.skip_buffer(); // Nulls
- self.skip_buffer(); // Indices
+ self.skip_buffer()?; // Nulls
+ self.skip_buffer()?; // Indices
}
Union(fields, mode) => {
if self.version < MetadataVersion::V5 {
- self.skip_buffer(); // Null buffer
+ self.skip_buffer()?; // Null buffer
}
- self.skip_buffer(); // Type ids
+ self.skip_buffer()?; // Type ids
match mode {
- UnionMode::Dense => self.skip_buffer(), // Offsets
+ UnionMode::Dense => self.skip_buffer()?, // Offsets
UnionMode::Sparse => {}
}
@@ -751,14 +748,45 @@ impl<'a> RecordBatchDecoder<'a> {
| Decimal64(_, _)
| Decimal128(_, _)
| Decimal256(_, _) => {
- self.skip_buffer();
- self.skip_buffer();
+ self.skip_buffer()?;
+ self.skip_buffer()?;
}
}
Ok(())
}
}
+impl RecordBatchDecoder<'_> {
+ /// Takes the number of variadic buffers declared for one `BinaryView` or
`Utf8View`
+ /// column, and returns the total number of buffers to read for it.
+ ///
+ /// The count comes from the IPC message, so it may be missing, negative,
or larger
+ /// than the number of buffers the message actually has.
+ fn next_variadic_buffer_count(
+ &self,
+ variadic_counts: &mut VecDeque<i64>,
+ data_type: &DataType,
+ ) -> Result<usize, ArrowError> {
+ let count = variadic_counts.pop_front().ok_or_else(|| {
+ ArrowError::IpcError(format!("Missing variadic count for
{data_type} column"))
+ })?;
+
+ let remaining = self.buffers.len();
+
+ // The view buffer and the null buffer are not counted as variadic.
+ usize::try_from(count)
+ .ok()
+ .and_then(|count| count.checked_add(2))
+ .filter(|total| *total <= remaining)
+ .ok_or_else(|| {
+ ArrowError::IpcError(format!(
+ "Invalid variadic count {count} for {data_type} column, \
+ with {remaining} buffer(s) left in the message"
+ ))
+ })
+ }
+}
+
/// Creates a record batch from binary data using the `crate::RecordBatch`
indexes and the `Schema`.
///
/// If `require_alignment` is true, this function will return an error if any
array data in the
@@ -930,6 +958,22 @@ fn read_block<R: Read + Seek>(mut reader: R, block:
&Block) -> Result<Buffer, Ar
Ok(buf.into())
}
+/// One variadic buffer count is consumed per `BinaryView` or `Utf8View`
column in
+/// the schema, so any count left over means the message and the schema
disagree.
+///
+/// The opposite case, too few counts, is reported by
[`RecordBatchDecoder::create_array`].
+fn check_variadic_counts_consumed(variadic_counts: &VecDeque<i64>) ->
Result<(), ArrowError> {
+ if variadic_counts.is_empty() {
+ Ok(())
+ } else {
+ Err(ArrowError::IpcError(format!(
+ "Mismatch between schema and data: the IPC message declares {}
more variadic \
+ buffer count(s) than the schema has BinaryView or Utf8View
columns",
+ variadic_counts.len()
+ )))
+ }
+}
+
/// Parse an encapsulated message
///
///
<https://arrow.apache.org/docs/format/Columnar.html#encapsulated-message-format>
@@ -1273,10 +1317,12 @@ impl FileReaderBuilder {
let mut custom_metadata = HashMap::new();
if let Some(fb_custom_metadata) = footer.custom_metadata() {
for kv in fb_custom_metadata {
- custom_metadata.insert(
- kv.key().unwrap().to_string(),
- kv.value().unwrap().to_string(),
- );
+ let (Some(key), Some(value)) = (kv.key(), kv.value()) else {
+ return Err(ArrowError::ParseError(
+ "Custom metadata in the IPC footer is missing a key or
a value".to_string(),
+ ));
+ };
+ custom_metadata.insert(key.to_string(), value.to_string());
}
}
@@ -2200,6 +2246,80 @@ mod tests {
}
}
+ /// A `Utf8View` batch whose variadic buffer count is `count`, with two
buffers
+ /// in the message. Returns the error from reading it with the given
projection.
+ fn read_batch_with_variadic_count(count: i64, projection:
Option<&[usize]>) -> ArrowError {
+ use crate::r#gen::Message::*;
+ use flatbuffers::FlatBufferBuilder;
+
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "col",
+ DataType::Utf8View,
+ true,
+ )]));
+
+ let mut fbb = FlatBufferBuilder::new();
+ let nodes = fbb.create_vector(&[FieldNode::new(1, 0)]);
+ let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8),
crate::Buffer::new(8, 8)]);
+ let variadic_buffer_counts = fbb.create_vector(&[count]);
+ let batch_offset = RecordBatch::create(
+ &mut fbb,
+ &RecordBatchArgs {
+ length: 1,
+ nodes: Some(nodes),
+ buffers: Some(buffers),
+ compression: None,
+ variadicBufferCounts: Some(variadic_buffer_counts),
+ },
+ );
+ fbb.finish_minimal(batch_offset);
+ let batch_bytes = fbb.finished_data().to_vec();
+ let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
+
+ let data_buffer = Buffer::from(vec![0u8; 16]);
+ let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
+
+ RecordBatchDecoder::try_new(
+ &data_buffer,
+ batch,
+ schema,
+ &dictionaries,
+ &MetadataVersion::V5,
+ )
+ .unwrap()
+ .with_projection(projection)
+ .read_record_batch()
+ .expect_err("should get error")
+ }
+
+ /// A variadic count the message cannot honour used to panic while slicing
the
+ /// buffers it did not read, both when reading the column and when
skipping it.
+ #[test]
+ fn test_invalid_variadic_buffer_count_error() {
+ // -2 leaves no buffers at all, -1 leaves too few, and 1 asks for more
than the
+ // message has. The projection selects nothing, so the column is
skipped instead.
+ for count in [-2, -1, 1, i64::MAX] {
+ for projection in [None, Some([].as_slice())] {
+ let err = read_batch_with_variadic_count(count, projection);
+ assert_eq!(
+ err.to_string(),
+ format!(
+ "Ipc error: Invalid variadic count {count} for
Utf8View column, \
+ with 2 buffer(s) left in the message"
+ ),
+ "count {count}, projection {projection:?}"
+ );
+ }
+ }
+ }
+
+ /// The valid count for a message with two buffers is zero.
+ #[test]
+ fn test_valid_variadic_buffer_count_is_accepted() {
+ let err = read_batch_with_variadic_count(0, None);
+ assert!(!err.to_string().contains("Invalid variadic count"), "{err}");
+ }
+
#[test]
fn test_missing_footer_schema_error() {
use crate::r#gen::File::{Footer, FooterArgs};
diff --git a/arrow-json/src/reader/mod.rs b/arrow-json/src/reader/mod.rs
index 643e45f238..ad17402ff8 100644
--- a/arrow-json/src/reader/mod.rs
+++ b/arrow-json/src/reader/mod.rs
@@ -714,12 +714,12 @@ impl Decoder {
// First offset is null sentinel
let mut next_object = 1;
- let pos: Vec<_> = (0..tape.num_rows())
+ let pos = (0..tape.num_rows())
.map(|_| {
- let next = tape.next(next_object, "row").unwrap();
- std::mem::replace(&mut next_object, next)
+ let next = tape.next(next_object, "row")?;
+ Ok(std::mem::replace(&mut next_object, next))
})
- .collect();
+ .collect::<Result<Vec<_>, ArrowError>>()?;
let decoded = self.decoder.decode(&tape, &pos)?;
self.tape_decoder.clear();
diff --git a/arrow-schema/src/ffi.rs b/arrow-schema/src/ffi.rs
index 679ced42d5..064b4ac16d 100644
--- a/arrow-schema/src/ffi.rs
+++ b/arrow-schema/src/ffi.rs
@@ -433,54 +433,47 @@ impl FFI_ArrowSchema {
// On some platforms, c_char = u8, and on some, c_char = i8.
let buffer = self.metadata.cast::<u8>();
- fn next_four_bytes(buffer: *const u8, pos: &mut isize) -> [u8; 4] {
+ fn next_four_bytes(buffer: *const u8, pos: &mut usize) -> [u8; 4] {
// Safety: the caller advances `pos` only by the number of
bytes consumed,
// so `*pos..*pos+4` is always within the bounds of the
metadata buffer.
let out = unsafe {
[
- *buffer.offset(*pos),
- *buffer.offset(*pos + 1),
- *buffer.offset(*pos + 2),
- *buffer.offset(*pos + 3),
+ *buffer.add(*pos),
+ *buffer.add(*pos + 1),
+ *buffer.add(*pos + 2),
+ *buffer.add(*pos + 3),
]
};
*pos += 4;
out
}
- fn next_n_bytes(buffer: *const u8, pos: &mut isize, n: i32) ->
&[u8] {
+ fn next_n_bytes(buffer: *const u8, pos: &mut usize, n: usize) ->
&[u8] {
// Safety: same as `next_four_bytes`; `*pos..*pos+n` is within
the metadata buffer.
- let out = unsafe {
- std::slice::from_raw_parts(buffer.offset(*pos),
n.try_into().unwrap())
- };
- *pos += isize::try_from(n).unwrap();
+ let out = unsafe {
std::slice::from_raw_parts(buffer.add(*pos), n) };
+ *pos += n;
out
}
- let num_entries = i32::from_ne_bytes(next_four_bytes(buffer, &mut
pos));
- if num_entries < 0 {
- return Err(ArrowError::CDataInterface(
- "Negative number of metadata entries".to_string(),
- ));
+ /// A length read from the metadata, which the producer may have
got wrong.
+ fn checked_length(what: &str, length: i32) -> Result<usize,
ArrowError> {
+ usize::try_from(length).map_err(|_| {
+ ArrowError::CDataInterface(format!("Invalid {what} in
metadata: {length}"))
+ })
}
- let mut metadata =
- HashMap::with_capacity(num_entries.try_into().expect("Too many
metadata entries"));
+ let num_entries = i32::from_ne_bytes(next_four_bytes(buffer, &mut
pos));
+ let num_entries = checked_length("number of entries",
num_entries)?;
+
+ // The count comes from the producer, so do not preallocate all of
it
+ let mut metadata = HashMap::with_capacity(num_entries.min(128));
for _ in 0..num_entries {
let key_length = i32::from_ne_bytes(next_four_bytes(buffer,
&mut pos));
- if key_length < 0 {
- return Err(ArrowError::CDataInterface(
- "Negative key length in metadata".to_string(),
- ));
- }
+ let key_length = checked_length("key length", key_length)?;
let key = String::from_utf8(next_n_bytes(buffer, &mut pos,
key_length).to_vec())?;
let value_length = i32::from_ne_bytes(next_four_bytes(buffer,
&mut pos));
- if value_length < 0 {
- return Err(ArrowError::CDataInterface(
- "Negative value length in metadata".to_string(),
- ));
- }
+ let value_length = checked_length("value length",
value_length)?;
let value =
String::from_utf8(next_n_bytes(buffer, &mut pos,
value_length).to_vec())?;
metadata.insert(key, value);
diff --git a/arrow-select/src/dictionary.rs b/arrow-select/src/dictionary.rs
index 2f0418e1dc..80a18abb04 100644
--- a/arrow-select/src/dictionary.rs
+++ b/arrow-select/src/dictionary.rs
@@ -58,8 +58,8 @@ pub fn garbage_collect_dictionary<K: ArrowDictionaryKeyType>(
// Create a mapping from the old keys to the new keys, use a Vec for easy
indexing
let mut key_remap = vec![K::Native::ZERO; values.len()];
for (new_idx, old_idx) in mask.set_indices().enumerate() {
- key_remap[old_idx] = K::Native::from_usize(new_idx)
- .expect("new index should fit in K::Native, as old index was in
range");
+ key_remap[old_idx] =
+
K::Native::from_usize(new_idx).ok_or(ArrowError::DictionaryKeyOverflowError)?;
}
// ... and then build the new keys array
diff --git a/arrow-string/src/concat_elements.rs
b/arrow-string/src/concat_elements.rs
index 5abe66fd05..8934f62987 100644
--- a/arrow-string/src/concat_elements.rs
+++ b/arrow-string/src/concat_elements.rs
@@ -28,6 +28,11 @@ use arrow_data::{ArrayDataBuilder, MAX_INLINE_VIEW_LEN};
use arrow_schema::{ArrowError, DataType};
/// Returns the elementwise concatenation of a [`GenericByteArray`].
+///
+/// # Errors
+///
+/// Returns an error if the arrays have different lengths, or if the
concatenated
+/// data is too long for the offset type.
pub fn concat_elements_bytes<T: ByteArrayType>(
left: &GenericByteArray<T>,
right: &GenericByteArray<T>,
@@ -61,7 +66,10 @@ pub fn concat_elements_bytes<T: ByteArrayType>(
.extend_from_slice(&left_values[left_idx[0].as_usize()..left_idx[1].as_usize()]);
output_values
.extend_from_slice(&right_values[right_idx[0].as_usize()..right_idx[1].as_usize()]);
-
output_offsets.push(T::Offset::from_usize(output_values.len()).unwrap());
+ let output_len = output_values.len();
+ let offset =
+
T::Offset::from_usize(output_len).ok_or(ArrowError::OffsetOverflowError(output_len))?;
+ output_offsets.push(offset);
}
let builder = ArrayDataBuilder::new(T::DATA_TYPE)
@@ -110,6 +118,11 @@ pub fn concat_element_binary<Offset: OffsetSizeTrait>(
/// ```
///
/// An error will be returned if the [`StringArray`] are of different lengths
+///
+/// # Errors
+///
+/// Returns an error if the arrays have different lengths, or if the
concatenated
+/// data is too long for the offset type.
pub fn concat_elements_utf8_many<Offset: OffsetSizeTrait>(
arrays: &[&GenericStringArray<Offset>],
) -> Result<GenericStringArray<Offset>, ArrowError> {
@@ -159,7 +172,10 @@ pub fn concat_elements_utf8_many<Offset: OffsetSizeTrait>(
let index_end = offset.peek().unwrap().as_usize();
output_values.extend_from_slice(&values[index_start..index_end]);
});
- output_offsets.push(Offset::from_usize(output_values.len()).unwrap());
+ let output_len = output_values.len();
+ let offset =
+
Offset::from_usize(output_len).ok_or(ArrowError::OffsetOverflowError(output_len))?;
+ output_offsets.push(offset);
}
let builder =
ArrayDataBuilder::new(GenericStringArray::<Offset>::DATA_TYPE)
diff --git a/parquet/src/arrow/arrow_reader/mod.rs
b/parquet/src/arrow/arrow_reader/mod.rs
index ab3e015455..fed4541e14 100644
--- a/parquet/src/arrow/arrow_reader/mod.rs
+++ b/parquet/src/arrow/arrow_reader/mod.rs
@@ -1300,10 +1300,15 @@ impl<T: ChunkReader + 'static>
ParquetRecordBatchReaderBuilder<T> {
}
let bitset = match column_metadata.bloom_filter_length() {
- Some(_) => buffer.slice(
- (TryInto::<usize>::try_into(bitset_offset).unwrap()
- - TryInto::<usize>::try_into(offset).unwrap())..,
- ),
+ Some(_) => {
+ let bitset_start = bitset_offset
+ .checked_sub(offset)
+ .and_then(|start| usize::try_from(start).ok())
+ .ok_or_else(|| {
+ ParquetError::General("Bloom filter offset is
invalid".to_string())
+ })?;
+ buffer.slice(bitset_start..)
+ }
None => {
let bitset_length: usize =
header.num_bytes.try_into().map_err(|_| {
ParquetError::General("Bloom filter length is
invalid".to_string())
diff --git a/parquet/src/arrow/arrow_writer/mod.rs
b/parquet/src/arrow/arrow_writer/mod.rs
index ac38a11f10..267a52d4c9 100644
--- a/parquet/src/arrow/arrow_writer/mod.rs
+++ b/parquet/src/arrow/arrow_writer/mod.rs
@@ -1186,6 +1186,11 @@ impl ArrowColumnWriter {
}
/// Close this column returning the written [`ArrowColumnChunk`]
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if the column could not be finalised, or if another
thread
+ /// panicked while holding the column chunk. The caller cannot cause
either.
pub fn close(self) -> Result<ArrowColumnChunk> {
let distinct_count = self
.distinct_values_seen
@@ -1206,8 +1211,12 @@ impl ArrowColumnWriter {
c.close()?
}
};
- let chunk = Arc::try_unwrap(self.chunk).ok().unwrap();
- let data = chunk.into_inner().unwrap();
+ // Closing the writer above dropped the only other handle on the chunk.
+ let chunk = Arc::try_unwrap(self.chunk)
+ .map_err(|_| general_err!("Internal Error: the column chunk is
still shared"))?;
+ let data = chunk
+ .into_inner()
+ .map_err(|_| general_err!("The column chunk lock is poisoned"))?;
Ok(ArrowColumnChunk { data, close })
}
diff --git a/parquet/src/arrow/async_reader/mod.rs
b/parquet/src/arrow/async_reader/mod.rs
index 3eddd54935..16cbea14c2 100644
--- a/parquet/src/arrow/async_reader/mod.rs
+++ b/parquet/src/arrow/async_reader/mod.rs
@@ -630,10 +630,15 @@ impl<T: AsyncFileReader + Send + 'static>
ParquetRecordBatchStreamBuilder<T> {
}
let bitset = match column_metadata.bloom_filter_length() {
- Some(_) => buffer.slice(
- (TryInto::<usize>::try_into(bitset_offset).unwrap()
- - TryInto::<usize>::try_into(offset).unwrap())..,
- ),
+ Some(_) => {
+ let bitset_start = bitset_offset
+ .checked_sub(offset)
+ .and_then(|start| usize::try_from(start).ok())
+ .ok_or_else(|| {
+ ParquetError::General("Bloom filter offset is
invalid".to_string())
+ })?;
+ buffer.slice(bitset_start..)
+ }
None => {
let bitset_length: u64 =
header.num_bytes.try_into().map_err(|_| {
ParquetError::General("Bloom filter length is
invalid".to_string())
diff --git a/parquet/src/column/reader.rs b/parquet/src/column/reader.rs
index 42e8ccdc7a..9c2345b097 100644
--- a/parquet/src/column/reader.rs
+++ b/parquet/src/column/reader.rs
@@ -254,9 +254,7 @@ where
));
}
if levels_read == remaining_levels &&
self.has_record_delimiter {
- // Reached end of page, which implies records_read <
remaining_records
- // as otherwise would have stopped reading before
reaching the end
- assert!(records_read < remaining_records); // Sanity
check
+ check_partial_record_fits(records_read,
remaining_records)?;
records_read += reader.flush_partial() as usize;
}
(records_read, levels_read)
@@ -359,9 +357,7 @@ where
decoder.skip_rep_levels(remaining_records,
remaining_levels)?;
if levels_read == remaining_levels &&
self.has_record_delimiter {
- // Reached end of page, which implies records_read <
remaining_records
- // as otherwise would have stopped reading before
reaching the end
- assert!(records_read < remaining_records); // Sanity
check
+ check_partial_record_fits(records_read,
remaining_records)?;
records_read += decoder.flush_partial() as usize;
}
@@ -585,6 +581,22 @@ where
}
}
+/// Checks that a partial record can still be flushed into the caller's record
budget.
+///
+/// Reaching the end of a page with a record delimiter means the decoder
stopped because
+/// it ran out of levels, not records, so it cannot have used up the whole
budget. A page
+/// whose repetition levels disagree with its record count breaks that, and
flushing the
+/// partial record would then take the count past what was asked for.
+fn check_partial_record_fits(records_read: usize, remaining_records: usize) ->
Result<()> {
+ if remaining_records <= records_read {
+ return Err(general_err!(
+ "page ended after {records_read} record(s), which is already all
of the \
+ {remaining_records} record(s) asked for, so there is no partial
record to flush"
+ ));
+ }
+ Ok(())
+}
+
fn parse_v1_level(
max_level: i16,
num_buffered_values: u32,
diff --git a/parquet/src/file/writer.rs b/parquet/src/file/writer.rs
index 62c1157b7f..4b7a12c274 100644
--- a/parquet/src/file/writer.rs
+++ b/parquet/src/file/writer.rs
@@ -260,7 +260,7 @@ impl<W: Write + Send> SerializedFileWriter<W> {
self.row_group_index = self
.row_group_index
.checked_add(1)
- .expect("SerializedFileWriter::row_group_index overflowed");
+ .ok_or_else(|| ParquetError::General("Row group index
overflowed".to_string()))?;
let bloom_filter_position = self.properties().bloom_filter_position();
let row_groups = &mut self.row_groups;
