Skip to content

support sum/avg agg for decimal, change sum(float32) --> float64 #1408

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

Merged
alamb merged 7 commits into from
Dec 17, 2021
Merged

support sum/avg agg for decimal, change sum(float32) --> float64 #1408

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
2095685
support sum/avg agg for decimal
liukun4515 Dec 7, 2021
e3b005f
Merge remote-tracking branch 'upstream/master' into agg_sum/avg_decimal
liukun4515 Dec 8, 2021
3c19c41
support sum/avg agg for decimal
liukun4515 Dec 8, 2021
ff61c04
Merge remote-tracking branch 'upstream/master' into agg_sum/avg_decimal
liukun4515 Dec 13, 2021
9675e30
suppor the avg and add test
liukun4515 Dec 14, 2021
f4b5655
add comments and const
liukun4515 Dec 17, 2021
2d247ce
Merge remote-tracking branch 'upstream/master' into agg_sum/avg_decimal
liukun4515 Dec 17, 2021
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
59 changes: 58 additions & 1 deletion datafusion/src/execution/context.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1845,9 +1845,9 @@ mod tests {
#[tokio::test]
async fn aggregate_decimal_min() -> Result<()> {
let mut ctx = ExecutionContext::new();
// the data type of c1 is decimal(10,3)
ctx.register_table("d_table", test::table_with_decimal())
.unwrap();

let result = plan_and_collect(&mut ctx, "select min(c1) from d_table")
.await
.unwrap();
Expand All @@ -1858,13 +1858,18 @@ mod tests {
"| -100.009 |",
"+-----------------+",
];
assert_eq!(
&DataType::Decimal(10, 3),
result[0].schema().field(0).data_type()
);
assert_batches_sorted_eq!(expected, &result);
Ok(())
}

#[tokio::test]
async fn aggregate_decimal_max() -> Result<()> {
let mut ctx = ExecutionContext::new();
// the data type of c1 is decimal(10,3)
ctx.register_table("d_table", test::table_with_decimal())
.unwrap();

Expand All @@ -1878,6 +1883,58 @@ mod tests {
"| 110.009 |",
"+-----------------+",
];
assert_eq!(
&DataType::Decimal(10, 3),
result[0].schema().field(0).data_type()
);
assert_batches_sorted_eq!(expected, &result);
Ok(())
}

#[tokio::test]
async fn aggregate_decimal_sum() -> Result<()> {
let mut ctx = ExecutionContext::new();
// the data type of c1 is decimal(10,3)
ctx.register_table("d_table", test::table_with_decimal())
.unwrap();
let result = plan_and_collect(&mut ctx, "select sum(c1) from d_table")
.await
.unwrap();
let expected = vec![
"+-----------------+",
"| SUM(d_table.c1) |",
"+-----------------+",
"| 100.000 |",
"+-----------------+",
];
assert_eq!(
&DataType::Decimal(20, 3),

@alamb alamb Dec 14, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

👍

result[0].schema().field(0).data_type()
);
assert_batches_sorted_eq!(expected, &result);
Ok(())
}

#[tokio::test]
async fn aggregate_decimal_avg() -> Result<()> {
let mut ctx = ExecutionContext::new();
// the data type of c1 is decimal(10,3)
ctx.register_table("d_table", test::table_with_decimal())
.unwrap();
let result = plan_and_collect(&mut ctx, "select avg(c1) from d_table")
.await
.unwrap();
let expected = vec![
"+-----------------+",
"| AVG(d_table.c1) |",
"+-----------------+",
"| 5.0000000 |",
"+-----------------+",
];
assert_eq!(
&DataType::Decimal(14, 7),
result[0].schema().field(0).data_type()
);
assert_batches_sorted_eq!(expected, &result);
Ok(())
}
Expand Down
34 changes: 33 additions & 1 deletion datafusion/src/physical_plan/aggregates.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -426,7 +426,7 @@ mod tests {
| DataType::Int16
| DataType::Int32
| DataType::Int64 => DataType::Int64,
DataType::Float32 | DataType::Float64 => data_type.clone(),
DataType::Float32 | DataType::Float64 => DataType::Float64,
_ => data_type.clone(),
};

Expand Down Expand Up @@ -470,6 +470,29 @@ mod tests {
Ok(())
}

#[test]
fn test_sum_return_type() -> Result<()> {
let observed = return_type(&AggregateFunction::Sum, &[DataType::Int32])?;
assert_eq!(DataType::Int64, observed);

let observed = return_type(&AggregateFunction::Sum, &[DataType::UInt8])?;
assert_eq!(DataType::UInt64, observed);

let observed = return_type(&AggregateFunction::Sum, &[DataType::Float32])?;
assert_eq!(DataType::Float64, observed);

let observed = return_type(&AggregateFunction::Sum, &[DataType::Float64])?;
assert_eq!(DataType::Float64, observed);

let observed = return_type(&AggregateFunction::Sum, &[DataType::Decimal(10, 5)])?;
assert_eq!(DataType::Decimal(20, 5), observed);

let observed = return_type(&AggregateFunction::Sum, &[DataType::Decimal(35, 5)])?;
assert_eq!(DataType::Decimal(38, 5), observed);

Ok(())
}

#[test]
fn test_sum_no_utf8() {
let observed = return_type(&AggregateFunction::Sum, &[DataType::Utf8]);
Expand Down Expand Up @@ -504,6 +527,15 @@ mod tests {

let observed = return_type(&AggregateFunction::Avg, &[DataType::Float64])?;
assert_eq!(DataType::Float64, observed);

let observed = return_type(&AggregateFunction::Avg, &[DataType::Int32])?;
assert_eq!(DataType::Float64, observed);

let observed = return_type(&AggregateFunction::Avg, &[DataType::Decimal(10, 6)])?;
assert_eq!(DataType::Decimal(14, 10), observed);

let observed = return_type(&AggregateFunction::Avg, &[DataType::Decimal(36, 6)])?;
assert_eq!(DataType::Decimal(38, 10), observed);
Ok(())
}

Expand Down
3 changes: 1 addition & 2 deletions datafusion/src/physical_plan/coercion_rule/aggregate_rule.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,8 +193,7 @@ mod tests {
let input_types = vec![
vec![DataType::Int32],
vec![DataType::Float32],
// support the decimal data type
// vec![DataType::Decimal(20, 3)],
vec![DataType::Decimal(20, 3)],
];
for fun in funs {
for input_type in &input_types {
Expand Down
119 changes: 105 additions & 14 deletions datafusion/src/physical_plan/expressions/average.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,11 +38,18 @@ use super::{format_state_name, sum};
pub struct Avg {
name: String,
expr: Arc<dyn PhysicalExpr>,
data_type: DataType,
}

/// function return type of an average
pub fn avg_return_type(arg_type: &DataType) -> Result<DataType> {
match arg_type {
DataType::Decimal(precision, scale) => {
// the new precision and scale for return type of avg function

@alamb alamb Dec 14, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Can you please document document the rationale for the 4 and 38 constants below (or even better pull them into named constants somewhere)?

I also don't understand where the additional 4 came from. I tried to see if it was what postgres did, but when I checked the output schema for avg(numeric(10,3)) appears to be numeric without the precision or scale specified 🤔

(arrow_dev) alamb@MacBook-Pro-2:~/Downloads$ psql
psql (14.1)
Type "help" for help.

alamb=# create table test(x decimal(10, 3));
CREATE TABLE
alamb=# insert into test values (1.02);
INSERT 0 1
alamb=# create table test2 as select avg(x) from test;
SELECT 1

alamb=# select table_name, column_name, numeric_precision, numeric_scale, data_type from information_schema.columns where table_name='test2';
 table_name | column_name | numeric_precision | numeric_scale | data_type 
------------+-------------+-------------------+---------------+-----------
 test2      | avg         |                   |               | numeric
(1 row)

@liukun4515 liukun4515 Dec 15, 2021 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

This is the intention for #1418.
In the PG, we can create bigger precision for decimal than that in datafusion.
The decimal in datafusion, whose behavior may be different from that in PG.
So we should discuss them.
Now for the promotion precision or scale, I just follow the spark behavior.
@alamb

@alamb alamb Dec 15, 2021 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I think following the spark behavior is reasonable, but it should be documented (aka that the constant 4 came from spark). Otherwise in 3 months we'll be 🤔 why were the particular constants picked

@liukun4515 liukun4515 Dec 17, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

add comments to the code and issue #1461 to track this rule.
We can add the rule to our document later in the follow-up pull request.

let new_precision = 38.min(*precision + 4);
let new_scale = 38.min(*scale + 4);
Ok(DataType::Decimal(new_precision, new_scale))

@alamb alamb Dec 17, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Could potentially use make_decimal_type : https://github.com/apache/arrow-datafusion/blob/4b454d0363b034e3ebcdf88f6add2403cd77a23b/datafusion/src/sql/utils.rs#L510 here too

}
DataType::Int8
| DataType::Int16
| DataType::Int32
Expand Down Expand Up @@ -73,6 +80,7 @@ pub(crate) fn is_avg_support_arg_type(arg_type: &DataType) -> bool {
| DataType::Int64
| DataType::Float32
| DataType::Float64
| DataType::Decimal(_, _)
)
}

Expand All @@ -86,11 +94,15 @@ impl Avg {
// Average is always Float64, but Avg::new() has a data_type
// parameter to keep a consistent signature with the other
// Aggregate expressions.
assert_eq!(data_type, DataType::Float64);

Self {
name: name.into(),
expr,
match data_type {

@alamb alamb Dec 14, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

The comment above looks out of date -- I think it should simply be removed.

And perhaps we can change this code so it doesn't use unreachable as I think it would be fairly easy to reach this code by calling Avg::new(..) with some incorrect paramters

How about something like

assert!(matches!(data_type, DataType::Float64 | DataType::Decimal(_, _)));

Which I think might be easier to diagnose if anyone hits it

@liukun4515 liukun4515 Dec 17, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

grate comments

DataType::Float64 | DataType::Decimal(_, _) => Self {
name: name.into(),
expr,
data_type,
},
_ => {
unreachable!();
}
}
}
}
Expand All @@ -102,7 +114,14 @@ impl AggregateExpr for Avg {
}

fn field(&self) -> Result<Field> {
Ok(Field::new(&self.name, DataType::Float64, true))
Ok(Field::new(&self.name, self.data_type.clone(), true))
}

fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
Ok(Box::new(AvgAccumulator::try_new(
// avg is f64 or decimal
&self.data_type,

@alamb alamb Dec 14, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

if a sum of decimal(10,2) can be decimal(20,2) shouldn't the accumulator state also be decimal(20,2) to avoid overflow?

I think handling overflow is probably fine to for a later date / PR, but it is strange to me that there is a discrepancy between the type for sum and the accumulator type for computing avg

@liukun4515 liukun4515 Dec 17, 2021 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

The result type of phy expr(sum/avg) is same with each Accumulator, and it was decided by sum_return_type and avg_return_type.

If the column is decimal(8,2), the avg of this column must be less than 10^8-1, but we need more digits to represent the decimal part. For example, The avg of 3,4,6 is 4.333333....., we should increase the scale part.

For the sum agg, we just should increase the precision part, and the rule of adding 10 to precision is spark coercion rule for sum decimal. We can have our rules for decimal if we want.
@alamb
We can just follow the spark now, and change the rules if we want to define own rules.

@liukun4515 liukun4515 Dec 17, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

add issue to track the overflow.
#1460

)?))
}

fn state_fields(&self) -> Result<Vec<Field>> {
Expand All @@ -114,19 +133,12 @@ impl AggregateExpr for Avg {
),
Field::new(
&format_state_name(&self.name, "sum"),
DataType::Float64,
self.data_type.clone(),
true,
),
])
}

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>> {
vec![self.expr.clone()]
}
Expand Down Expand Up @@ -205,6 +217,17 @@ impl Accumulator for AvgAccumulator {
ScalarValue::Float64(e) => {
Ok(ScalarValue::Float64(e.map(|f| f / self.count as f64)))
}
ScalarValue::Decimal128(value, precision, scale) => {
Ok(match value {
None => ScalarValue::Decimal128(None, precision, scale),
// TODO add the checker for overflow the precision

@alamb alamb Dec 14, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

👍

Some(v) => ScalarValue::Decimal128(
Some(v / self.count as i128),
precision,
scale,
),
})
}
_ => Err(DataFusionError::Internal(
"Sum should be f64 on average".to_string(),
)),
Expand All @@ -220,6 +243,74 @@ mod tests {
use arrow::record_batch::RecordBatch;
use arrow::{array::*, datatypes::*};

#[test]
fn test_avg_return_data_type() -> Result<()> {
let data_type = DataType::Decimal(10, 5);
let result_type = avg_return_type(&data_type)?;
assert_eq!(DataType::Decimal(14, 9), result_type);

let data_type = DataType::Decimal(36, 10);
let result_type = avg_return_type(&data_type)?;
assert_eq!(DataType::Decimal(38, 14), result_type);
Ok(())
}

#[test]
fn avg_decimal() -> Result<()> {
// test agg
let mut decimal_builder = DecimalBuilder::new(6, 10, 0);
for i in 1..7 {
// the avg is 3.5, but we get the result of 3

@alamb alamb Dec 14, 2021

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I don't understand this comment: the result is 3.5 below (35000 with precision 14 and scale 4)

@liukun4515 liukun4515 Dec 15, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

This comment is out of data, I will remove this.
Sorry for that which make you confused.

decimal_builder.append_value(i as i128)?;
}
let array: ArrayRef = Arc::new(decimal_builder.finish());

generic_test_op!(
array,
DataType::Decimal(10, 0),
Avg,
ScalarValue::Decimal128(Some(35000), 14, 4),
DataType::Decimal(14, 4)
)
}

#[test]
fn avg_decimal_with_nulls() -> Result<()> {
let mut decimal_builder = DecimalBuilder::new(5, 10, 0);
for i in 1..6 {
if i == 2 {
decimal_builder.append_null()?;
} else {
decimal_builder.append_value(i)?;
}
}
let array: ArrayRef = Arc::new(decimal_builder.finish());
generic_test_op!(
array,
DataType::Decimal(10, 0),
Avg,
ScalarValue::Decimal128(Some(32500), 14, 4),
DataType::Decimal(14, 4)
)
}

#[test]
fn avg_decimal_all_nulls() -> Result<()> {
// test agg
let mut decimal_builder = DecimalBuilder::new(5, 10, 0);
for _i in 1..6 {
decimal_builder.append_null()?;
}
let array: ArrayRef = Arc::new(decimal_builder.finish());
generic_test_op!(
array,
DataType::Decimal(10, 0),
Avg,
ScalarValue::Decimal128(None, 14, 4),
DataType::Decimal(14, 4)
)
}

#[test]
fn avg_i32() -> Result<()> {
let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]));
Expand Down
Loading