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/10/11 15:08:10 UTC
(datafusion-comet) 01/01: perf: presize collections and reuse buffers in native loops (#6861)
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-6861-0691bd39f75da61e02aa23ce188de141bd623252
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 6296625b9081b1ec0186d23c9c7f4222cd905f30
Author: Oleks V <co...@users.noreply.github.com>
AuthorDate: Sun Oct 11 15:08:00 2026 +0000
perf: presize collections and reuse buffers in native loops (#6861)
Give collections a capacity where their final size is known, and stop
allocating a fresh structure in every iteration of per-row and per-batch
loops where one can be reused, borrowed or moved instead.
---
native/common/src/bin/analyze_trace.rs | 10 ++--
.../src/execution/operators/iceberg_dictionary.rs | 2 +-
native/core/src/execution/operators/merge_rows.rs | 3 +-
.../core/src/execution/operators/parquet_writer.rs | 3 +-
native/core/src/execution/planner.rs | 4 +-
native/core/src/execution/spark_plan.rs | 6 +--
native/core/src/parquet/parquet_support.rs | 5 +-
native/core/src/parquet/schema_adapter.rs | 6 +--
native/operators/src/expand.rs | 2 +-
native/operators/src/explode.rs | 53 ++++++++++------------
native/shuffle/src/partitioners/multi_partition.rs | 10 ++--
native/shuffle/src/schema_align.rs | 4 +-
native/shuffle/src/spark_unsafe/row.rs | 46 +++++++++++++++----
native/spark-expr/src/agg_funcs/percentile.rs | 6 ++-
.../src/bloom_filter/bloom_filter_agg.rs | 8 ++++
.../spark-expr/src/bloom_filter/spark_bit_array.rs | 4 +-
.../src/bloom_filter/spark_bloom_filter.rs | 17 ++++---
.../spark-expr/src/json_funcs/json_array_length.rs | 4 +-
native/spark-expr/src/jvm_udf/mod.rs | 3 +-
native/spark-expr/src/string_funcs/split.rs | 26 +++++++----
20 files changed, 129 insertions(+), 93 deletions(-)
diff --git a/native/common/src/bin/analyze_trace.rs b/native/common/src/bin/analyze_trace.rs
index b2cdb9eba9..644fd1d474 100644
--- a/native/common/src/bin/analyze_trace.rs
+++ b/native/common/src/bin/analyze_trace.rs
@@ -121,7 +121,7 @@ fn main() {
}
let file = File::open(&args[1]).expect("Failed to open trace file");
- let reader = BufReader::new(file);
+ let mut reader = BufReader::new(file);
// Index into ALLOCATED_COUNTERS of the counter being analyzed, once one has been seen
let mut source: Option<usize> = None;
@@ -147,8 +147,12 @@ fn main() {
// Each line is one JSON event, possibly with a trailing comma.
// The file starts with "[ " on the first event line or as a prefix.
- for line in reader.lines() {
- let line = line.expect("Failed to read line");
+ let mut line = String::new();
+ loop {
+ line.clear();
+ if reader.read_line(&mut line).expect("Failed to read line") == 0 {
+ break;
+ }
let trimmed = line.trim();
// Skip empty lines or bare array brackets
diff --git a/native/core/src/execution/operators/iceberg_dictionary.rs b/native/core/src/execution/operators/iceberg_dictionary.rs
index 0d6dc8165e..1540e3fd74 100644
--- a/native/core/src/execution/operators/iceberg_dictionary.rs
+++ b/native/core/src/execution/operators/iceberg_dictionary.rs
@@ -122,7 +122,7 @@ impl DictionaryChooser {
.with_coerce_types(base.coerce_types())
.convert(schema)
.map_err(DataFusionError::from)?;
- let mut leaves = Vec::new();
+ let mut leaves = Vec::with_capacity(parquet_schema.num_columns());
for (column, field) in schema.fields().iter().enumerate() {
let mut leaf = Leaf {
column,
diff --git a/native/core/src/execution/operators/merge_rows.rs b/native/core/src/execution/operators/merge_rows.rs
index 9166127168..869f6141b9 100644
--- a/native/core/src/execution/operators/merge_rows.rs
+++ b/native/core/src/execution/operators/merge_rows.rs
@@ -16,6 +16,7 @@
// under the License.
use arrow::array::{Array, ArrayRef, BooleanArray, Int64Array, RecordBatch};
+use arrow::buffer::BooleanBuffer;
use arrow::compute::kernels::boolean::{and, and_not, not};
use arrow::compute::{filter_record_batch, prep_null_mask_filter};
use arrow::datatypes::{DataType, SchemaRef};
@@ -503,7 +504,7 @@ fn run_group(
// predicates that Spark would never evaluate after an earlier clause matched.
let fire = match instr.condition.evaluate(¤t)? {
ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))) => {
- BooleanArray::from(vec![true; current.num_rows()])
+ BooleanArray::new(BooleanBuffer::new_set(current.num_rows()), None)
}
ColumnarValue::Scalar(ScalarValue::Boolean(Some(false) | None)) => continue,
value => value
diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs
index cc2da616e5..81fed5e599 100644
--- a/native/core/src/execution/operators/parquet_writer.rs
+++ b/native/core/src/execution/operators/parquet_writer.rs
@@ -547,9 +547,10 @@ impl ExecutionPlan for ParquetWriterExec {
// the placeholder Scan. Both schemas use the same Catalyst data types, and
// disabling field-name matching still recursively validates nested nullability;
// only nested field names and metadata are ignored.
+ let (_, columns, _) = batch.into_parts();
RecordBatch::try_new_with_options(
Arc::clone(&schema_for_write),
- batch.columns().to_vec(),
+ columns,
&RecordBatchOptions::new().with_match_field_names(false),
)
.map_err(|e| {
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index cc9157f846..ae5916092b 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -504,6 +504,8 @@ impl PhysicalPlanner {
object_store_options: &HashMap<String, String>,
) -> Result<Vec<PartitionedFile>, ExecutionError> {
let mut files = Vec::with_capacity(partition.partitioned_file.len());
+ // Create an empty input schema for partition values because they are all literals.
+ let empty_schema = Arc::new(Schema::empty());
partition.partitioned_file.iter().try_for_each(|file| {
assert!(file.start + file.length <= file.file_size);
@@ -523,8 +525,6 @@ impl PhysicalPlanner {
partitioned_file.object_meta.location = path;
// Process partition values
- // Create an empty input schema for partition values because they are all literals.
- let empty_schema = Arc::new(Schema::empty());
let partition_values: Result<Vec<_>, _> = file
.partition_values
.iter()
diff --git a/native/core/src/execution/spark_plan.rs b/native/core/src/execution/spark_plan.rs
index 2487c76a14..7f5c645a19 100644
--- a/native/core/src/execution/spark_plan.rs
+++ b/native/core/src/execution/spark_plan.rs
@@ -57,15 +57,11 @@ impl SparkPlan {
children: Vec<Arc<SparkPlan>>,
additional_native_plans: Vec<Arc<dyn ExecutionPlan>>,
) -> Self {
- let mut accum: Vec<Arc<dyn ExecutionPlan>> = vec![];
- for plan in &additional_native_plans {
- accum.push(Arc::clone(plan));
- }
Self {
plan_id,
native_plan,
children,
- additional_native_plans: accum,
+ additional_native_plans,
}
}
diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs
index cc17429a72..2a2c173747 100644
--- a/native/core/src/parquet/parquet_support.rs
+++ b/native/core/src/parquet/parquet_support.rs
@@ -443,7 +443,7 @@ pub(crate) fn match_struct_fields(
// `None` marks an id that more than one file field carries.
let from_id_to_index: HashMap<i32, Option<usize>> = if should_match_by_id {
- let mut map = HashMap::new();
+ let mut map = HashMap::with_capacity(from_fields.len());
for (i, field) in from_fields.iter().enumerate() {
if let Some(id) = field_id(field) {
map.entry(id).and_modify(|m| *m = None).or_insert(Some(i));
@@ -465,7 +465,8 @@ pub(crate) fn match_struct_fields(
// Group file field indices by folded name so a case-insensitive collision is detected
// (Spark's `caseInsensitiveParquetFieldMap`) rather than silently overwritten.
- let mut folded_to_indices: HashMap<&str, Vec<usize>> = HashMap::new();
+ let mut folded_to_indices: HashMap<&str, Vec<usize>> =
+ HashMap::with_capacity(from_folded.len());
for (i, folded) in from_folded.iter().enumerate() {
folded_to_indices
.entry(folded.as_str())
diff --git a/native/core/src/parquet/schema_adapter.rs b/native/core/src/parquet/schema_adapter.rs
index 495eeb545e..fb37111ec5 100644
--- a/native/core/src/parquet/schema_adapter.rs
+++ b/native/core/src/parquet/schema_adapter.rs
@@ -229,7 +229,7 @@ fn remap_physical_schema(
// Pre-build id -> first matching logical field for the per-physical rename pass below.
let id_to_logical: HashMap<i32, &FieldRef> = if should_match_by_id {
- let mut map = HashMap::new();
+ let mut map = HashMap::with_capacity(logical_schema.fields().len());
for lf in logical_schema.fields() {
if let Some(id) = field_id(lf) {
map.entry(id).or_insert(lf);
@@ -944,7 +944,7 @@ impl PhysicalExprAdapterFactory for SparkPhysicalExprAdapterFactory {
let original_folded = original_folded.as_ref().unwrap_or(&physical_folded);
// Only allocate per-column index vectors when a folded name actually collides.
- let mut seen = HashSet::new();
+ let mut seen = HashSet::with_capacity(original_folded.len());
let mut collisions = HashSet::new();
for name in original_folded {
if !seen.insert(name.as_str()) {
@@ -982,7 +982,7 @@ impl PhysicalExprAdapterFactory for SparkPhysicalExprAdapterFactory {
};
let id_duplicate_roots = if should_match_by_id {
- let mut exact_names = HashSet::new();
+ let mut exact_names = HashSet::with_capacity(physical_file_schema.fields().len());
let mut duplicate_names = HashSet::new();
for field in physical_file_schema.fields() {
if !exact_names.insert(field.name()) {
diff --git a/native/operators/src/expand.rs b/native/operators/src/expand.rs
index 22bdc6d219..13d0851897 100644
--- a/native/operators/src/expand.rs
+++ b/native/operators/src/expand.rs
@@ -220,7 +220,7 @@ impl ExpandStream {
batch: &RecordBatch,
projection: &[Arc<dyn PhysicalExpr>],
) -> Result<RecordBatch, DataFusionError> {
- let mut columns = vec![];
+ let mut columns = Vec::with_capacity(projection.len());
projection.iter().try_for_each(|expr| {
let column = expr.evaluate(batch)?;
diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs
index 5af1e052d7..6aae0e2d34 100644
--- a/native/operators/src/explode.rs
+++ b/native/operators/src/explode.rs
@@ -567,25 +567,23 @@ fn flatten_struct_cols(
struct_column_indices: &HashSet<usize>,
) -> Result<RecordBatch> {
// horizontal expansion because of struct unnest
- let columns_expanded = input_batch
- .iter()
- .enumerate()
- .map(|(idx, column_data)| match struct_column_indices.get(&idx) {
+ let mut columns_expanded = Vec::with_capacity(schema.fields().len());
+ for (idx, column_data) in input_batch.iter().enumerate() {
+ match struct_column_indices.get(&idx) {
Some(_) => match column_data.data_type() {
DataType::Struct(_) => {
let struct_arr = column_data.as_any().downcast_ref::<StructArray>().unwrap();
- Ok(struct_arr.columns().to_vec())
+ columns_expanded.extend_from_slice(struct_arr.columns());
+ }
+ data_type => {
+ return internal_err!(
+ "expecting column {idx} from input plan to be a struct, got {data_type}"
+ )
}
- data_type => internal_err!(
- "expecting column {idx} from input plan to be a struct, got {data_type}"
- ),
},
- None => Ok(vec![Arc::clone(column_data)]),
- })
- .collect::<Result<Vec<_>>>()?
- .into_iter()
- .flatten()
- .collect();
+ None => columns_expanded.push(Arc::clone(column_data)),
+ }
+ }
Ok(RecordBatch::try_new(Arc::clone(schema), columns_expanded)?)
}
@@ -791,7 +789,7 @@ fn build_batch(
let transformed = match list_type_columns.len() {
0 => flatten_struct_cols(batch.columns(), schema, struct_column_indices),
_ => {
- let mut temp_unnested_result = HashMap::new();
+ let mut temp_unnested_result = HashMap::with_capacity(list_type_columns.len());
let max_recursion = list_type_columns
.iter()
.fold(0, |highest_depth, ListUnnest { depth, .. }| {
@@ -831,7 +829,7 @@ fn build_batch(
}
let unnested_array_map: HashMap<usize, Vec<UnnestingResult>> =
temp_unnested_result.into_iter().fold(
- HashMap::new(),
+ HashMap::with_capacity(list_type_columns.len()),
|mut acc,
(
ListUnnest {
@@ -894,19 +892,16 @@ fn build_batch(
)
.collect::<HashMap<_, _>>();
- let ret = flatten_arrs
- .into_iter()
- .enumerate()
- .flat_map(|(col_idx, arr)| {
- // Convert original column into its unnested version(s)
- // Plural because one column can be unnested with different recursion level
- // and into separate output columns
- match multi_unnested_per_original_index.remove(&col_idx) {
- Some(unnested_arrays) => unnested_arrays,
- None => vec![arr],
- }
- })
- .collect::<Vec<_>>();
+ let mut ret = Vec::with_capacity(flatten_arrs.len() + list_type_columns.len());
+ for (col_idx, arr) in flatten_arrs.into_iter().enumerate() {
+ // Convert original column into its unnested version(s)
+ // Plural because one column can be unnested with different recursion level
+ // and into separate output columns
+ match multi_unnested_per_original_index.remove(&col_idx) {
+ Some(unnested_arrays) => ret.extend(unnested_arrays),
+ None => ret.push(arr),
+ }
+ }
flatten_struct_cols(&ret, schema, struct_column_indices)
}
diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs
index 05b14b5324..f217165d38 100644
--- a/native/shuffle/src/partitioners/multi_partition.rs
+++ b/native/shuffle/src/partitioners/multi_partition.rs
@@ -23,7 +23,7 @@ use crate::partitioners::partitioned_batch_iterator::{
use crate::partitioners::ShufflePartitioner;
use crate::writers::PartitionWriter;
use crate::{comet_partitioning, CometPartitioning, RoundRobinStrategy};
-use arrow::array::{Array, ArrayData, ArrayRef, RecordBatch};
+use arrow::array::{Array, ArrayData, RecordBatch};
use datafusion::common::utils::proxy::VecAllocExt;
use datafusion::common::{DataFusionError, HashSet};
use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation};
@@ -424,16 +424,14 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
} else {
(*max_hash_columns).min(input.num_columns())
};
- let columns_to_hash: Vec<ArrayRef> = (0..num_columns_to_hash)
- .map(|i| Arc::clone(input.column(i)))
- .collect();
+ let columns_to_hash = &input.columns()[..num_columns_to_hash];
// Use identical seed as Spark hash partitioning.
let hashes_buf = &mut scratch.hashes_buf[..num_rows];
hashes_buf.fill(42_u32);
// Compute hash for selected columns
- create_murmur3_hashes(&columns_to_hash, hashes_buf)?;
+ create_murmur3_hashes(columns_to_hash, hashes_buf)?;
// Assign partition IDs based on hash (same as hash partitioning)
let partition_ids = &mut scratch.partition_ids[..num_rows];
@@ -736,7 +734,7 @@ impl<T: PartitionWriter> Debug for MultiPartitionShuffleRepartitioner<T> {
#[cfg(test)]
mod tests {
use super::*;
- use arrow::array::Int64Array;
+ use arrow::array::{ArrayRef, Int64Array};
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
#[derive(Default)]
diff --git a/native/shuffle/src/schema_align.rs b/native/shuffle/src/schema_align.rs
index e07e92e582..d9089fd1f1 100644
--- a/native/shuffle/src/schema_align.rs
+++ b/native/shuffle/src/schema_align.rs
@@ -232,11 +232,11 @@ struct SchemaAlignStream {
impl SchemaAlignStream {
fn align(&self, batch: RecordBatch) -> Result<RecordBatch, DataFusionError> {
- let num_rows = batch.num_rows();
+ let (_, columns, num_rows) = batch.into_parts();
cast_and_stamp_schema(
"CometSchemaAlignExec",
&self.target_schema,
- batch.columns().to_vec(),
+ columns,
num_rows,
)
}
diff --git a/native/shuffle/src/spark_unsafe/row.rs b/native/shuffle/src/spark_unsafe/row.rs
index ce18a398df..9246665f45 100644
--- a/native/shuffle/src/spark_unsafe/row.rs
+++ b/native/shuffle/src/spark_unsafe/row.rs
@@ -37,7 +37,7 @@ use arrow::array::{
Array, ArrayRef, RecordBatch, RecordBatchOptions,
};
use arrow::compute::cast;
-use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
use arrow::error::ArrowError;
use datafusion::physical_plan::metrics::Time;
use datafusion_comet_jni_bridge::errors::CometError;
@@ -380,6 +380,10 @@ fn append_nested_struct_fields_field_major(
) -> Result<(), CometError> {
let num_rows = row_addresses.len();
let mut row = SparkUnsafeRow::new_with_num_fields(fields.len());
+ // Reused by every nested struct field. The recursive call only borrows them.
+ let mut nested_addresses: Vec<jlong> = Vec::new();
+ let mut nested_sizes: Vec<jint> = Vec::new();
+ let mut nested_is_null: Vec<bool> = Vec::new();
// Helper macro for processing primitive fields
macro_rules! process_field {
@@ -517,9 +521,12 @@ fn append_nested_struct_fields_field_major(
let nested_builder = get_field_builder!(struct_builder, StructBuilder, field_idx);
// Collect nested struct addresses and sizes in one pass, building validity
- let mut nested_addresses: Vec<jlong> = Vec::with_capacity(num_rows);
- let mut nested_sizes: Vec<jint> = Vec::with_capacity(num_rows);
- let mut nested_is_null: Vec<bool> = Vec::with_capacity(num_rows);
+ nested_addresses.clear();
+ nested_sizes.clear();
+ nested_is_null.clear();
+ nested_addresses.reserve(num_rows);
+ nested_sizes.reserve(num_rows);
+ nested_is_null.reserve(num_rows);
for row_idx in 0..num_rows {
if struct_is_null[row_idx] {
@@ -837,6 +844,11 @@ fn append_struct_fields_field_major(
}
}
+ // Reused by every nested struct field. The recursive call only borrows them.
+ let mut nested_addresses: Vec<jlong> = Vec::new();
+ let mut nested_sizes: Vec<jint> = Vec::new();
+ let mut nested_is_null: Vec<bool> = Vec::new();
+
// Helper macro for processing primitive fields
macro_rules! process_field {
($builder_type:ty, $field_idx:expr, $get_value:expr) => {{
@@ -970,9 +982,12 @@ fn append_struct_fields_field_major(
let nested_builder = get_field_builder!(struct_builder, StructBuilder, field_idx);
// Collect nested struct addresses and sizes in one pass, building validity
- let mut nested_addresses: Vec<jlong> = Vec::with_capacity(num_rows);
- let mut nested_sizes: Vec<jint> = Vec::with_capacity(num_rows);
- let mut nested_is_null: Vec<bool> = Vec::with_capacity(num_rows);
+ nested_addresses.clear();
+ nested_sizes.clear();
+ nested_is_null.clear();
+ nested_addresses.reserve(num_rows);
+ nested_sizes.reserve(num_rows);
+ nested_is_null.reserve(num_rows);
for (row_idx, i) in (row_start..row_end).enumerate() {
if struct_is_null[row_idx] {
@@ -1374,7 +1389,7 @@ pub fn process_sorted_row_partition(
// Create builders once and reuse them across batches.
// After finish() is called, builders are reset and can be reused.
- let mut data_builders: Vec<Box<dyn ArrayBuilder>> = vec![];
+ let mut data_builders: Vec<Box<dyn ArrayBuilder>> = Vec::with_capacity(schema.len());
schema.iter().try_for_each(|dt| {
make_builders(dt, batch_size, prefer_dictionary_ratio)
.map(|builder| data_builders.push(builder))?;
@@ -1395,6 +1410,8 @@ pub fn process_sorted_row_partition(
// One context for every batch this call encodes; the JVM calls in once per sorted
// partition, so there is no wider native scope to hoist it to.
let mut codec_context = ShuffleCodecContext::default();
+ // The block writer pre-encodes the batch schema, so it is kept until the schema changes.
+ let mut block_writer: Option<(SchemaRef, ShuffleBlockWriter)> = None;
while current_row < row_num {
let n = std::cmp::min(batch_size, row_num - current_row);
@@ -1427,8 +1444,17 @@ pub fn process_sorted_row_partition(
frozen.clear();
let mut cursor = Cursor::new(&mut frozen);
- let block_writer = ShuffleBlockWriter::try_new(batch.schema().as_ref(), codec.clone())?;
- written += block_writer.write_batch(&batch, &mut cursor, &mut codec_context, &ipc_time)?;
+ // The schema changes between batches when a dictionary-encoded column falls back to a
+ // plain array (see `builder_to_array`).
+ if block_writer
+ .as_ref()
+ .is_none_or(|(cached, _)| cached != batch.schema_ref())
+ {
+ let writer = ShuffleBlockWriter::try_new(batch.schema_ref(), codec.clone())?;
+ block_writer = Some((batch.schema(), writer));
+ }
+ let (_, writer) = block_writer.as_ref().expect("block writer is set above");
+ written += writer.write_batch(&batch, &mut cursor, &mut codec_context, &ipc_time)?;
if let Some(checksum) = &mut current_checksum {
checksum.update(&mut cursor)?;
diff --git a/native/spark-expr/src/agg_funcs/percentile.rs b/native/spark-expr/src/agg_funcs/percentile.rs
index 87219db253..7fd66a4a69 100644
--- a/native/spark-expr/src/agg_funcs/percentile.rs
+++ b/native/spark-expr/src/agg_funcs/percentile.rs
@@ -247,7 +247,11 @@ impl GroupsAccumulator for SparkPercentileGroupsAccumulator {
offsets.push(len);
}
- let values = emit_group_values.into_iter().flatten().collect::<Vec<_>>();
+ // `len` is the total number of values, which `flatten().collect()` has no size hint for.
+ let mut values = Vec::with_capacity(len as usize);
+ for group in emit_group_values {
+ values.extend_from_slice(&group);
+ }
let values = Float64Array::new(ScalarBuffer::from(values), None);
let list = ListArray::new(
Arc::new(Field::new_list_field(DataType::Float64, true)),
diff --git a/native/spark-expr/src/bloom_filter/bloom_filter_agg.rs b/native/spark-expr/src/bloom_filter/bloom_filter_agg.rs
index 8920b30c9d..adce71233a 100644
--- a/native/spark-expr/src/bloom_filter/bloom_filter_agg.rs
+++ b/native/spark-expr/src/bloom_filter/bloom_filter_agg.rs
@@ -24,6 +24,7 @@ use crate::bloom_filter::spark_bloom_filter::{SparkBloomFilter, SparkBloomFilter
use arrow::array::ArrayRef;
use arrow::array::BinaryArray;
+use arrow::array::StringArray;
use datafusion::common::{downcast_value, ScalarValue};
use datafusion::error::{DataFusionError, Result};
use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
@@ -118,6 +119,13 @@ impl Accumulator for SparkBloomFilter {
return Ok(());
}
let arr = &values[0];
+ // Hash strings in place rather than materializing a `ScalarValue::Utf8` per row.
+ if let Some(strings) = arr.as_any().downcast_ref::<StringArray>() {
+ for value in strings.iter().flatten() {
+ self.put_binary(value.as_bytes());
+ }
+ return Ok(());
+ }
(0..arr.len()).try_for_each(|index| {
let v = ScalarValue::try_from_array(arr, index)?;
diff --git a/native/spark-expr/src/bloom_filter/spark_bit_array.rs b/native/spark-expr/src/bloom_filter/spark_bit_array.rs
index b55567402b..214a832c06 100644
--- a/native/spark-expr/src/bloom_filter/spark_bit_array.rs
+++ b/native/spark-expr/src/bloom_filter/spark_bit_array.rs
@@ -62,8 +62,8 @@ impl SparkBitArray {
self.data.len()
}
- pub fn data(&self) -> Vec<u64> {
- self.data.clone()
+ pub fn data(&self) -> &[u64] {
+ &self.data
}
/// Number of set bits in the array. Mirrors Spark's `BitArray.cardinality()`.
diff --git a/native/spark-expr/src/bloom_filter/spark_bloom_filter.rs b/native/spark-expr/src/bloom_filter/spark_bloom_filter.rs
index 0c5dbea6c8..114b6be404 100644
--- a/native/spark-expr/src/bloom_filter/spark_bloom_filter.rs
+++ b/native/spark-expr/src/bloom_filter/spark_bloom_filter.rs
@@ -16,7 +16,6 @@
// under the License.
use arrow::array::{ArrowNativeTypeOp, BooleanArray, Int64Array};
-use arrow::datatypes::ToByteSlice;
use datafusion::common::{DataFusionError, Result as DFResult};
use std::cmp;
@@ -157,19 +156,19 @@ impl SparkBloomFilter {
/// Serializes a SparkBloomFilter to a byte array conforming to Spark's BloomFilter
/// binary format. The output format follows the filter's `version`.
pub fn spark_serialization(&self) -> Vec<u8> {
- let mut out: Vec<u8> = (self.version.to_int() as u32).to_be_bytes().to_vec();
- out.append(&mut self.num_hash_functions.to_be_bytes().to_vec());
+ // At most a 16-byte header (version, numHashFunctions, seed, numWords), then the words.
+ let mut out: Vec<u8> = Vec::with_capacity(16 + self.bits.word_size() * 8);
+ out.extend_from_slice(&(self.version.to_int() as u32).to_be_bytes());
+ out.extend_from_slice(&self.num_hash_functions.to_be_bytes());
if let SparkBloomFilterVersion::V2 = self.version {
// Spark's BloomFilterImplV2.writeTo writes the seed between
// numHashFunctions and the bit array.
- out.append(&mut (self.seed as u32).to_be_bytes().to_vec());
+ out.extend_from_slice(&(self.seed as u32).to_be_bytes());
}
- out.append(&mut (self.bits.word_size() as u32).to_be_bytes().to_vec());
- let mut filter_state: Vec<u64> = self.bits.data();
- for i in filter_state.iter_mut() {
- *i = i.to_be();
+ out.extend_from_slice(&(self.bits.word_size() as u32).to_be_bytes());
+ for word in self.bits.data() {
+ out.extend_from_slice(&word.to_be_bytes());
}
- out.append(&mut Vec::from(filter_state.to_byte_slice()));
out
}
diff --git a/native/spark-expr/src/json_funcs/json_array_length.rs b/native/spark-expr/src/json_funcs/json_array_length.rs
index 272b0b72d5..468fff426d 100644
--- a/native/spark-expr/src/json_funcs/json_array_length.rs
+++ b/native/spark-expr/src/json_funcs/json_array_length.rs
@@ -105,9 +105,7 @@ fn spark_json_array_length_scalar(scalar: &ScalarValue) -> Result<ScalarValue> {
}
fn spark_json_array_length_scalar_inner(json_str: &Option<String>) -> Result<ScalarValue> {
- let array_length = json_str
- .clone()
- .and_then(|json_str| get_json_array_length(&json_str));
+ let array_length = json_str.as_deref().and_then(get_json_array_length);
Ok(ScalarValue::Int32(array_length))
}
diff --git a/native/spark-expr/src/jvm_udf/mod.rs b/native/spark-expr/src/jvm_udf/mod.rs
index 6355b1a9da..e0de5fe369 100644
--- a/native/spark-expr/src/jvm_udf/mod.rs
+++ b/native/spark-expr/src/jvm_udf/mod.rs
@@ -208,7 +208,6 @@ impl PhysicalExpr for JvmScalarUdfExpr {
let out_arr_ptr = out_array.as_mut() as *mut FFI_ArrowArray as i64;
let out_sch_ptr = out_schema.as_mut() as *mut FFI_ArrowSchema as i64;
- let class_name = self.class_name.clone();
let n_args = arrays.len();
JVMClasses::with_env(|env| {
@@ -221,7 +220,7 @@ impl PhysicalExpr for JvmScalarUdfExpr {
})?;
let jclass_name = env
- .new_string(&class_name)
+ .new_string(&self.class_name)
.map_err(|e| CometError::JNI { source: e })?;
let in_arr_java = env
diff --git a/native/spark-expr/src/string_funcs/split.rs b/native/spark-expr/src/string_funcs/split.rs
index 64033440ca..685e0a4356 100644
--- a/native/spark-expr/src/string_funcs/split.rs
+++ b/native/spark-expr/src/string_funcs/split.rs
@@ -445,18 +445,24 @@ fn push_split_parts<O: OffsetSizeTrait>(
builder: &mut GenericStringBuilder<O>,
) {
if limit == 0 {
- // limit = 0: split all, drop trailing empties. Need to know the end
- // before pushing, so collect borrowed slices first (no string copies).
- let mut parts: Vec<&str> = regex.split(string).collect();
- while parts.last().is_some_and(|s| s.is_empty()) {
- parts.pop();
+ // limit = 0: split all, drop trailing empties. Empty parts are held back until a
+ // non-empty part follows, so the trailing ones are never appended.
+ let mut pending_empty = 0;
+ let mut appended = false;
+ for part in regex.split(string) {
+ if part.is_empty() {
+ pending_empty += 1;
+ } else {
+ for _ in 0..pending_empty {
+ builder.append_value("");
+ }
+ pending_empty = 0;
+ builder.append_value(part);
+ appended = true;
+ }
}
- if parts.is_empty() {
+ if !appended {
builder.append_value("");
- } else {
- for p in parts {
- builder.append_value(p);
- }
}
} else if limit > 0 {
// limit > 0: at most limit-1 splits.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org