Skip to content
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

Jump to bottom

[Minor] Clean up DecimalArray API Usage #1869

Merged
merged 11 commits into from
Mar 6, 2022
Merged

[Minor] Clean up DecimalArray API Usage #1869

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
65209a3
Clean up decimal array creation
alamb Feb 17, 2022
6324846
Refactor a bit more
alamb Feb 17, 2022
7cf30eb
Update the next
alamb Feb 17, 2022
9c8f7e2
cleanup
alamb Feb 17, 2022
613364c
update
alamb Feb 17, 2022
3eadf00
port over min/max
alamb Feb 18, 2022
9757ffc
Update sum
alamb Feb 18, 2022
0dd9cc2
Use max scale / precision from arrow
alamb Feb 18, 2022
c0a4784
fmt
alamb Mar 5, 2022
b42933a
Fixup
alamb Mar 5, 2022
009dd7c
clippy
alamb Mar 5, 2022
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
4 changes: 1 addition & 3 deletions datafusion-common/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,4 @@ mod scalar;
pub use column::Column;
pub use dfschema::{DFField, DFSchema, DFSchemaRef, ExprSchema, ToDFSchema};
pub use error::{DataFusionError, Result};
pub use scalar::{
ScalarType, ScalarValue, MAX_PRECISION_FOR_DECIMAL128, MAX_SCALE_FOR_DECIMAL128,
};
pub use scalar::{ScalarType, ScalarValue};
45 changes: 10 additions & 35 deletions datafusion-common/src/scalar.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ use arrow::{
Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, IntervalUnit, TimeUnit,
TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
TimestampSecondType, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
DECIMAL_MAX_PRECISION,
},
};
use ordered_float::OrderedFloat;
Expand All @@ -34,11 +35,6 @@ use std::convert::{Infallible, TryInto};
use std::str::FromStr;
use std::{convert::TryFrom, fmt, iter::repeat, sync::Arc};

// TODO may need to be moved to arrow-rs
/// The max precision and scale for decimal128
pub const MAX_PRECISION_FOR_DECIMAL128: usize = 38;
pub const MAX_SCALE_FOR_DECIMAL128: usize = 38;

