You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by ks...@apache.org on 2020/10/19 18:01:00 UTC
[arrow] 05/30: ARROW-10295 [Rust] [DataFusion] Replace
Rc> by Box<> in accumulators.
This is an automated email from the ASF dual-hosted git repository.
kszucs pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
commit a030fc5d1bc5b1d6dac18b4d4e68492d00b94edf
Author: Jorge C. Leitao <jo...@gmail.com>
AuthorDate: Tue Oct 13 08:54:21 2020 -0600
ARROW-10295 [Rust] [DataFusion] Replace Rc<RefCell<>> by Box<> in accumulators.
This PR replaces `Rc<RefCell<>>` by `Box<>`. We do not need interior mutability on the accumulations.
Closes #8456 from jorgecarleitao/box
Authored-by: Jorge C. Leitao <jo...@gmail.com>
Signed-off-by: Andy Grove <an...@nvidia.com>
---
rust/datafusion/examples/simple_udaf.rs | 4 +--
rust/datafusion/src/execution/context.rs | 8 ++---
rust/datafusion/src/physical_plan/aggregates.rs | 4 +--
.../src/physical_plan/distinct_expressions.rs | 17 ++++-------
rust/datafusion/src/physical_plan/expressions.rs | 34 +++++++++-------------
.../datafusion/src/physical_plan/hash_aggregate.rs | 34 ++++++++--------------
rust/datafusion/src/physical_plan/mod.rs | 4 +--
rust/datafusion/src/physical_plan/udaf.rs | 4 +--
8 files changed, 40 insertions(+), 69 deletions(-)
diff --git a/rust/datafusion/examples/simple_udaf.rs b/rust/datafusion/examples/simple_udaf.rs
index 4d3cc23..1f41f0d 100644
--- a/rust/datafusion/examples/simple_udaf.rs
+++ b/rust/datafusion/examples/simple_udaf.rs
@@ -24,7 +24,7 @@ use arrow::{
use datafusion::{error::Result, logical_plan::create_udaf, physical_plan::Accumulator};
use datafusion::{prelude::*, scalar::ScalarValue};
-use std::{cell::RefCell, rc::Rc, sync::Arc};
+use std::sync::Arc;
// create local execution context with an in-memory table
fn create_context() -> Result<ExecutionContext> {
@@ -138,7 +138,7 @@ async fn main() -> Result<()> {
// the return type; DataFusion expects this to match the type returned by `evaluate`.
Arc::new(DataType::Float64),
// This is the accumulator factory; DataFusion uses it to create new accumulators.
- Arc::new(|| Ok(Rc::new(RefCell::new(GeometricMean::new())))),
+ Arc::new(|| Ok(Box::new(GeometricMean::new()))),
// This is the description of the state. `state()` must match the types here.
Arc::new(vec![DataType::Float64, DataType::UInt32]),
);
diff --git a/rust/datafusion/src/execution/context.rs b/rust/datafusion/src/execution/context.rs
index eabc779..8df18c2 100644
--- a/rust/datafusion/src/execution/context.rs
+++ b/rust/datafusion/src/execution/context.rs
@@ -537,8 +537,8 @@ mod tests {
ArrayRef, Float64Array, Int32Array, PrimitiveArrayOps, StringArray,
};
use arrow::compute::add;
+ use std::fs::File;
use std::thread::{self, JoinHandle};
- use std::{cell::RefCell, fs::File, rc::Rc};
use std::{io::prelude::*, sync::Mutex};
use tempfile::TempDir;
use test::*;
@@ -1371,11 +1371,7 @@ mod tests {
"MY_AVG",
DataType::Float64,
Arc::new(DataType::Float64),
- Arc::new(|| {
- Ok(Rc::new(RefCell::new(AvgAccumulator::try_new(
- &DataType::Float64,
- )?)))
- }),
+ Arc::new(|| Ok(Box::new(AvgAccumulator::try_new(&DataType::Float64)?))),
Arc::new(vec![DataType::UInt64, DataType::Float64]),
);
diff --git a/rust/datafusion/src/physical_plan/aggregates.rs b/rust/datafusion/src/physical_plan/aggregates.rs
index 40bb562..d417c41 100644
--- a/rust/datafusion/src/physical_plan/aggregates.rs
+++ b/rust/datafusion/src/physical_plan/aggregates.rs
@@ -36,11 +36,11 @@ use crate::physical_plan::distinct_expressions;
use crate::physical_plan::expressions;
use arrow::datatypes::{DataType, Schema};
use expressions::{avg_return_type, sum_return_type};
-use std::{cell::RefCell, fmt, rc::Rc, str::FromStr, sync::Arc};
+use std::{fmt, str::FromStr, sync::Arc};
/// the implementation of an aggregate function
pub type AccumulatorFunctionImplementation =
- Arc<dyn Fn() -> Result<Rc<RefCell<dyn Accumulator>>> + Send + Sync>;
+ Arc<dyn Fn() -> Result<Box<dyn Accumulator>> + Send + Sync>;
/// This signature corresponds to which types an aggregator serializes
/// its state, given its return datatype.
diff --git a/rust/datafusion/src/physical_plan/distinct_expressions.rs b/rust/datafusion/src/physical_plan/distinct_expressions.rs
index 2d2ab62..cc77107 100644
--- a/rust/datafusion/src/physical_plan/distinct_expressions.rs
+++ b/rust/datafusion/src/physical_plan/distinct_expressions.rs
@@ -17,11 +17,9 @@
//! Implementations for DISTINCT expressions, e.g. `COUNT(DISTINCT c)`
-use std::cell::RefCell;
use std::convert::TryFrom;
use std::fmt::Debug;
use std::hash::Hash;
-use std::rc::Rc;
use std::sync::Arc;
use arrow::datatypes::{DataType, Field};
@@ -93,12 +91,12 @@ impl AggregateExpr for DistinctCount {
self.exprs.clone()
}
- fn create_accumulator(&self) -> Result<Rc<RefCell<dyn Accumulator>>> {
- Ok(Rc::new(RefCell::new(DistinctCountAccumulator {
+ fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
+ Ok(Box::new(DistinctCountAccumulator {
values: FnvHashSet::default(),
data_types: self.input_data_types.clone(),
count_data_type: self.data_type.clone(),
- })))
+ }))
}
}
@@ -282,8 +280,7 @@ mod tests {
DataType::UInt64,
);
- let accum = agg.create_accumulator()?;
- let mut accum = accum.borrow_mut();
+ let mut accum = agg.create_accumulator()?;
accum.update_batch(arrays)?;
Ok((accum.state()?, accum.evaluate()?))
@@ -300,8 +297,7 @@ mod tests {
DataType::UInt64,
);
- let accum = agg.create_accumulator()?;
- let mut accum = accum.borrow_mut();
+ let mut accum = agg.create_accumulator()?;
for row in rows.iter() {
accum.update(row)?
@@ -324,8 +320,7 @@ mod tests {
DataType::UInt64,
);
- let accum = agg.create_accumulator()?;
- let mut accum = accum.borrow_mut();
+ let mut accum = agg.create_accumulator()?;
accum.merge_batch(arrays)?;
Ok((accum.state()?, accum.evaluate()?))
diff --git a/rust/datafusion/src/physical_plan/expressions.rs b/rust/datafusion/src/physical_plan/expressions.rs
index 4c9029e..1f5dafd 100644
--- a/rust/datafusion/src/physical_plan/expressions.rs
+++ b/rust/datafusion/src/physical_plan/expressions.rs
@@ -17,10 +17,9 @@
//! Defines physical expressions that can evaluated at runtime during query execution
+use std::convert::TryFrom;
use std::fmt;
-use std::rc::Rc;
use std::sync::Arc;
-use std::{cell::RefCell, convert::TryFrom};
use crate::error::{ExecutionError, Result};
use crate::logical_plan::Operator;
@@ -162,10 +161,8 @@ impl AggregateExpr for Sum {
vec![self.expr.clone()]
}
- fn create_accumulator(&self) -> Result<Rc<RefCell<dyn Accumulator>>> {
- Ok(Rc::new(RefCell::new(SumAccumulator::try_new(
- &self.data_type,
- )?)))
+ fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
+ Ok(Box::new(SumAccumulator::try_new(&self.data_type)?))
}
}
@@ -391,11 +388,11 @@ impl AggregateExpr for Avg {
])
}
- fn create_accumulator(&self) -> Result<Rc<RefCell<dyn Accumulator>>> {
- Ok(Rc::new(RefCell::new(AvgAccumulator::try_new(
+ fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
+ Ok(Box::new(AvgAccumulator::try_new(
// avg is f64
&DataType::Float64,
- )?)))
+ )?))
}
fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> {
@@ -521,10 +518,8 @@ impl AggregateExpr for Max {
vec![self.expr.clone()]
}
- fn create_accumulator(&self) -> Result<Rc<RefCell<dyn Accumulator>>> {
- Ok(Rc::new(RefCell::new(MaxAccumulator::try_new(
- &self.data_type,
- )?)))
+ fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
+ Ok(Box::new(MaxAccumulator::try_new(&self.data_type)?))
}
}
@@ -774,10 +769,8 @@ impl AggregateExpr for Min {
vec![self.expr.clone()]
}
- fn create_accumulator(&self) -> Result<Rc<RefCell<dyn Accumulator>>> {
- Ok(Rc::new(RefCell::new(MinAccumulator::try_new(
- &self.data_type,
- )?)))
+ fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
+ Ok(Box::new(MinAccumulator::try_new(&self.data_type)?))
}
}
@@ -869,8 +862,8 @@ impl AggregateExpr for Count {
vec![self.expr.clone()]
}
- fn create_accumulator(&self) -> Result<Rc<RefCell<dyn Accumulator>>> {
- Ok(Rc::new(RefCell::new(CountAccumulator::new())))
+ fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
+ Ok(Box::new(CountAccumulator::new()))
}
}
@@ -2476,13 +2469,12 @@ mod tests {
batch: &RecordBatch,
agg: Arc<dyn AggregateExpr>,
) -> Result<ScalarValue> {
- let accum = agg.create_accumulator()?;
+ let mut accum = agg.create_accumulator()?;
let expr = agg.expressions();
let values = expr
.iter()
.map(|e| e.evaluate(batch))
.collect::<Result<Vec<_>>>()?;
- let mut accum = accum.borrow_mut();
accum.update_batch(&values)?;
accum.evaluate()
}
diff --git a/rust/datafusion/src/physical_plan/hash_aggregate.rs b/rust/datafusion/src/physical_plan/hash_aggregate.rs
index 5f4fe98..53b74c2 100644
--- a/rust/datafusion/src/physical_plan/hash_aggregate.rs
+++ b/rust/datafusion/src/physical_plan/hash_aggregate.rs
@@ -18,8 +18,6 @@
//! Defines the execution plan for the hash aggregate operation
use std::any::Any;
-use std::cell::RefCell;
-use std::rc::Rc;
use std::sync::Arc;
use crate::error::{ExecutionError, Result};
@@ -278,9 +276,8 @@ fn group_aggregate_batch(
.map(|(_, (accumulator_set, indices))| {
// 2.2
accumulator_set
- .iter()
- .zip(&aggr_input_values)
.into_iter()
+ .zip(&aggr_input_values)
.map(|(accumulator, aggr_array)| {
(
accumulator,
@@ -300,12 +297,10 @@ fn group_aggregate_batch(
})
// 2.4
.map(|(accumulator, values)| match mode {
- AggregateMode::Partial => {
- accumulator.borrow_mut().update_batch(&values)
- }
+ AggregateMode::Partial => accumulator.update_batch(&values),
AggregateMode::Final => {
// note: the aggregation here is over states, not values, thus the merge
- accumulator.borrow_mut().merge_batch(&values)
+ accumulator.merge_batch(&values)
}
})
.collect::<Result<()>>()
@@ -335,7 +330,7 @@ impl GroupedHashAggregateIterator {
}
}
-type AccumulatorSet = Vec<Rc<RefCell<dyn Accumulator>>>;
+type AccumulatorSet = Vec<Box<dyn Accumulator>>;
impl Iterator for GroupedHashAggregateIterator {
type Item = ArrowResult<RecordBatch>;
@@ -490,7 +485,7 @@ impl HashAggregateIterator {
fn aggregate_batch(
mode: &AggregateMode,
batch: &RecordBatch,
- accumulators: &AccumulatorSet,
+ accumulators: &mut AccumulatorSet,
expressions: &Vec<Vec<Arc<dyn PhysicalExpr>>>,
) -> Result<()> {
// 1.1 iterate accumulators and respective expressions together
@@ -499,7 +494,7 @@ fn aggregate_batch(
// 1.1
accumulators
- .iter()
+ .into_iter()
.zip(expressions)
.map(|(accum, expr)| {
// 1.2
@@ -510,8 +505,8 @@ fn aggregate_batch(
// 1.3
match mode {
- AggregateMode::Partial => accum.borrow_mut().update_batch(values),
- AggregateMode::Final => accum.borrow_mut().merge_batch(values),
+ AggregateMode::Partial => accum.update_batch(values),
+ AggregateMode::Final => accum.merge_batch(values),
}
})
.collect::<Result<()>>()
@@ -528,7 +523,7 @@ impl Iterator for HashAggregateIterator {
// return single batch
self.finished = true;
- let accumulators = match create_accumulators(&self.aggr_expr) {
+ let mut accumulators = match create_accumulators(&self.aggr_expr) {
Ok(e) => e,
Err(e) => return Some(Err(ExecutionError::into_arrow_external_error(e))),
};
@@ -547,7 +542,7 @@ impl Iterator for HashAggregateIterator {
.as_mut()
.into_iter()
.map(|batch| {
- aggregate_batch(&mode, &batch?, &accumulators, &expressions)
+ aggregate_batch(&mode, &batch?, &mut accumulators, &expressions)
.map_err(ExecutionError::into_arrow_external_error)
})
.collect::<ArrowResult<()>>()
@@ -655,7 +650,7 @@ fn finalize_aggregation(
// build the vector of states
let a = accumulators
.iter()
- .map(|accumulator| accumulator.borrow_mut().state())
+ .map(|accumulator| accumulator.state())
.map(|value| {
value.and_then(|e| {
Ok(e.iter().map(|v| v.to_array()).collect::<Vec<ArrayRef>>())
@@ -668,12 +663,7 @@ fn finalize_aggregation(
// merge the state to the final value
accumulators
.iter()
- .map(|accumulator| {
- accumulator
- .borrow_mut()
- .evaluate()
- .and_then(|v| Ok(v.to_array()))
- })
+ .map(|accumulator| accumulator.evaluate().and_then(|v| Ok(v.to_array())))
.collect::<Result<Vec<ArrayRef>>>()
}
}
diff --git a/rust/datafusion/src/physical_plan/mod.rs b/rust/datafusion/src/physical_plan/mod.rs
index ac33c67..1d6c46a 100644
--- a/rust/datafusion/src/physical_plan/mod.rs
+++ b/rust/datafusion/src/physical_plan/mod.rs
@@ -18,9 +18,7 @@
//! Traits for physical query plan, supporting parallel execution for partitioned relations.
use std::any::Any;
-use std::cell::RefCell;
use std::fmt::{Debug, Display};
-use std::rc::Rc;
use std::sync::Arc;
use crate::execution::context::ExecutionContextState;
@@ -122,7 +120,7 @@ pub trait AggregateExpr: Send + Sync + Debug {
/// the accumulator used to accumulate values from the expressions.
/// the accumulator expects the same number of arguments as `expressions` and must
/// return states with the same description as `state_fields`
- fn create_accumulator(&self) -> Result<Rc<RefCell<dyn Accumulator>>>;
+ fn create_accumulator(&self) -> Result<Box<dyn Accumulator>>;
/// the fields that encapsulate the Accumulator's state
/// the number of fields here equals the number of states that the accumulator contains
diff --git a/rust/datafusion/src/physical_plan/udaf.rs b/rust/datafusion/src/physical_plan/udaf.rs
index 933fd23..db86e14 100644
--- a/rust/datafusion/src/physical_plan/udaf.rs
+++ b/rust/datafusion/src/physical_plan/udaf.rs
@@ -18,7 +18,7 @@
//! This module contains functions and structs supporting user-defined aggregate functions.
use fmt::{Debug, Formatter};
-use std::{cell::RefCell, fmt, rc::Rc};
+use std::fmt;
use arrow::{
datatypes::Field,
@@ -150,7 +150,7 @@ impl AggregateExpr for AggregateFunctionExpr {
Ok(Field::new(&self.name, self.data_type.clone(), true))
}
- fn create_accumulator(&self) -> Result<Rc<RefCell<dyn Accumulator>>> {
+ fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
(self.fun.accumulator)()
}
}