Skip to content

ARROW-9679: [Rust] [DataFusion] More efficient creation of final batch from HashAggregateExec #7936

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
andygrove wants to merge 8 commits into from
Closed

ARROW-9679: [Rust] [DataFusion] More efficient creation of final batch from HashAggregateExec #7936

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
e1d63ed
rough out new approach
andygrove Aug 12, 2020
432cdb4
Min and Max should always return same data type as input since can ne…
andygrove Aug 12, 2020
f7fb7d0
update tests and improve error handling
andygrove Aug 12, 2020
d71cda1
bug fix + formatting
andygrove Aug 12, 2020
c103c8c
rebase
andygrove Aug 12, 2020
ca175ac
remove panic, code clean up
andygrove Aug 12, 2020
ede721a
add comment
andygrove Aug 12, 2020
c9b4974
address more style feedback
andygrove Aug 12, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 40 additions & 66 deletions rust/datafusion/src/execution/physical_plan/expressions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -419,20 +419,7 @@ impl Max {

impl AggregateExpr for Max {
fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
match self.expr.data_type(input_schema)? {
Copy link
Contributor

@alamb alamb Aug 12, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree that the type of Max and Min should be the same as the input (as those functions are never going to overflow the value type of their input, the way Sum or Avg could.

I actually think the change results in cleaner code too, which is a nice bonus

DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => {
Ok(DataType::Int64)
}
DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => {
Ok(DataType::UInt64)
}
DataType::Float32 => Ok(DataType::Float32),
DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"MAX does not support {:?}",
other
))),
}
self.expr.data_type(input_schema)
}

