Incorporate dyn scalar kernels #1685
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -16,97 +16,98 @@ | |
| # under the License. | ||
|
|
||
| [package] | ||
| authors = ["Apache Arrow <[email protected]>"] | ||
| description = "DataFusion is an in-memory query engine that uses Apache Arrow as the memory model" | ||
| edition = "2021" | ||
| homepage = "https://github.com/apache/arrow-datafusion" | ||
| include = [ | ||
| "benches/*.rs", | ||
| "src/**/*.rs", | ||
| "Cargo.toml", | ||
| ] | ||
| keywords = ["arrow", "query", "sql"] | ||
| license = "Apache-2.0" | ||
| name = "datafusion" | ||
| readme = "../README.md" | ||
| repository = "https://github.com/apache/arrow-datafusion" | ||
| rust-version = "1.58" | ||
| version = "6.0.0" | ||
|
|
||
| [lib] | ||
| name = "datafusion" | ||
| path = "src/lib.rs" | ||
|
|
||
| [features] | ||
| crypto_expressions = ["md-5", "sha2", "blake2", "blake3"] | ||
| default = ["crypto_expressions", "regex_expressions", "unicode_expressions"] | ||
| pyarrow = ["pyo3", "arrow/pyarrow"] | ||
| regex_expressions = ["regex"] | ||
| simd = ["arrow/simd"] | ||
| unicode_expressions = ["unicode-segmentation"] | ||
| # Used for testing ONLY: causes all values to hash to the same value (test for collisions) | ||
| force_hash_collisions = [] | ||
| # Used to enable the avro format | ||
| avro = ["avro-rs", "num-traits"] | ||
|
|
||
| [dependencies] | ||
| ahash = {version = "0.7", default-features = false} | ||
| arrow = {version = "8.0.0", features = ["prettyprint"]} | ||
| async-trait = "0.1.41" | ||
| avro-rs = {version = "0.13", features = ["snappy"], optional = true} | ||
| blake2 = {version = "^0.10.2", optional = true} | ||
| blake3 = {version = "1.0", optional = true} | ||
| chrono = {version = "0.4", default-features = false} | ||
| futures = "0.3" | ||
| hashbrown = {version = "0.12", features = ["raw"]} | ||
| lazy_static = {version = "^1.4.0"} | ||
| log = "^0.4" | ||
| md-5 = {version = "^0.10.0", optional = true} | ||
| num = "0.4" | ||
| num-traits = {version = "0.2", optional = true} | ||
| num_cpus = "1.13.0" | ||
| ordered-float = "2.0" | ||
| parquet = {version = "8.0.0", features = ["arrow"]} | ||
| paste = "^1.0" | ||
| pin-project-lite = "^0.2.7" | ||
| pyo3 = {version = "0.15", optional = true} | ||
| rand = "0.8" | ||
| regex = {version = "^1.4.3", optional = true} | ||
| sha2 = {version = "^0.10.1", optional = true} | ||
| smallvec = {version = "1.6", features = ["union"]} | ||
| sqlparser = "0.13" | ||
| tempfile = "3" | ||
| tokio = {version = "1.0", features = ["macros", "rt", "rt-multi-thread", "sync", "fs"]} | ||
| tokio-stream = "0.1" | ||
| unicode-segmentation = {version = "^1.7.1", optional = true} | ||
|
|
||
| [dev-dependencies] | ||
| criterion = "0.3" | ||
| doc-comment = "0.3" | ||
|
|
||
| [[bench]] | ||
| harness = false | ||
| name = "aggregate_query_sql" | ||
|
|
||
| [[bench]] | ||
| harness = false | ||
| name = "sort_limit_query_sql" | ||
|
|
||
| [[bench]] | ||
| harness = false | ||
| name = "math_query_sql" | ||
|
|
||
| [[bench]] | ||
| harness = false | ||
| name = "filter_query_sql" | ||
|
|
||
| [[bench]] | ||
| harness = false | ||
| name = "window_query_sql" | ||
|
|
||
| [[bench]] | ||
| harness = false | ||
| name = "scalar" | ||
|
|
||
| [[bench]] | ||
| harness = false | ||
| name = "physical_plan" | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -15,6 +15,7 @@ | |
| // specific language governing permissions and limitations | ||
| // under the License. | ||
|
|
||
| use std::convert::TryInto; | ||
| use std::{any::Any, sync::Arc}; | ||
|
|
||
| use arrow::array::TimestampMillisecondArray; | ||
|
|
@@ -28,6 +29,18 @@ use arrow::compute::kernels::comparison::{ | |
| eq_bool, eq_bool_scalar, gt_bool, gt_bool_scalar, gt_eq_bool, gt_eq_bool_scalar, | ||
| lt_bool, lt_bool_scalar, lt_eq_bool, lt_eq_bool_scalar, neq_bool, neq_bool_scalar, | ||
| }; | ||
| use arrow::compute::kernels::comparison::{ | ||
| eq_dyn_bool_scalar, gt_dyn_bool_scalar, gt_eq_dyn_bool_scalar, lt_dyn_bool_scalar, | ||
| lt_eq_dyn_bool_scalar, neq_dyn_bool_scalar, | ||
| }; | ||
| use arrow::compute::kernels::comparison::{ | ||
| eq_dyn_scalar, gt_dyn_scalar, gt_eq_dyn_scalar, lt_dyn_scalar, lt_eq_dyn_scalar, | ||
| neq_dyn_scalar, | ||
| }; | ||
| use arrow::compute::kernels::comparison::{ | ||
| eq_dyn_utf8_scalar, gt_dyn_utf8_scalar, gt_eq_dyn_utf8_scalar, lt_dyn_utf8_scalar, | ||
| lt_eq_dyn_utf8_scalar, neq_dyn_utf8_scalar, | ||
| }; | ||
| use arrow::compute::kernels::comparison::{ | ||
| eq_scalar, gt_eq_scalar, gt_scalar, lt_eq_scalar, lt_scalar, neq_scalar, | ||
| }; | ||
|
|
@@ -44,6 +57,8 @@ use arrow::datatypes::{ArrowNumericType, DataType, Schema, TimeUnit}; | |
| use arrow::error::ArrowError::DivideByZero; | ||
| use arrow::record_batch::RecordBatch; | ||
|
|
||
| use num::ToPrimitive; | ||
|
|
||
| use crate::error::{DataFusionError, Result}; | ||
| use crate::logical_plan::Operator; | ||
| use crate::physical_plan::coercion_rule::binary_rule::coerce_types; | ||
|
|
@@ -429,6 +444,24 @@ macro_rules! compute_utf8_op_scalar { | |
| }}; | ||
| } | ||
|
|
||
| /// Invoke a compute kernel on a data array and a scalar value | ||
| macro_rules! compute_utf8_op_dyn_scalar { | ||
| ($LEFT:expr, $RIGHT:expr, $OP:ident) => {{ | ||
| if let ScalarValue::Utf8(Some(string_value)) = $RIGHT { | ||
| Ok(Arc::new(paste::expr! {[<$OP _dyn_utf8_scalar>]}( | ||
| $LEFT, | ||
| &string_value, | ||
| )?)) | ||
| } else { | ||
| Err(DataFusionError::Internal(format!( | ||
| "compute_utf8_op_scalar for '{}' failed to cast literal value {}", | ||
| stringify!($OP), | ||
| $RIGHT | ||
| ))) | ||
| } | ||
| }}; | ||
| } | ||
|
|
||
| /// Invoke a compute kernel on a boolean data array and a scalar value | ||
| macro_rules! compute_bool_op_scalar { | ||
| ($LEFT:expr, $RIGHT:expr, $OP:ident, $DT:ident) => {{ | ||
|
|
@@ -446,6 +479,18 @@ macro_rules! compute_bool_op_scalar { | |
| }}; | ||
| } | ||
|
|
||
| /// Invoke a compute kernel on a boolean data array and a scalar value | ||
| macro_rules! compute_bool_op_dyn_scalar { | ||
| ($LEFT:expr, $RIGHT:expr, $OP:ident) => {{ | ||
| // generate the scalar function name, such as lt_dyn_bool_scalar, from the $OP parameter | ||
| // (which could have a value of lt) and the suffix _scalar | ||
| Ok(Arc::new(paste::expr! {[<$OP _dyn_bool_scalar>]}( | ||
| $LEFT, | ||
| $RIGHT.try_into()?, | ||
| )?)) | ||
| }}; | ||
| } | ||
|
|
||
| /// Invoke a bool compute kernel on array(s) | ||
| macro_rules! compute_bool_op { | ||
| // invoke binary operator | ||
|
|
@@ -474,7 +519,7 @@ macro_rules! compute_bool_op { | |
| /// LEFT is array, RIGHT is scalar value | ||
| macro_rules! compute_op_scalar { | ||
| ($LEFT:expr, $RIGHT:expr, $OP:ident, $DT:ident) => {{ | ||
| // use std::convert::TryInto; | ||
|
|
||
| let ll = $LEFT | ||
| .as_any() | ||
| .downcast_ref::<$DT>() | ||
|
|
@@ -488,6 +533,19 @@ macro_rules! compute_op_scalar { | |
| }}; | ||
| } | ||
|
|
||
| /// Invoke a dyn compute kernel on a data array and a scalar value | ||
| /// LEFT is Primitive or Dictionart array of numeric values, RIGHT is scalar value | ||
| macro_rules! compute_op_dyn_scalar { | ||
| ($LEFT:expr, $RIGHT:expr, $OP:ident) => {{ | ||
| // generate the scalar function name, such as lt_dyn_scalar, from the $OP parameter | ||
| // (which could have a value of lt_dyn) and the suffix _scalar | ||
| Ok(Arc::new(paste::expr! {[<$OP _dyn_scalar>]}( | ||
| $LEFT, | ||
| $RIGHT, | ||
| )?)) | ||
| }}; | ||
| } | ||
|
|
||
| /// Invoke a compute kernel on array(s) | ||
| macro_rules! compute_op { | ||
| // invoke binary operator | ||
|
|
@@ -878,26 +936,92 @@ impl PhysicalExpr for BinaryExpr { | |
| } | ||
| } | ||
|
|
||
| /// The binary_array_op_scalar macro includes types that extend beyond the primitive, | ||
| /// such as Utf8 strings. | ||
| #[macro_export] | ||
| macro_rules! binary_array_op_dyn_scalar { | ||
| ($LEFT:expr, $RIGHT:expr, $OP:ident) => {{ | ||
| let is_numeric = DataType::is_numeric($LEFT.data_type()); | ||
| let is_numeric_dict = match $LEFT.data_type() { | ||
| DataType::Dictionary(_, val_type) => DataType::is_numeric(val_type), | ||
| _ => false | ||
| }; | ||
| let numeric_like = is_numeric | is_numeric_dict; | ||
|
|
||
| let is_string = ($LEFT.data_type() == &DataType::Utf8) | ($LEFT.data_type() == &DataType::LargeUtf8); | ||
|
There was a problem hiding this comment. It seems to me that the I wonder if we could do something like match $RIGHT {
ScalarValue::Utf8(v) => {
Ok(Arc::new(paste::expr! {[<$OP _dyn_scalar_utf8>]}(
$LEFT,
v,
)?))
}
..
ScalarValue::Int8(v) => {
Ok(Arc::new(paste::expr! {[<$OP _dyn_scalar>]}(
$LEFT,
v,
)?))
}
...
}Though we will probably need some sort of wrapper to handle types not yet supported in arrow-rs 🤔 There was a problem hiding this comment. @alamb this basically solved everything i think...we'll see how CI comes out |
||
| let is_string_dict = match $LEFT.data_type() { | ||
| DataType::Dictionary(_, val_type) => match **val_type { | ||
| DataType::Utf8 | DataType::LargeUtf8 => true, | ||
| _ => false | ||
| } | ||
| _ => false | ||
| }; | ||
| let string_like = is_string | is_string_dict; | ||
|
|
||
| let result: Result<Arc<dyn Array>> = if numeric_like { | ||
| compute_op_dyn_scalar!($LEFT, $RIGHT.try_into()?, $OP) | ||
| } else if string_like { | ||
| compute_utf8_op_dyn_scalar!($LEFT, $RIGHT, $OP) | ||
| } else { | ||
| let r: Result<Arc<dyn Array>> = match $LEFT.data_type() { | ||
|
|
||
| DataType::Decimal(_,_) => compute_decimal_op_scalar!($LEFT, $RIGHT, $OP, DecimalArray), | ||
| DataType::Boolean => compute_bool_op_dyn_scalar!($LEFT, $RIGHT, $OP), | ||
| DataType::Timestamp(TimeUnit::Nanosecond, _) => { | ||
| compute_op_scalar!($LEFT, $RIGHT, $OP, TimestampNanosecondArray) | ||
| } | ||
| DataType::Timestamp(TimeUnit::Microsecond, _) => { | ||
| compute_op_scalar!($LEFT, $RIGHT, $OP, TimestampMicrosecondArray) | ||
| } | ||
| DataType::Timestamp(TimeUnit::Millisecond, _) => { | ||
| compute_op_scalar!($LEFT, $RIGHT, $OP, TimestampMillisecondArray) | ||
| } | ||
| DataType::Timestamp(TimeUnit::Second, _) => { | ||
| compute_op_scalar!($LEFT, $RIGHT, $OP, TimestampSecondArray) | ||
| } | ||
| DataType::Date32 => { | ||
| compute_op_scalar!($LEFT, $RIGHT, $OP, Date32Array) | ||
| } | ||
| DataType::Date64 => { | ||
| compute_op_scalar!($LEFT, $RIGHT, $OP, Date64Array) | ||
| } | ||
| other => Err(DataFusionError::Internal(format!( | ||
| "Data type {:?} not supported for scalar operation '{}' on dyn array", | ||
| other, stringify!($OP) | ||
| ))), | ||
| }; | ||
| r | ||
| }; | ||
| Some(result) | ||
| }} | ||
| } | ||
|
|
||
| impl BinaryExpr { | ||
| /// Evaluate the expression of the left input is an array and | ||
| /// right is literal - use scalar operations | ||
| fn evaluate_array_scalar( | ||
| &self, | ||
| array: &dyn Array, | ||
| scalar: &ScalarValue, | ||
| ) -> Result<Option<Result<ArrayRef>>> { | ||
| let scalar_result = match &self.op { | ||
| Operator::Lt => { | ||
| binary_array_op_dyn_scalar!(array, scalar.clone(), lt) | ||
| } | ||
| Operator::LtEq => { | ||
| binary_array_op_dyn_scalar!(array, scalar.clone(), lt_eq) | ||
| } | ||
| Operator::Gt => { | ||
| binary_array_op_dyn_scalar!(array, scalar.clone(), gt) | ||
| } | ||
| Operator::GtEq => { | ||
| binary_array_op_dyn_scalar!(array, scalar.clone(), gt_eq) | ||
| } | ||
| Operator::Eq => { | ||
| binary_array_op_dyn_scalar!(array, scalar.clone(), eq) | ||
| } | ||
| Operator::NotEq => { | ||
| binary_array_op_dyn_scalar!(array, scalar.clone(), neq) | ||
| } | ||
| Operator::Like => { | ||
| binary_string_array_op_scalar!(array, scalar.clone(), like) | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
😭why are there so many meaningless changes?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I believe I accidentally started auto sorting my Cargo.toml. Or maybe it was related to upgrading cargo / rust? I'm not sure to be honest - trying to figure out what caused the change.