You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/19 14:43:48 UTC
(datafusion) branch gh-readonly-queue/main/pr-25481-de8803cd3f5703335c7d1dd9675329de80c56ae1 created (now 378a986c04)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-25481-de8803cd3f5703335c7d1dd9675329de80c56ae1
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 378a986c04 chore: Add helpers for computing `OffsetBuffer` span, length (#25481)
This branch includes the following new commits:
new 378a986c04 chore: Add helpers for computing `OffsetBuffer` span, length (#25481)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: chore: Add helpers for computing `OffsetBuffer` span, length (#25481)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25481-de8803cd3f5703335c7d1dd9675329de80c56ae1
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 378a986c041a672e658896eb50c7dab50992d895
Author: Neil Conway <ne...@gmail.com>
AuthorDate: Sat Sep 19 14:43:28 2026 +0000
chore: Add helpers for computing `OffsetBuffer` span, length (#25481)
## Which issue does this PR close?
- N/A
## Rationale for this change
Code that works with Arrow arrays often neglects to account for slicing:
that is, all of the array's backing values might not be visible, so code
should generally not iterate over values outside the visible range, or
assume that the raw length of the backing array is equivalent to the
logical length of the array.
To simplify common coding patterns in this area and to encourage correct
usage, this PR introduces two helper functions:
* `offset_span` returns `(first_offset, value_len)`, the "visible" span
described by an `OffsetBuffer`
* `offset_span_len` returns `value_len`, the number of values covered by
an `OffsetBuffer`
This replaces typical boilerplate required to handle sliced arrays.
## What changes are included in this PR?
* Introduce two new helper functions
* Adopt those helper functions through DataFusion, replacing
hand-written offset arithmetic
* Simplify or eliminate several helpers that did similar things already
(e.g., `estimate_byte_data_size` in `encoding/inner.rs`)
* Add unit tests
## What is the testing strategy for this PR?
Existing tests pass; new unit tests added.
## Are there any user-facing changes?
No. No behavior changes.
---
datafusion/common/src/hash_utils.rs | 24 +++++------
datafusion/common/src/utils/mod.rs | 56 +++++++++++++++++++++----
datafusion/functions-aggregate/src/array_agg.rs | 12 ++----
datafusion/functions-nested/src/position.rs | 27 +++++-------
datafusion/functions-nested/src/remove.rs | 6 +--
datafusion/functions-nested/src/replace.rs | 9 ++--
datafusion/functions-nested/src/set_ops.rs | 5 +--
datafusion/functions-nested/src/sort.rs | 17 ++++----
datafusion/functions/src/core/overlay.rs | 18 ++------
datafusion/functions/src/encoding/inner.rs | 20 ++++-----
datafusion/functions/src/regex/regexpreplace.rs | 9 ++--
datafusion/functions/src/string/common.rs | 18 +++-----
datafusion/functions/src/unicode/initcap.rs | 5 +--
datafusion/functions/src/unicode/substrindex.rs | 21 ++++------
datafusion/spark/src/function/string/quote.rs | 9 +---
15 files changed, 117 insertions(+), 139 deletions(-)
diff --git a/datafusion/common/src/hash_utils.rs b/datafusion/common/src/hash_utils.rs
index 1e6d96ecfd..8da165e8fb 100644
--- a/datafusion/common/src/hash_utils.rs
+++ b/datafusion/common/src/hash_utils.rs
@@ -89,6 +89,8 @@ use crate::cast::{
};
use crate::error::Result;
use crate::error::{_internal_datafusion_err, _internal_err};
+#[cfg(not(feature = "force_hash_collisions"))]
+use crate::utils::offset_span;
use std::cell::RefCell;
mod build_hasher;
@@ -689,9 +691,7 @@ fn hash_map_array(
let offsets = array.offsets();
// Create hashes for each entry in each row
- let first_offset = offsets.first() as usize;
- let last_offset = offsets.last() as usize;
- let entries_len = last_offset - first_offset;
+ let (first_offset, entries_len) = offset_span(offsets);
// Only hash the entries that are actually referenced
let mut values_hashes = vec![0u64; entries_len];
@@ -739,15 +739,11 @@ fn hash_list_array<OffsetSize>(
where
OffsetSize: OffsetSizeTrait,
{
- // In case values is sliced, hash only the bytes used by the offsets of this ListArray
- let first_offset = array.value_offsets().first().copied().unwrap_or_default();
- let last_offset = array.value_offsets().last().copied().unwrap_or_default();
- let value_bytes_len = (last_offset - first_offset).as_usize();
- let mut values_hashes = vec![0u64; value_bytes_len];
+ // Hash only the child values referenced by this ListArray's offsets.
+ let (first_offset, values_len) = offset_span(array.offsets());
+ let mut values_hashes = vec![0u64; values_len];
child_hashing.create_hashes(
- [array
- .values()
- .slice(first_offset.as_usize(), value_bytes_len)],
+ [array.values().slice(first_offset, values_len)],
&mut values_hashes,
)?;
@@ -756,8 +752,8 @@ where
{
if array.is_valid(i) {
let hash = &mut hashes_buffer[i];
- for values_hash in &values_hashes[(*start - first_offset).as_usize()
- ..(*stop - first_offset).as_usize()]
+ for values_hash in &values_hashes
+ [start.as_usize() - first_offset..stop.as_usize() - first_offset]
{
*hash = combine_hashes(*hash, *values_hash);
}
@@ -771,7 +767,7 @@ where
.zip(hashes_buffer.iter_mut())
{
for values_hash in &values_hashes
- [(*start - first_offset).as_usize()..(*stop - first_offset).as_usize()]
+ [start.as_usize() - first_offset..stop.as_usize() - first_offset]
{
*hash = combine_hashes(*hash, *values_hash);
}
diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs
index fba76f949f..bf87b9a788 100644
--- a/datafusion/common/src/utils/mod.rs
+++ b/datafusion/common/src/utils/mod.rs
@@ -1227,6 +1227,36 @@ pub fn take_function_args<const N: usize, T>(
})
}
+/// Returns the number of values covered by an offset buffer.
+///
+/// Slicing a variable-length array narrows its offsets but retains its entire
+/// values buffer. Use this span to size output buffers: it counts child elements
+/// for lists and bytes for strings/binary arrays, including values in null rows.
+/// An empty array still has one offset and therefore a span of zero.
+#[inline]
+pub fn offset_span_len<O: ArrowNativeType>(offsets: &OffsetBuffer<O>) -> usize {
+ offset_span(offsets).1
+}
+
+/// Returns the start and length of the values covered by an offset buffer.
+///
+/// The returned pair can be passed to [`Array::slice`] to select the visible
+/// child values of a sliced list or map. For strings/binary arrays, the start
+/// and length are measured in bytes. Values in null rows are included.
+/// An empty array has a length of zero but may have a nonzero start.
+///
+/// ```
+/// # use arrow::buffer::OffsetBuffer;
+/// # use datafusion_common::utils::offset_span;
+/// let offsets = OffsetBuffer::new(vec![100_i32, 103, 108].into());
+/// assert_eq!(offset_span(&offsets), (100, 8));
+/// ```
+#[inline]
+pub fn offset_span<O: ArrowNativeType>(offsets: &OffsetBuffer<O>) -> (usize, usize) {
+ let start = offsets.first().as_usize();
+ (start, offsets.last().as_usize() - start)
+}
+
/// Returns the inner values of a list, or an error otherwise
/// For [`ListArray`] and [`LargeListArray`], if it's sliced, it returns a
/// sliced array too. Therefore, too reconstruct a list using it,
@@ -1244,13 +1274,10 @@ pub fn list_values(array: &dyn Array) -> Result<ArrayRef> {
fn sliced_list_values<O: OffsetSizeTrait>(list: &GenericListArray<O>) -> ArrayRef {
let values = list.values();
- let offsets = list.offsets();
-
- let first = offsets.first().as_usize();
- let last = offsets.last().as_usize();
+ let (start, len) = offset_span(list.offsets());
- if first != 0 || last != values.len() {
- return values.slice(first, last - first);
+ if start != 0 || len != values.len() {
+ return values.slice(start, len);
}
Arc::clone(values)
@@ -1300,9 +1327,9 @@ fn truncate_list_nulls<O: OffsetSizeTrait>(
if valid_or_empty.has_false() {
let array_data = list.values().to_data();
let offsets = list.offsets();
- let capacity = offsets[offsets.len() - 1] - offsets[0];
+ let capacity = offset_span_len(offsets);
let mut mutable_array_data =
- MutableArrayData::new(vec![&array_data], false, capacity.as_usize());
+ MutableArrayData::new(vec![&array_data], false, capacity);
let (valid_or_empty, _nulls) = valid_or_empty.into_parts();
@@ -1508,6 +1535,19 @@ mod tests {
#[cfg(feature = "sql")]
use sqlparser::ast::Ident;
+ #[test]
+ fn test_offset_span() {
+ let offsets = OffsetBuffer::new(vec![0_i32, 5, 8, 8, 12].into());
+ assert_eq!(offset_span(&offsets), (0, 12));
+ assert_eq!(offset_span(&offsets.slice(1, 2)), (5, 3));
+ assert_eq!(offset_span_len(&offsets.slice(1, 2)), 3);
+ assert_eq!(offset_span(&offsets.slice(2, 1)), (8, 0));
+ assert_eq!(offset_span(&offsets.slice(4, 0)), (12, 0));
+ assert_eq!(offset_span(&OffsetBuffer::<i64>::new_empty()), (0, 0));
+ let large = OffsetBuffer::new(vec![i64::MAX - 10, i64::MAX].into());
+ assert_eq!(offset_span(&large), ((i64::MAX - 10) as usize, 10));
+ }
+
#[test]
fn test_usize_wire_conversions() {
assert_eq!(usize_from_wire(42_u64, "SomeExec", "fetch").unwrap(), 42);
diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs
index cc00e8f45b..546fad9738 100644
--- a/datafusion/functions-aggregate/src/array_agg.rs
+++ b/datafusion/functions-aggregate/src/array_agg.rs
@@ -35,8 +35,8 @@ use arrow_select::dictionary::garbage_collect_any_dictionary;
use datafusion_common::cast::as_list_array;
use datafusion_common::hash_utils::{RandomState, create_hashes};
use datafusion_common::scalar::copy_array_data;
-use datafusion_common::utils::SingleRowListArrayBuilder;
use datafusion_common::utils::proxy::HashTableAllocExt;
+use datafusion_common::utils::{SingleRowListArrayBuilder, offset_span};
use datafusion_common::{
Result, ScalarValue, assert_eq_or_internal_err, exec_err, internal_err,
};
@@ -284,19 +284,15 @@ impl ArrayAggAccumulator {
/// This function will return the underlying list array values if all valid values are consecutive without gaps (i.e. no null value point to a non-empty list)
/// If there are gaps but only in the end of the list array, the function will return the values without the null values in the end
fn get_optional_values_to_merge_as_is(list_array: &ListArray) -> Option<ArrayRef> {
- let offsets = list_array.value_offsets();
- // Offsets always have at least 1 value
- let initial_offset = offsets[0];
+ let offsets = list_array.offsets();
let null_count = list_array.null_count();
// If no nulls than just use the fast path
// This is ok as the state is a ListArray rather than a ListViewArray so all the values are consecutive
if null_count == 0 {
// According to Arrow specification, the first offset can be non-zero
- let list_values = list_array.values().slice(
- initial_offset as usize,
- (offsets[offsets.len() - 1] - initial_offset) as usize,
- );
+ let (start, len) = offset_span(offsets);
+ let list_values = list_array.values().slice(start, len);
return Some(list_values);
}
diff --git a/datafusion/functions-nested/src/position.rs b/datafusion/functions-nested/src/position.rs
index 9ef8142fa4..9ed9c2651e 100644
--- a/datafusion/functions-nested/src/position.rs
+++ b/datafusion/functions-nested/src/position.rs
@@ -40,6 +40,7 @@ use arrow::array::{
use datafusion_common::cast::{
as_generic_list_array, as_int64_array, as_large_list_array, as_list_array,
};
+use datafusion_common::utils::offset_span;
use datafusion_common::{
Result, exec_datafusion_err, exec_err, utils::take_function_args,
};
@@ -255,11 +256,8 @@ fn array_position_scalar<O: OffsetSizeTrait>(
// ListArrays, values() returns the full underlying array but only
// elements between the first and last offset are referenced.
let offsets = haystack.offsets();
- let first_offset = offsets[0].as_usize();
- let last_offset = offsets[haystack.len()].as_usize();
- let visible_values = haystack
- .values()
- .slice(first_offset, last_offset - first_offset);
+ let (first_offset, values_len) = offset_span(offsets);
+ let visible_values = haystack.values().slice(first_offset, values_len);
// `not_distinct` treats NULL=NULL as true, matching the semantics of
// `array_position`.
@@ -271,9 +269,9 @@ fn array_position_scalar<O: OffsetSizeTrait>(
// Match positions are relative to visible_values (0-based), so
// subtract first_offset from each offset when comparing.
- for i in 0..haystack.len() {
- let start = offsets[i].as_usize() - first_offset;
- let end = offsets[i + 1].as_usize() - first_offset;
+ for (i, window) in offsets.windows(2).enumerate() {
+ let start = window[0].as_usize() - first_offset;
+ let end = window[1].as_usize() - first_offset;
if validity.is_some_and(|v| v.is_null(i)) {
// Null row -> null output; advance past matches in range
@@ -554,11 +552,8 @@ fn array_positions_scalar<O: OffsetSizeTrait>(
// ListArrays, values() returns the full underlying array but only
// elements between the first and last offset are referenced.
let offsets = haystack.offsets();
- let first_offset = offsets[0].as_usize();
- let last_offset = offsets[num_rows].as_usize();
- let visible_values = haystack
- .values()
- .slice(first_offset, last_offset - first_offset);
+ let (first_offset, values_len) = offset_span(offsets);
+ let visible_values = haystack.values().slice(first_offset, values_len);
// `not_distinct` treats NULL=NULL as true, matching the semantics of
// `array_positions`.
@@ -573,9 +568,9 @@ fn array_positions_scalar<O: OffsetSizeTrait>(
// Match positions are relative to visible_values (0-based), so
// subtract first_offset from each offset when comparing.
- for i in 0..num_rows {
- let start = offsets[i].as_usize() - first_offset;
- let end = offsets[i + 1].as_usize() - first_offset;
+ for (i, window) in offsets.windows(2).enumerate() {
+ let start = window[0].as_usize() - first_offset;
+ let end = window[1].as_usize() - first_offset;
if validity.is_some_and(|v| v.is_null(i)) {
// Null row -> null output; advance past matches in range.
diff --git a/datafusion/functions-nested/src/remove.rs b/datafusion/functions-nested/src/remove.rs
index fa4127c17c..7470c60b85 100644
--- a/datafusion/functions-nested/src/remove.rs
+++ b/datafusion/functions-nested/src/remove.rs
@@ -25,7 +25,7 @@ use arrow::array::{
use arrow::buffer::OffsetBuffer;
use arrow::datatypes::{DataType, FieldRef};
use datafusion_common::cast::as_int64_array;
-use datafusion_common::utils::ListCoercion;
+use datafusion_common::utils::{ListCoercion, offset_span};
use datafusion_common::{
Result, ScalarValue, exec_err, internal_err, utils::take_function_args,
};
@@ -576,9 +576,7 @@ fn general_remove_with_scalar<OffsetSize: OffsetSizeTrait>(
};
let list_offsets = list_array.offsets();
- let first_offset = list_offsets[0].to_usize().unwrap();
- let last_offset = list_offsets[list_offsets.len() - 1].to_usize().unwrap();
- let values_range_len = last_offset - first_offset;
+ let (first_offset, values_range_len) = offset_span(list_offsets);
let values_slice = list_array.values().slice(first_offset, values_range_len);
let original_data = values_slice.to_data();
let mut offsets = Vec::<OffsetSize>::with_capacity(list_array.len() + 1);
diff --git a/datafusion/functions-nested/src/replace.rs b/datafusion/functions-nested/src/replace.rs
index 4bfd0c0dbe..af45058801 100644
--- a/datafusion/functions-nested/src/replace.rs
+++ b/datafusion/functions-nested/src/replace.rs
@@ -24,7 +24,7 @@ use arrow::array::{
use arrow::buffer::OffsetBuffer;
use arrow::datatypes::{DataType, Field, FieldRef};
use datafusion_common::cast::as_int64_array;
-use datafusion_common::utils::ListCoercion;
+use datafusion_common::utils::{ListCoercion, offset_span};
use datafusion_common::{
Result, ScalarValue, exec_err, internal_err, utils::take_function_args,
};
@@ -582,11 +582,8 @@ fn general_replace_with_scalar<O: OffsetSizeTrait>(
)?));
}
- let first_offset = list_array.offsets()[0].to_usize().unwrap();
- let last_offset = list_array.offsets()[list_array.len()].to_usize().unwrap();
- let visible_values = list_array
- .values()
- .slice(first_offset, last_offset - first_offset);
+ let (first_offset, values_len) = offset_span(list_array.offsets());
+ let visible_values = list_array.values().slice(first_offset, values_len);
let to_array = scalar_to.to_array_of_size(1)?;
let original_data = visible_values.to_data();
diff --git a/datafusion/functions-nested/src/set_ops.rs b/datafusion/functions-nested/src/set_ops.rs
index 5cd4819cae..66cce98b09 100644
--- a/datafusion/functions-nested/src/set_ops.rs
+++ b/datafusion/functions-nested/src/set_ops.rs
@@ -28,7 +28,7 @@ use arrow::datatypes::DataType::{LargeList, List, Null};
use arrow::datatypes::{DataType, Field, FieldRef};
use arrow::row::{RowConverter, SortField};
use datafusion_common::cast::{as_large_list_array, as_list_array};
-use datafusion_common::utils::{ListCoercion, normalize_float_zero};
+use datafusion_common::utils::{ListCoercion, normalize_float_zero, offset_span};
use datafusion_common::{
Result, assert_eq_or_internal_err, exec_err, internal_err, utils::take_function_args,
};
@@ -332,8 +332,7 @@ impl Display for SetOp {
pub(crate) fn normalize_visible_values<OffsetSize: OffsetSizeTrait>(
array: &GenericListArray<OffsetSize>,
) -> ArrayRef {
- let first = array.offsets()[0].as_usize();
- let len = array.offsets()[array.len()].as_usize() - first;
+ let (first, len) = offset_span(array.offsets());
if first == 0 && len == array.values().len() {
normalize_float_zero(array.values())
} else {
diff --git a/datafusion/functions-nested/src/sort.rs b/datafusion/functions-nested/src/sort.rs
index fadc5da4b4..bccea76867 100644
--- a/datafusion/functions-nested/src/sort.rs
+++ b/datafusion/functions-nested/src/sort.rs
@@ -28,7 +28,7 @@ use arrow::datatypes::{ArrowNativeTypeOp, DataType, FieldRef};
use arrow::row::{RowConverter, Rows, SortField};
use arrow::{compute, compute::SortOptions, downcast_primitive_array};
use datafusion_common::cast::{as_large_list_array, as_list_array, as_string_array};
-use datafusion_common::utils::ListCoercion;
+use datafusion_common::utils::{ListCoercion, offset_span};
use datafusion_common::{Result, exec_err};
use datafusion_expr::{
ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue, Documentation,
@@ -292,15 +292,14 @@ fn sort_row<N: ArrowNativeTypeOp>(row: &mut [N], descending: bool) {
/// values into a single `Vec` and sorts each row's slice in-place.
fn sort_rows_no_nulls<N: ArrowNativeTypeOp, O: OffsetSizeTrait>(
src_values: &[N],
- offsets: &[O],
+ offsets: &OffsetBuffer<O>,
list_nulls: Option<&NullBuffer>,
descending: bool,
) -> Vec<N> {
- let values_start = offsets[0].as_usize();
- let values_end = offsets[offsets.len() - 1].as_usize();
+ let (values_start, total_values) = offset_span(offsets);
// Copy all values into a mutable buffer
- let mut values = src_values[values_start..values_end].to_vec();
+ let mut values = src_values[values_start..values_start + total_values].to_vec();
for (row_index, window) in offsets.windows(2).enumerate() {
if list_nulls.is_some_and(|n| n.is_null(row_index)) {
@@ -318,13 +317,12 @@ fn sort_rows_no_nulls<N: ArrowNativeTypeOp, O: OffsetSizeTrait>(
fn sort_rows_with_nulls<N: ArrowNativeTypeOp, O: OffsetSizeTrait>(
src_values: &[N],
src_nulls: &NullBuffer,
- offsets: &[O],
+ offsets: &OffsetBuffer<O>,
list_nulls: Option<&NullBuffer>,
descending: bool,
nulls_first: bool,
) -> (Vec<N>, NullBuffer) {
- let values_start = offsets[0].as_usize();
- let total_values = offsets[offsets.len() - 1].as_usize() - values_start;
+ let (values_start, total_values) = offset_span(offsets);
let mut out_values: Vec<N> = vec![N::default(); total_values];
let mut validity = BooleanBufferBuilder::new(total_values);
@@ -385,8 +383,7 @@ fn array_sort_non_primitive<OffsetSize: OffsetSizeTrait>(
let row_count = list_array.len();
let values = list_array.values();
let offsets = list_array.offsets();
- let values_start = offsets[0].as_usize();
- let total_values = offsets[row_count].as_usize() - values_start;
+ let (values_start, total_values) = offset_span(offsets);
let converter = RowConverter::new(vec![SortField::new_with_options(
values.data_type().clone(),
diff --git a/datafusion/functions/src/core/overlay.rs b/datafusion/functions/src/core/overlay.rs
index c1f3353a8f..ae5bc989e8 100644
--- a/datafusion/functions/src/core/overlay.rs
+++ b/datafusion/functions/src/core/overlay.rs
@@ -16,8 +16,7 @@
// under the License.
use arrow::array::{
- Array, ArrayRef, GenericStringArray, Int64Array, OffsetSizeTrait, StringArrayType,
- StringViewArray,
+ Array, ArrayRef, Int64Array, OffsetSizeTrait, StringArrayType, StringViewArray,
};
use arrow::buffer::NullBuffer;
use arrow::datatypes::DataType;
@@ -29,6 +28,7 @@ use crate::utils::{make_scalar_function, utf8_to_str_type};
use datafusion_common::cast::{
as_generic_string_array, as_int64_array, as_string_view_array,
};
+use datafusion_common::utils::offset_span_len;
use datafusion_common::{Result, exec_err};
use datafusion_expr::{ColumnarValue, Documentation, TypeSignature, Volatility};
use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature};
@@ -236,8 +236,8 @@ fn overlay<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
} else {
let string_array = as_generic_string_array::<T>(&args[0])?;
let characters_array = as_generic_string_array::<T>(&args[1])?;
- let data_capacity = visible_offset_bytes(string_array)
- .saturating_add(visible_offset_bytes(characters_array));
+ let data_capacity = offset_span_len(string_array.offsets())
+ .saturating_add(offset_span_len(characters_array.offsets()));
let builder = GenericStringArrayBuilder::<T>::with_capacity(
string_array.len(),
data_capacity,
@@ -311,16 +311,6 @@ fn visible_view_bytes(array: &StringViewArray) -> usize {
array.lengths().map(|l| l as usize).sum()
}
-/// Bytes referenced by the visible window of `array`, derived from the offset
-/// buffer.
-fn visible_offset_bytes<T: OffsetSizeTrait>(array: &GenericStringArray<T>) -> usize {
- let offsets = array.value_offsets();
- // `value_offsets()` always has `array.len() + 1` entries (≥1).
- let first = offsets.first().copied().unwrap_or_default();
- let last = offsets.last().copied().unwrap_or_default();
- last.as_usize() - first.as_usize()
-}
-
#[cfg(test)]
mod tests {
use std::sync::Arc;
diff --git a/datafusion/functions/src/encoding/inner.rs b/datafusion/functions/src/encoding/inner.rs
index 8b57033fa0..43d75f3061 100644
--- a/datafusion/functions/src/encoding/inner.rs
+++ b/datafusion/functions/src/encoding/inner.rs
@@ -35,7 +35,7 @@ use datafusion_common::{
types::{NativeType, logical_string},
utils::{
hex::{HexCase, encode_bytes, encode_bytes_to_slice},
- take_function_args,
+ offset_span_len, take_function_args,
},
};
use datafusion_expr::{
@@ -272,19 +272,13 @@ fn decode_scalar(value: &ScalarValue, encoding: Encoding) -> Result<ColumnarValu
}
}
-/// Estimate how many bytes are actually represented by the array; in case the
-/// the array slices it's internal buffer, this returns the byte size of that slice
-/// but not the byte size of the entire buffer.
+/// Estimates the number of input bytes covered by the array's visible rows.
///
-/// This is an estimation only as it can estimate higher if null slots are non-zero
-/// sized.
+/// Sliced arrays exclude bytes outside their visible offset span. This can
+/// overestimate the bytes processed because null rows may cover nonempty byte
+/// ranges.
fn estimate_byte_data_size<O: OffsetSizeTrait>(array: &GenericBinaryArray<O>) -> usize {
- let offsets = array.value_offsets();
- // Unwraps are safe as should always have 1 element in offset buffer
- let start = *offsets.first().unwrap();
- let end = *offsets.last().unwrap();
- let data_size = end - start;
- data_size.as_usize()
+ offset_span_len(array.offsets())
}
fn decode_array(array: &ArrayRef, encoding: Encoding) -> Result<ColumnarValue> {
@@ -544,7 +538,7 @@ mod tests {
let size = estimate_byte_data_size(&array);
assert_eq!(size, 15);
- // Offsets starting at 0, but don't count entire data buffer size
+ // Nonzero starting offset, including bytes covered by null slots
let array = BinaryArray::new(
OffsetBuffer::new(vec![50, 51, 51, 60, 80, 81].into()),
vec![0; 100].into(),
diff --git a/datafusion/functions/src/regex/regexpreplace.rs b/datafusion/functions/src/regex/regexpreplace.rs
index 6ad138de15..647abfed59 100644
--- a/datafusion/functions/src/regex/regexpreplace.rs
+++ b/datafusion/functions/src/regex/regexpreplace.rs
@@ -32,6 +32,7 @@ use datafusion_common::cast::{
};
use datafusion_common::exec_err;
use datafusion_common::plan_err;
+use datafusion_common::utils::offset_span_len;
use datafusion_common::{Result, cast::as_generic_string_array, internal_err};
use datafusion_expr::ColumnarValue;
use datafusion_expr::TypeSignature;
@@ -554,12 +555,8 @@ fn regexp_replace_static_pattern_replace<T: OffsetSizeTrait>(
// We are going to create the underlying string buffer from its parts
// to be able to re-use the existing null buffer for sparse arrays.
- let mut vals = BufferBuilder::<u8>::new({
- let offsets = string_array.value_offsets();
- (offsets[string_array.len()] - offsets[0])
- .to_usize()
- .unwrap()
- });
+ let mut vals =
+ BufferBuilder::<u8>::new(offset_span_len(string_array.offsets()));
let mut new_offsets = BufferBuilder::<T>::new(string_array.len() + 1);
new_offsets.append(T::zero());
diff --git a/datafusion/functions/src/string/common.rs b/datafusion/functions/src/string/common.rs
index c13f724648..0bc63afa6a 100644
--- a/datafusion/functions/src/string/common.rs
+++ b/datafusion/functions/src/string/common.rs
@@ -31,6 +31,7 @@ use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer};
use arrow::datatypes::DataType;
use datafusion_common::Result;
use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
+use datafusion_common::utils::{offset_span, offset_span_len};
use datafusion_common::{ScalarValue, exec_err};
use datafusion_expr::ColumnarValue;
@@ -282,11 +283,7 @@ where
F: for<'a> FnMut(usize, &'a str) -> &'a str,
{
let len = string_array.len();
- let input_offsets = string_array.value_offsets();
- let start = input_offsets.first().unwrap().as_usize();
- let end = input_offsets.last().unwrap().as_usize();
-
- let mut values: Vec<u8> = Vec::with_capacity(end - start);
+ let mut values: Vec<u8> = Vec::with_capacity(offset_span_len(string_array.offsets()));
let mut offsets: Vec<T> = Vec::with_capacity(len + 1);
offsets.push(T::usize_as(0));
@@ -540,10 +537,7 @@ fn case_conversion_array<O: OffsetSizeTrait>(
// Values contain non-ASCII.
let item_len = string_array.len();
- let offsets = string_array.value_offsets();
- let start = offsets.first().unwrap().as_usize();
- let end = offsets.last().unwrap().as_usize();
- let capacity = (end - start) + PRE_ALLOC_BYTES;
+ let capacity = offset_span_len(string_array.offsets()) + PRE_ALLOC_BYTES;
// Null-preserving: reuse the input null buffer as the output null buffer.
let nulls = string_array.nulls().cloned();
let mut builder = GenericStringArrayBuilder::<O>::with_capacity(item_len, capacity);
@@ -705,10 +699,8 @@ fn case_conversion_ascii_array<O: OffsetSizeTrait>(
string_array: &GenericStringArray<O>,
lower: bool,
) -> Result<ArrayRef> {
- let value_offsets = string_array.value_offsets();
- let start = value_offsets.first().unwrap().as_usize();
- let end = value_offsets.last().unwrap().as_usize();
- let relevant = &string_array.value_data()[start..end];
+ let (start, len) = offset_span(string_array.offsets());
+ let relevant = &string_array.value_data()[start..start + len];
let converted: Vec<u8> = if lower {
relevant.iter().map(u8::to_ascii_lowercase).collect()
diff --git a/datafusion/functions/src/unicode/initcap.rs b/datafusion/functions/src/unicode/initcap.rs
index d18aaa945e..d521ca0411 100644
--- a/datafusion/functions/src/unicode/initcap.rs
+++ b/datafusion/functions/src/unicode/initcap.rs
@@ -24,6 +24,7 @@ use arrow::datatypes::DataType;
use crate::strings::{GenericStringArrayBuilder, StringViewArrayBuilder};
use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
use datafusion_common::types::logical_string;
+use datafusion_common::utils::offset_span_len;
use datafusion_common::{Result, ScalarValue, exec_err};
use datafusion_expr::{
Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs,
@@ -204,12 +205,10 @@ fn initcap_ascii_array<T: OffsetSizeTrait>(
) -> ArrayRef {
let offsets = string_array.offsets();
let src = string_array.value_data();
- let first_offset = offsets.first().as_usize();
- let last_offset = offsets.last().as_usize();
// For sliced arrays, only convert the visible bytes, not the entire input
// buffer.
- let mut out = Vec::with_capacity(last_offset - first_offset);
+ let mut out = Vec::with_capacity(offset_span_len(offsets));
for window in offsets.windows(2) {
let start = window[0].as_usize();
diff --git a/datafusion/functions/src/unicode/substrindex.rs b/datafusion/functions/src/unicode/substrindex.rs
index 51cd615b92..614e5db32b 100644
--- a/datafusion/functions/src/unicode/substrindex.rs
+++ b/datafusion/functions/src/unicode/substrindex.rs
@@ -18,8 +18,8 @@
use std::sync::Arc;
use arrow::array::{
- Array, ArrayRef, AsArray, ByteView, GenericStringArray, OffsetSizeTrait,
- PrimitiveArray, StringArrayType, StringViewArray, make_view, new_null_array,
+ Array, ArrayRef, AsArray, ByteView, OffsetSizeTrait, PrimitiveArray, StringArrayType,
+ StringViewArray, make_view, new_null_array,
};
use arrow::buffer::ScalarBuffer;
use arrow::datatypes::{DataType, Int64Type};
@@ -27,6 +27,7 @@ use arrow_buffer::NullBuffer;
use crate::strings::GenericStringArrayBuilder;
use crate::utils::make_scalar_function;
+use datafusion_common::utils::offset_span_len;
use datafusion_common::{
Result, ScalarValue, exec_datafusion_err, exec_err, utils::take_function_args,
};
@@ -153,7 +154,7 @@ fn substr_index(args: &[ArrayRef]) -> Result<ArrayRef> {
count_array,
GenericStringArrayBuilder::<i32>::with_capacity(
string_array.len(),
- visible_string_bytes(string_array),
+ offset_span_len(string_array.offsets()),
),
)
}
@@ -167,7 +168,7 @@ fn substr_index(args: &[ArrayRef]) -> Result<ArrayRef> {
count_array,
GenericStringArrayBuilder::<i64>::with_capacity(
string_array.len(),
- visible_string_bytes(string_array),
+ offset_span_len(string_array.offsets()),
),
)
}
@@ -231,7 +232,7 @@ fn substr_index_scalar(
count,
GenericStringArrayBuilder::<i32>::with_capacity(
arr.len(),
- visible_string_bytes(arr),
+ offset_span_len(arr.offsets()),
),
)
}
@@ -243,7 +244,7 @@ fn substr_index_scalar(
count,
GenericStringArrayBuilder::<i64>::with_capacity(
arr.len(),
- visible_string_bytes(arr),
+ offset_span_len(arr.offsets()),
),
)
}
@@ -253,14 +254,6 @@ fn substr_index_scalar(
Ok(ColumnarValue::Array(result))
}
-#[inline]
-fn visible_string_bytes<T: OffsetSizeTrait>(
- string_array: &GenericStringArray<T>,
-) -> usize {
- let offsets = string_array.value_offsets();
- offsets[offsets.len() - 1].as_usize() - offsets[0].as_usize()
-}
-
fn substr_index_general<'a, S, O>(
string_array: S,
delimiter_array: S,
diff --git a/datafusion/spark/src/function/string/quote.rs b/datafusion/spark/src/function/string/quote.rs
index b4fd5e3827..8fafc42dfd 100644
--- a/datafusion/spark/src/function/string/quote.rs
+++ b/datafusion/spark/src/function/string/quote.rs
@@ -19,7 +19,7 @@ use arrow::array::{ArrayRef, GenericStringBuilder, OffsetSizeTrait};
use arrow::datatypes::DataType;
use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
use datafusion_common::types::{NativeType, logical_string};
-use datafusion_common::utils::take_function_args;
+use datafusion_common::utils::{offset_span_len, take_function_args};
use datafusion_common::{Result, exec_err};
use datafusion_expr::{
Coercion, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature,
@@ -91,12 +91,7 @@ fn spark_quote_inner(arg: &[ArrayRef]) -> Result<ArrayRef> {
fn quote_array<T: OffsetSizeTrait>(array: &ArrayRef) -> Result<ArrayRef> {
let str_array = as_generic_string_array::<T>(array)?;
- // Slicing an array keeps the whole value buffer and narrows only the
- // offsets, so measure the data through the offsets rather than through
- // `value_data()`. An offset buffer holds one more entry than the array has
- // rows, so it is never empty.
- let offsets = str_array.value_offsets();
- let data_len = offsets.last().unwrap().as_usize() - offsets[0].as_usize();
+ let data_len = offset_span_len(str_array.offsets());
Ok(quote_impl::<T, _>(str_array.iter(), data_len))
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org