fn evaluate_input(&self, batch: &RecordBatch) -> Result<ArrayRef> {
Expand All @@ -449,21 +436,21 @@ impl AggregateExpr for Max {
}

macro_rules! max_accumulate {
($SELF:ident, $VALUE:expr, $ARRAY_TYPE:ident, $SCALAR_VARIANT:ident, $TY:ty) => {{
($SELF:ident, $VALUE:expr, $ARRAY_TYPE:ident, $SCALAR_VARIANT:ident) => {{
$SELF.max = match $SELF.max {
Some(ScalarValue::$SCALAR_VARIANT(n)) => {
if n > ($VALUE as $TY) {
if n > ($VALUE) {
Some(ScalarValue::$SCALAR_VARIANT(n))
} else {
Some(ScalarValue::$SCALAR_VARIANT($VALUE as $TY))
Some(ScalarValue::$SCALAR_VARIANT($VALUE))
}
}
Some(_) => {
return Err(ExecutionError::InternalError(
"Unexpected ScalarValue variant".to_string(),
))
}
None => Some(ScalarValue::$SCALAR_VARIANT($VALUE as $TY)),
None => Some(ScalarValue::$SCALAR_VARIANT($VALUE)),
};
}};
}
Expand All @@ -477,34 +464,34 @@ impl Accumulator for MaxAccumulator {
if let Some(value) = value {
match value {
ScalarValue::Int8(value) => {
max_accumulate!(self, value, Int8Array, Int64, i64);
max_accumulate!(self, value, Int8Array, Int8);
}
ScalarValue::Int16(value) => {
max_accumulate!(self, value, Int16Array, Int64, i64)
max_accumulate!(self, value, Int16Array, Int16)
}
ScalarValue::Int32(value) => {
max_accumulate!(self, value, Int32Array, Int64, i64)
max_accumulate!(self, value, Int32Array, Int32)
}
ScalarValue::Int64(value) => {
max_accumulate!(self, value, Int64Array, Int64, i64)
max_accumulate!(self, value, Int64Array, Int64)
}
ScalarValue::UInt8(value) => {
max_accumulate!(self, value, UInt8Array, UInt64, u64)
max_accumulate!(self, value, UInt8Array, UInt8)
}
ScalarValue::UInt16(value) => {
max_accumulate!(self, value, UInt16Array, UInt64, u64)
max_accumulate!(self, value, UInt16Array, UInt16)
}
ScalarValue::UInt32(value) => {
max_accumulate!(self, value, UInt32Array, UInt64, u64)
max_accumulate!(self, value, UInt32Array, UInt32)
}
ScalarValue::UInt64(value) => {
max_accumulate!(self, value, UInt64Array, UInt64, u64)
max_accumulate!(self, value, UInt64Array, UInt64)
}
ScalarValue::Float32(value) => {
max_accumulate!(self, value, Float32Array, Float32, f32)
max_accumulate!(self, value, Float32Array, Float32)
}
ScalarValue::Float64(value) => {
max_accumulate!(self, value, Float64Array, Float64, f64)
max_accumulate!(self, value, Float64Array, Float64)
}
other => {
return Err(ExecutionError::General(format!(
Expand Down Expand Up @@ -616,20 +603,7 @@ impl Min {

impl AggregateExpr for Min {
fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
match self.expr.data_type(input_schema)? {
DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => {
Ok(DataType::Int64)
}
DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => {
Ok(DataType::UInt64)
}
DataType::Float32 => Ok(DataType::Float32),
DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"MIN does not support {:?}",
other
))),
}
self.expr.data_type(input_schema)
}

fn evaluate_input(&self, batch: &RecordBatch) -> Result<ArrayRef> {
Expand All @@ -646,21 +620,21 @@ impl AggregateExpr for Min {
}

macro_rules! min_accumulate {
($SELF:ident, $VALUE:expr, $ARRAY_TYPE:ident, $SCALAR_VARIANT:ident, $TY:ty) => {{
($SELF:ident, $VALUE:expr, $ARRAY_TYPE:ident, $SCALAR_VARIANT:ident) => {{
$SELF.min = match $SELF.min {
Some(ScalarValue::$SCALAR_VARIANT(n)) => {
if n < ($VALUE as $TY) {
if n < ($VALUE) {
Some(ScalarValue::$SCALAR_VARIANT(n))
} else {
Some(ScalarValue::$SCALAR_VARIANT($VALUE as $TY))
Some(ScalarValue::$SCALAR_VARIANT($VALUE))
}
}
Some(_) => {
return Err(ExecutionError::InternalError(
"Unexpected ScalarValue variant".to_string(),
))
}
None => Some(ScalarValue::$SCALAR_VARIANT($VALUE as $TY)),
None => Some(ScalarValue::$SCALAR_VARIANT($VALUE)),
};
}};
}
Expand All @@ -674,34 +648,34 @@ impl Accumulator for MinAccumulator {
if let Some(value) = value {
match value {
ScalarValue::Int8(value) => {
min_accumulate!(self, value, Int8Array, Int64, i64);
min_accumulate!(self, value, Int8Array, Int8);
}
ScalarValue::Int16(value) => {
min_accumulate!(self, value, Int16Array, Int64, i64)
min_accumulate!(self, value, Int16Array, Int16)
}
ScalarValue::Int32(value) => {
min_accumulate!(self, value, Int32Array, Int64, i64)
min_accumulate!(self, value, Int32Array, Int32)
}
ScalarValue::Int64(value) => {
min_accumulate!(self, value, Int64Array, Int64, i64)
min_accumulate!(self, value, Int64Array, Int64)
}
ScalarValue::UInt8(value) => {
min_accumulate!(self, value, UInt8Array, UInt64, u64)
min_accumulate!(self, value, UInt8Array, UInt8)
}
ScalarValue::UInt16(value) => {
min_accumulate!(self, value, UInt16Array, UInt64, u64)
min_accumulate!(self, value, UInt16Array, UInt16)
}
ScalarValue::UInt32(value) => {
min_accumulate!(self, value, UInt32Array, UInt64, u64)
min_accumulate!(self, value, UInt32Array, UInt32)
}
ScalarValue::UInt64(value) => {
min_accumulate!(self, value, UInt64Array, UInt64, u64)
min_accumulate!(self, value, UInt64Array, UInt64)
}
ScalarValue::Float32(value) => {
min_accumulate!(self, value, Float32Array, Float32, f32)
min_accumulate!(self, value, Float32Array, Float32)
}
ScalarValue::Float64(value) => {
min_accumulate!(self, value, Float64Array, Float64, f64)
min_accumulate!(self, value, Float64Array, Float64)
}
other => {
return Err(ExecutionError::General(format!(
Expand Down Expand Up @@ -1481,7 +1455,7 @@ mod tests {
let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);

let max = max(col("a"));
assert_eq!(DataType::Int64, max.data_type(&schema)?);
assert_eq!(DataType::Int32, max.data_type(&schema)?);

// after the aggr expression is applied, the schema changes to:
let schema = Schema::new(vec![
Expand All @@ -1490,7 +1464,7 @@ mod tests {
]);

let combiner = max.create_reducer("Max(a)");
assert_eq!(DataType::Int64, combiner.data_type(&schema)?);
assert_eq!(DataType::Int32, combiner.data_type(&schema)?);

Ok(())
}
Expand All @@ -1500,15 +1474,15 @@ mod tests {
let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);

let min = min(col("a"));
assert_eq!(DataType::Int64, min.data_type(&schema)?);
assert_eq!(DataType::Int32, min.data_type(&schema)?);

// after the aggr expression is applied, the schema changes to:
let schema = Schema::new(vec![
schema.field(0).clone(),
Field::new("MIN(a)", min.data_type(&schema)?, false),
]);
let combiner = min.create_reducer("MIN(a)");
assert_eq!(DataType::Int64, combiner.data_type(&schema)?);
assert_eq!(DataType::Int32, combiner.data_type(&schema)?);

Ok(())
}
Expand Down Expand Up @@ -1562,7 +1536,7 @@ mod tests {
let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(a)])?;

assert_eq!(do_max(&batch)?, Some(ScalarValue::Int64(5)));
assert_eq!(do_max(&batch)?, Some(ScalarValue::Int32(5)));

Ok(())
}
Expand All @@ -1574,7 +1548,7 @@ mod tests {
let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(a)])?;

assert_eq!(do_min(&batch)?, Some(ScalarValue::Int64(1)));
assert_eq!(do_min(&batch)?, Some(ScalarValue::Int32(1)));

Ok(())
}
Expand Down Expand Up @@ -1610,7 +1584,7 @@ mod tests {
let a = Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5)]);
let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(a)])?;

assert_eq!(do_max(&batch)?, Some(ScalarValue::Int64(5)));
assert_eq!(do_max(&batch)?, Some(ScalarValue::Int32(5)));

Ok(())
}
Expand All @@ -1622,7 +1596,7 @@ mod tests {
let a = Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5)]);
let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(a)])?;

assert_eq!(do_min(&batch)?, Some(ScalarValue::Int64(1)));
assert_eq!(do_min(&batch)?, Some(ScalarValue::Int32(1)));

Ok(())
}
Expand Down Expand Up @@ -1706,7 +1680,7 @@ mod tests {
let a = UInt32Array::from(vec![1_u32, 2_u32, 3_u32, 4_u32, 5_u32]);
let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(a)])?;

assert_eq!(do_max(&batch)?, Some(ScalarValue::UInt64(5_u64)));
assert_eq!(do_max(&batch)?, Some(ScalarValue::UInt32(5_u32)));

Ok(())
}
Expand All @@ -1718,7 +1692,7 @@ mod tests {
let a = UInt32Array::from(vec![1_u32, 2_u32, 3_u32, 4_u32, 5_u32]);
let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(a)])?;

assert_eq!(do_min(&batch)?, Some(ScalarValue::UInt64(1_u64)));
assert_eq!(do_min(&batch)?, Some(ScalarValue::UInt32(1_u32)));

Ok(())
}
Expand Down
Loading