/// Represents a dynamically typed, nullable single value.
/// This is the single-valued counter-part of arrow’s `Array`.
#[derive(Clone)]
Expand Down Expand Up @@ -542,7 +538,7 @@ impl ScalarValue {
scale: usize,
) -> Result<Self> {
// make sure the precision and scale is valid
if precision <= MAX_PRECISION_FOR_DECIMAL128 && scale <= precision {
if precision <= DECIMAL_MAX_PRECISION && scale <= precision {
return Ok(ScalarValue::Decimal128(Some(value), precision, scale));
}
return Err(DataFusionError::Internal(format!(
Expand Down Expand Up @@ -985,26 +981,15 @@ impl ScalarValue {
precision: &usize,
scale: &usize,
) -> Result<DecimalArray> {
// collect the value as Option<i128>
let array = scalars
.into_iter()
.map(|element: ScalarValue| match element {
ScalarValue::Decimal128(v1, _, _) => v1,
_ => unreachable!(),
})
.collect::<Vec<Option<i128>>>();

// build the decimal array using the Decimal Builder
let mut builder = DecimalBuilder::new(array.len(), *precision, *scale);
array.iter().for_each(|element| match element {
None => {
builder.append_null().unwrap();
}
Some(v) => {
builder.append_value(*v).unwrap();
}
});
Ok(builder.finish())
.collect::<DecimalArray>()
.with_precision_and_scale(*precision, *scale)?;
Ok(array)
}

fn iter_to_array_list(
Expand Down Expand Up @@ -1080,21 +1065,11 @@ impl ScalarValue {
scale: &usize,
size: usize,
) -> DecimalArray {
let mut builder = DecimalBuilder::new(size, *precision, *scale);
match value {
None => {
for _i in 0..size {
builder.append_null().unwrap();
}
}
Some(v) => {
let v = *v;
for _i in 0..size {
builder.append_value(v).unwrap();
}
}
};
builder.finish()
std::iter::repeat(value)
.take(size)
.collect::<DecimalArray>()
.with_precision_and_scale(*precision, *scale)
.unwrap()
}

/// Converts a scalar value into an array of `size` rows.
Expand Down
7 changes: 3 additions & 4 deletions datafusion-physical-expr/src/coercion_rule/binary_rule.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,9 @@

//! Coercion rules for matching argument types for binary operators

use arrow::datatypes::DataType;
use arrow::datatypes::{DataType, DECIMAL_MAX_PRECISION, DECIMAL_MAX_SCALE};
use datafusion_common::DataFusionError;
use datafusion_common::Result;
use datafusion_common::{MAX_PRECISION_FOR_DECIMAL128, MAX_SCALE_FOR_DECIMAL128};
use datafusion_expr::Operator;

/// Coercion rules for all binary operators. Returns the output type
Expand Down Expand Up @@ -261,8 +260,8 @@ fn mathematics_numerical_coercion(

fn create_decimal_type(precision: usize, scale: usize) -> DataType {
DataType::Decimal(
MAX_PRECISION_FOR_DECIMAL128.min(precision),
MAX_SCALE_FOR_DECIMAL128.min(scale),
DECIMAL_MAX_PRECISION.min(precision),
DECIMAL_MAX_SCALE.min(scale),
)
}

Expand Down
47 changes: 22 additions & 25 deletions datafusion-physical-expr/src/expressions/average.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,15 +23,13 @@ use std::sync::Arc;

use crate::{AggregateExpr, PhysicalExpr};
use arrow::compute;
use arrow::datatypes::DataType;
use arrow::datatypes::{DataType, DECIMAL_MAX_PRECISION, DECIMAL_MAX_SCALE};
use arrow::{
array::{ArrayRef, UInt64Array},
datatypes::Field,
};
use datafusion_common::ScalarValue;
use datafusion_common::{DataFusionError, Result};
use datafusion_common::{
ScalarValue, MAX_PRECISION_FOR_DECIMAL128, MAX_SCALE_FOR_DECIMAL128,
};
use datafusion_expr::Accumulator;

use super::{format_state_name, sum};
Expand All @@ -50,8 +48,8 @@ pub fn avg_return_type(arg_type: &DataType) -> Result<DataType> {
DataType::Decimal(precision, scale) => {
// in the spark, the result type is DECIMAL(min(38,precision+4), min(38,scale+4)).
// ref: https://github.com/apache/spark/blob/fcf636d9eb8d645c24be3db2d599aba2d7e2955a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala#L66
let new_precision = MAX_PRECISION_FOR_DECIMAL128.min(*precision + 4);
let new_scale = MAX_SCALE_FOR_DECIMAL128.min(*scale + 4);
let new_precision = DECIMAL_MAX_PRECISION.min(*precision + 4);
let new_scale = DECIMAL_MAX_SCALE.min(*scale + 4);
Ok(DataType::Decimal(new_precision, new_scale))
}
DataType::Int8
Expand Down Expand Up @@ -237,11 +235,12 @@ mod tests {
#[test]
fn avg_decimal() -> Result<()> {
// test agg
let mut decimal_builder = DecimalBuilder::new(6, 10, 0);
for i in 1..7 {
decimal_builder.append_value(i as i128)?;
}
let array: ArrayRef = Arc::new(decimal_builder.finish());
let array: ArrayRef = Arc::new(
(1..7)
.map(Some)
.collect::<DecimalArray>()
.with_precision_and_scale(10, 0)?,
);

generic_test_op!(
array,
Expand All @@ -254,15 +253,12 @@ mod tests {

#[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());
let array: ArrayRef = Arc::new(
(1..6)
.map(|i| if i == 2 { None } else { Some(i) })
.collect::<DecimalArray>()
.with_precision_and_scale(10, 0)?,
);
generic_test_op!(
array,
DataType::Decimal(10, 0),
Expand All @@ -275,11 +271,12 @@ mod tests {
#[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());
let array: ArrayRef = Arc::new(
std::iter::repeat(None)
.take(6)
.collect::<DecimalArray>()
.with_precision_and_scale(10, 0)?,
);
generic_test_op!(
array,
DataType::Decimal(10, 0),
Expand Down
80 changes: 54 additions & 26 deletions datafusion-physical-expr/src/expressions/cast.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,8 +161,8 @@ mod tests {
use crate::expressions::col;
use arrow::{
array::{
Array, DecimalArray, DecimalBuilder, Float32Array, Float64Array, Int16Array,
Int32Array, Int64Array, Int8Array, StringArray, Time64NanosecondArray,
Array, DecimalArray, Float32Array, Float64Array, Int16Array, Int32Array,
Int64Array, Int8Array, StringArray, Time64NanosecondArray,
TimestampNanosecondArray, UInt32Array,
},
datatypes::*,
Expand Down Expand Up @@ -268,23 +268,16 @@ mod tests {
}};
}

fn create_decimal_array(
array: &[i128],
precision: usize,
scale: usize,
) -> Result<DecimalArray> {
let mut decimal_builder = DecimalBuilder::new(array.len(), precision, scale);
for value in array {
decimal_builder.append_value(*value)?
}
decimal_builder.append_null()?;
Ok(decimal_builder.finish())
}

#[test]
fn test_cast_decimal_to_decimal() -> Result<()> {
let array: Vec<i128> = vec![1234, 2222, 3, 4000, 5000];
let decimal_array = create_decimal_array(&array, 10, 3)?;
let array = vec![1234, 2222, 3, 4000, 5000];

let decimal_array = array
.iter()
.map(|v| Some(*v))
.collect::<DecimalArray>()
.with_precision_and_scale(10, 3)?;

generic_decimal_to_other_test_cast!(
decimal_array,
DataType::Decimal(10, 3),
Expand All @@ -301,7 +294,12 @@ mod tests {
DEFAULT_DATAFUSION_CAST_OPTIONS
);

let decimal_array = create_decimal_array(&array, 10, 3)?;
let decimal_array = array
.iter()
.map(|v| Some(*v))
.collect::<DecimalArray>()
.with_precision_and_scale(10, 3)?;

generic_decimal_to_other_test_cast!(
decimal_array,
DataType::Decimal(10, 3),
Expand All @@ -323,9 +321,12 @@ mod tests {

#[test]
fn test_cast_decimal_to_numeric() -> Result<()> {
let array: Vec<i128> = vec![1, 2, 3, 4, 5];
let array = vec![Some(1), Some(2), Some(3), Some(4), Some(5), None];
// decimal to i8
let decimal_array = create_decimal_array(&array, 10, 0)?;
let decimal_array = array
.iter()
.collect::<DecimalArray>()
.with_precision_and_scale(10, 0)?;
generic_decimal_to_other_test_cast!(
decimal_array,
DataType::Decimal(10, 0),
Expand All @@ -341,8 +342,12 @@ mod tests {
],
DEFAULT_DATAFUSION_CAST_OPTIONS
);

// decimal to i16
let decimal_array = create_decimal_array(&array, 10, 0)?;
let decimal_array = array
.iter()
.collect::<DecimalArray>()
.with_precision_and_scale(10, 0)?;
generic_decimal_to_other_test_cast!(
decimal_array,
DataType::Decimal(10, 0),
Expand All @@ -358,8 +363,12 @@ mod tests {
],
DEFAULT_DATAFUSION_CAST_OPTIONS
);

// decimal to i32
let decimal_array = create_decimal_array(&array, 10, 0)?;
let decimal_array = array
.iter()
.collect::<DecimalArray>()
.with_precision_and_scale(10, 0)?;
generic_decimal_to_other_test_cast!(
decimal_array,
DataType::Decimal(10, 0),
Expand All @@ -375,8 +384,12 @@ mod tests {
],
DEFAULT_DATAFUSION_CAST_OPTIONS
);

// decimal to i64
let decimal_array = create_decimal_array(&array, 10, 0)?;
let decimal_array = array
.iter()
.collect::<DecimalArray>()
.with_precision_and_scale(10, 0)?;
generic_decimal_to_other_test_cast!(
decimal_array,
DataType::Decimal(10, 0),
Expand All @@ -392,9 +405,20 @@ mod tests {
],
DEFAULT_DATAFUSION_CAST_OPTIONS
);

// decimal to float32
let array: Vec<i128> = vec![1234, 2222, 3, 4000, 5000];
let decimal_array = create_decimal_array(&array, 10, 3)?;
let array = vec![
Some(1234),
Some(2222),
Some(3),
Some(4000),
Some(5000),
None,
];
let decimal_array = array
.iter()
.collect::<DecimalArray>()
.with_precision_and_scale(10, 3)?;
generic_decimal_to_other_test_cast!(
decimal_array,
DataType::Decimal(10, 3),
Expand All @@ -410,8 +434,12 @@ mod tests {
],
DEFAULT_DATAFUSION_CAST_OPTIONS
);

// decimal to float64
let decimal_array = create_decimal_array(&array, 20, 6)?;
let decimal_array = array
.into_iter()
.collect::<DecimalArray>()
.with_precision_and_scale(20, 6)?;
generic_decimal_to_other_test_cast!(
decimal_array,
DataType::Decimal(20, 6),
Expand Down
Loading