This is an automated email from the ASF dual-hosted git repository.
berkay pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 8adbc2324a [minor]: use arrow take_batch instead of
get_record_batch_indices (#13084)
8adbc2324a is described below
commit 8adbc2324afac66d8cb88b20cb1482913b190d4b
Author: Mustafa Akur <[email protected]>
AuthorDate: Thu Oct 24 00:10:09 2024 -0700
[minor]: use arrow take_batch instead of get_record_batch_indices (#13084)
* Initial commit
* Fix linter errors
* Minor changes
* Fix error
---
datafusion/common/src/utils/mod.rs | 24 ++++------------------
.../core/tests/fuzz_cases/equivalence/utils.rs | 17 +++++----------
.../src/windows/bounded_window_agg_exec.rs | 6 +++---
3 files changed, 12 insertions(+), 35 deletions(-)
diff --git a/datafusion/common/src/utils/mod.rs
b/datafusion/common/src/utils/mod.rs
index def1def985..dacf90af9b 100644
--- a/datafusion/common/src/utils/mod.rs
+++ b/datafusion/common/src/utils/mod.rs
@@ -23,16 +23,14 @@ pub mod proxy;
pub mod string_utils;
use crate::error::{_internal_datafusion_err, _internal_err};
-use crate::{arrow_datafusion_err, DataFusionError, Result, ScalarValue};
-use arrow::array::{ArrayRef, PrimitiveArray};
+use crate::{DataFusionError, Result, ScalarValue};
+use arrow::array::ArrayRef;
use arrow::buffer::OffsetBuffer;
-use arrow::compute::{partition, take_arrays, SortColumn, SortOptions};
-use arrow::datatypes::{Field, SchemaRef, UInt32Type};
-use arrow::record_batch::RecordBatch;
+use arrow::compute::{partition, SortColumn, SortOptions};
+use arrow::datatypes::{Field, SchemaRef};
use arrow_array::cast::AsArray;
use arrow_array::{
Array, FixedSizeListArray, LargeListArray, ListArray, OffsetSizeTrait,
- RecordBatchOptions,
};
use arrow_schema::DataType;
use sqlparser::ast::Ident;
@@ -92,20 +90,6 @@ pub fn get_row_at_idx(columns: &[ArrayRef], idx: usize) ->
Result<Vec<ScalarValu
.collect()
}
-/// Construct a new RecordBatch from the rows of the `record_batch` at the
`indices`.
-pub fn get_record_batch_at_indices(
- record_batch: &RecordBatch,
- indices: &PrimitiveArray<UInt32Type>,
-) -> Result<RecordBatch> {
- let new_columns = take_arrays(record_batch.columns(), indices, None)?;
- RecordBatch::try_new_with_options(
- record_batch.schema(),
- new_columns,
- &RecordBatchOptions::new().with_row_count(Some(indices.len())),
- )
- .map_err(|e| arrow_datafusion_err!(e))
-}
-
/// This function compares two tuples depending on the given sort options.
pub fn compare_rows(
x: &[ScalarValue],
diff --git a/datafusion/core/tests/fuzz_cases/equivalence/utils.rs
b/datafusion/core/tests/fuzz_cases/equivalence/utils.rs
index 61691311fe..acc45fe0e5 100644
--- a/datafusion/core/tests/fuzz_cases/equivalence/utils.rs
+++ b/datafusion/core/tests/fuzz_cases/equivalence/utils.rs
@@ -22,15 +22,11 @@ use std::any::Any;
use std::cmp::Ordering;
use std::sync::Arc;
-use arrow::compute::{lexsort_to_indices, SortColumn};
+use arrow::compute::{lexsort_to_indices, take_record_batch, SortColumn};
use arrow::datatypes::{DataType, Field, Schema};
-use arrow_array::{
- ArrayRef, Float32Array, Float64Array, PrimitiveArray, RecordBatch,
UInt32Array,
-};
+use arrow_array::{ArrayRef, Float32Array, Float64Array, RecordBatch,
UInt32Array};
use arrow_schema::{SchemaRef, SortOptions};
-use datafusion_common::utils::{
- compare_rows, get_record_batch_at_indices, get_row_at_idx,
-};
+use datafusion_common::utils::{compare_rows, get_row_at_idx};
use datafusion_common::{exec_err, plan_datafusion_err, DataFusionError,
Result};
use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
@@ -465,7 +461,7 @@ pub fn generate_table_for_orderings(
// Sort batch according to first ordering expression
let sort_columns = get_sort_columns(&batch, &orderings[0])?;
let sort_indices = lexsort_to_indices(&sort_columns, None)?;
- let mut batch = get_record_batch_at_indices(&batch, &sort_indices)?;
+ let mut batch = take_record_batch(&batch, &sort_indices)?;
// prune out rows that is invalid according to remaining orderings.
for ordering in orderings.iter().skip(1) {
@@ -490,10 +486,7 @@ pub fn generate_table_for_orderings(
}
}
// Only keep valid rows, that satisfies given ordering relation.
- batch = get_record_batch_at_indices(
- &batch,
- &PrimitiveArray::from_iter_values(keep_indices),
- )?;
+ batch = take_record_batch(&batch,
&UInt32Array::from_iter_values(keep_indices))?;
}
Ok(batch)
diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs
b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs
index 6254ae139a..6495657339 100644
--- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs
+++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs
@@ -40,6 +40,7 @@ use crate::{
SendableRecordBatchStream, Statistics, WindowExpr,
};
use ahash::RandomState;
+use arrow::compute::take_record_batch;
use arrow::{
array::{Array, ArrayRef, RecordBatchOptions, UInt32Builder},
compute::{concat, concat_batches, sort_to_indices, take_arrays},
@@ -49,8 +50,7 @@ use arrow::{
use datafusion_common::hash_utils::create_hashes;
use datafusion_common::stats::Precision;
use datafusion_common::utils::{
- evaluate_partition_ranges, get_at_indices, get_record_batch_at_indices,
- get_row_at_idx,
+ evaluate_partition_ranges, get_at_indices, get_row_at_idx,
};
use datafusion_common::{arrow_datafusion_err, exec_err, DataFusionError,
Result};
use datafusion_execution::TaskContext;
@@ -558,7 +558,7 @@ impl PartitionSearcher for LinearSearch {
let mut new_indices =
UInt32Builder::with_capacity(indices.len());
new_indices.append_slice(&indices);
let indices = new_indices.finish();
- Ok((row, get_record_batch_at_indices(record_batch, &indices)?))
+ Ok((row, take_record_batch(record_batch, &indices)?))
})
.collect()
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
