This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new b04baa55 chore: move scalar_funcs into spark-expr (#712)
b04baa55 is described below
commit b04baa55e548961fe641c6e385e65956880249f3
Author: Arttu <[email protected]>
AuthorDate: Sun Jul 28 15:44:32 2024 +0200
chore: move scalar_funcs into spark-expr (#712)
---
native/Cargo.lock | 5 +-
native/Cargo.toml | 1 +
native/core/Cargo.toml | 4 +-
native/core/benches/hash.rs | 2 +-
.../datafusion/expressions/comet_scalar_funcs.rs | 185 +++++++++++++++++++++
.../src/execution/datafusion/expressions/mod.rs | 3 +-
native/core/src/execution/datafusion/mod.rs | 1 -
native/core/src/execution/datafusion/planner.rs | 2 +-
.../src/execution/datafusion/shuffle_writer.rs | 20 ++-
.../datafusion/util/spark_bloom_filter.rs | 5 +-
native/spark-expr/Cargo.toml | 7 +-
native/spark-expr/src/lib.rs | 3 +
.../expressions => spark-expr/src}/scalar_funcs.rs | 177 ++------------------
.../src}/scalar_funcs/chr.rs | 28 ++--
.../src}/scalar_funcs/hash_expressions.rs | 32 +++-
.../src}/scalar_funcs/hex.rs | 3 +-
.../src}/scalar_funcs/unhex.rs | 3 +-
.../datafusion => spark-expr/src}/spark_hash.rs | 29 +---
.../expressions => spark-expr/src}/xxhash64.rs | 0
19 files changed, 283 insertions(+), 227 deletions(-)
diff --git a/native/Cargo.lock b/native/Cargo.lock
index 58061074..a6748fa5 100644
--- a/native/Cargo.lock
+++ b/native/Cargo.lock
@@ -896,8 +896,6 @@ dependencies = [
"tempfile",
"thiserror",
"tokio",
- "twox-hash",
- "unicode-segmentation",
"zstd",
]
@@ -915,6 +913,7 @@ version = "0.2.0"
dependencies = [
"arrow",
"arrow-array",
+ "arrow-data",
"arrow-schema",
"chrono",
"chrono-tz 0.8.6",
@@ -929,6 +928,8 @@ dependencies = [
"rand",
"regex",
"thiserror",
+ "twox-hash",
+ "unicode-segmentation",
]
[[package]]
diff --git a/native/Cargo.toml b/native/Cargo.toml
index 2b07b1e5..fc9b68c4 100644
--- a/native/Cargo.toml
+++ b/native/Cargo.toml
@@ -51,6 +51,7 @@ datafusion-comet-proto = { path = "proto", version = "0.2.0" }
chrono = { version = "0.4", default-features = false, features = ["clock"] }
chrono-tz = { version = "0.8" }
num = "0.4"
+rand = "0.8"
regex = "1.9.6"
thiserror = "1"
diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml
index e396976d..5b5c237c 100644
--- a/native/core/Cargo.toml
+++ b/native/core/Cargo.toml
@@ -58,7 +58,7 @@ brotli = "3.3"
flate2 = "1.0"
lz4 = "1.24"
zstd = "0.11"
-rand = "0.8"
+rand = { workspace = true}
num = { workspace = true }
bytes = "1.5.0"
tempfile = "3.8.0"
@@ -70,7 +70,6 @@ datafusion = { workspace = true }
datafusion-expr = { workspace = true }
datafusion-physical-expr-common = { workspace = true }
datafusion-physical-expr = { workspace = true }
-unicode-segmentation = "^1.10.1"
once_cell = "1.18.0"
regex = { workspace = true }
crc32fast = "1.3.2"
@@ -85,7 +84,6 @@ jni = { version = "0.21", features = ["invocation"] }
lazy_static = "1.4"
assertables = "7"
hex = "0.4.3"
-twox-hash = "1.6.3"
[features]
default = []
diff --git a/native/core/benches/hash.rs b/native/core/benches/hash.rs
index d66d5892..039a6d5d 100644
--- a/native/core/benches/hash.rs
+++ b/native/core/benches/hash.rs
@@ -19,10 +19,10 @@
mod common;
use arrow_array::ArrayRef;
-use
comet::execution::datafusion::expressions::scalar_funcs::{spark_murmur3_hash,
spark_xxhash64};
use comet::execution::kernels::hash;
use common::*;
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use datafusion_comet_spark_expr::scalar_funcs::{spark_murmur3_hash,
spark_xxhash64};
use datafusion_common::ScalarValue;
use datafusion_expr::ColumnarValue;
use std::sync::Arc;
diff --git
a/native/core/src/execution/datafusion/expressions/comet_scalar_funcs.rs
b/native/core/src/execution/datafusion/expressions/comet_scalar_funcs.rs
new file mode 100644
index 00000000..70cbdeba
--- /dev/null
+++ b/native/core/src/execution/datafusion/expressions/comet_scalar_funcs.rs
@@ -0,0 +1,185 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow_schema::DataType;
+use datafusion_comet_spark_expr::scalar_funcs::hash_expressions::{
+ spark_sha224, spark_sha256, spark_sha384, spark_sha512,
+};
+use datafusion_comet_spark_expr::scalar_funcs::{
+ spark_ceil, spark_decimal_div, spark_floor, spark_hex, spark_isnan,
spark_make_decimal,
+ spark_murmur3_hash, spark_round, spark_rpad, spark_unhex,
spark_unscaled_value, spark_xxhash64,
+ SparkChrFunc,
+};
+use datafusion_common::{DataFusionError, Result as DataFusionResult};
+use datafusion_expr::registry::FunctionRegistry;
+use datafusion_expr::{
+ ColumnarValue, ScalarFunctionImplementation, ScalarUDF, ScalarUDFImpl,
Signature, Volatility,
+};
+use std::any::Any;
+use std::fmt::Debug;
+use std::sync::Arc;
+
+macro_rules! make_comet_scalar_udf {
+ ($name:expr, $func:ident, $data_type:ident) => {{
+ let scalar_func = CometScalarFunction::new(
+ $name.to_string(),
+ Signature::variadic_any(Volatility::Immutable),
+ $data_type.clone(),
+ Arc::new(move |args| $func(args, &$data_type)),
+ );
+ Ok(Arc::new(ScalarUDF::new_from_impl(scalar_func)))
+ }};
+ ($name:expr, $func:expr, without $data_type:ident) => {{
+ let scalar_func = CometScalarFunction::new(
+ $name.to_string(),
+ Signature::variadic_any(Volatility::Immutable),
+ $data_type,
+ $func,
+ );
+ Ok(Arc::new(ScalarUDF::new_from_impl(scalar_func)))
+ }};
+}
+
+/// Create a physical scalar function.
+pub fn create_comet_physical_fun(
+ fun_name: &str,
+ data_type: DataType,
+ registry: &dyn FunctionRegistry,
+) -> Result<Arc<ScalarUDF>, DataFusionError> {
+ match fun_name {
+ "ceil" => {
+ make_comet_scalar_udf!("ceil", spark_ceil, data_type)
+ }
+ "floor" => {
+ make_comet_scalar_udf!("floor", spark_floor, data_type)
+ }
+ "rpad" => {
+ let func = Arc::new(spark_rpad);
+ make_comet_scalar_udf!("rpad", func, without data_type)
+ }
+ "round" => {
+ make_comet_scalar_udf!("round", spark_round, data_type)
+ }
+ "unscaled_value" => {
+ let func = Arc::new(spark_unscaled_value);
+ make_comet_scalar_udf!("unscaled_value", func, without data_type)
+ }
+ "make_decimal" => {
+ make_comet_scalar_udf!("make_decimal", spark_make_decimal,
data_type)
+ }
+ "hex" => {
+ let func = Arc::new(spark_hex);
+ make_comet_scalar_udf!("hex", func, without data_type)
+ }
+ "unhex" => {
+ let func = Arc::new(spark_unhex);
+ make_comet_scalar_udf!("unhex", func, without data_type)
+ }
+ "decimal_div" => {
+ make_comet_scalar_udf!("decimal_div", spark_decimal_div, data_type)
+ }
+ "murmur3_hash" => {
+ let func = Arc::new(spark_murmur3_hash);
+ make_comet_scalar_udf!("murmur3_hash", func, without data_type)
+ }
+ "xxhash64" => {
+ let func = Arc::new(spark_xxhash64);
+ make_comet_scalar_udf!("xxhash64", func, without data_type)
+ }
+ "chr" =>
Ok(Arc::new(ScalarUDF::new_from_impl(SparkChrFunc::default()))),
+ "isnan" => {
+ let func = Arc::new(spark_isnan);
+ make_comet_scalar_udf!("isnan", func, without data_type)
+ }
+ "sha224" => {
+ let func = Arc::new(spark_sha224);
+ make_comet_scalar_udf!("sha224", func, without data_type)
+ }
+ "sha256" => {
+ let func = Arc::new(spark_sha256);
+ make_comet_scalar_udf!("sha256", func, without data_type)
+ }
+ "sha384" => {
+ let func = Arc::new(spark_sha384);
+ make_comet_scalar_udf!("sha384", func, without data_type)
+ }
+ "sha512" => {
+ let func = Arc::new(spark_sha512);
+ make_comet_scalar_udf!("sha512", func, without data_type)
+ }
+ _ => registry.udf(fun_name).map_err(|e| {
+ DataFusionError::Execution(format!(
+ "Function {fun_name} not found in the registry: {e}",
+ ))
+ }),
+ }
+}
+
+struct CometScalarFunction {
+ name: String,
+ signature: Signature,
+ data_type: DataType,
+ func: ScalarFunctionImplementation,
+}
+
+impl Debug for CometScalarFunction {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ f.debug_struct("CometScalarFunction")
+ .field("name", &self.name)
+ .field("signature", &self.signature)
+ .field("data_type", &self.data_type)
+ .finish()
+ }
+}
+
+impl CometScalarFunction {
+ fn new(
+ name: String,
+ signature: Signature,
+ data_type: DataType,
+ func: ScalarFunctionImplementation,
+ ) -> Self {
+ Self {
+ name,
+ signature,
+ data_type,
+ func,
+ }
+ }
+}
+
+impl ScalarUDFImpl for CometScalarFunction {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ self.name.as_str()
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _: &[DataType]) -> DataFusionResult<DataType> {
+ Ok(self.data_type.clone())
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) ->
DataFusionResult<ColumnarValue> {
+ (self.func)(args)
+ }
+}
diff --git a/native/core/src/execution/datafusion/expressions/mod.rs
b/native/core/src/execution/datafusion/expressions/mod.rs
index 3c0a5b26..ecd04956 100644
--- a/native/core/src/execution/datafusion/expressions/mod.rs
+++ b/native/core/src/execution/datafusion/expressions/mod.rs
@@ -20,13 +20,13 @@
pub mod bitwise_not;
pub mod checkoverflow;
mod normalize_nan;
-pub mod scalar_funcs;
pub use normalize_nan::NormalizeNaNAndZero;
use crate::errors::CometError;
pub mod avg;
pub mod avg_decimal;
pub mod bloom_filter_might_contain;
+pub mod comet_scalar_funcs;
pub mod correlation;
pub mod covariance;
pub mod create_named_struct;
@@ -39,7 +39,6 @@ pub mod sum_decimal;
pub mod unbound;
mod utils;
pub mod variance;
-pub mod xxhash64;
pub use datafusion_comet_spark_expr::{EvalMode, SparkError};
diff --git a/native/core/src/execution/datafusion/mod.rs
b/native/core/src/execution/datafusion/mod.rs
index e98e4ed4..6f81ee91 100644
--- a/native/core/src/execution/datafusion/mod.rs
+++ b/native/core/src/execution/datafusion/mod.rs
@@ -21,5 +21,4 @@ pub mod expressions;
mod operators;
pub mod planner;
pub mod shuffle_writer;
-pub mod spark_hash;
mod util;
diff --git a/native/core/src/execution/datafusion/planner.rs
b/native/core/src/execution/datafusion/planner.rs
index f3bc146d..c283cebb 100644
--- a/native/core/src/execution/datafusion/planner.rs
+++ b/native/core/src/execution/datafusion/planner.rs
@@ -76,7 +76,6 @@ use crate::{
correlation::Correlation,
covariance::Covariance,
negative,
- scalar_funcs::create_comet_physical_fun,
stats::StatsType,
stddev::Stddev,
strings::{Contains, EndsWith, Like, StartsWith,
StringSpaceExpr, SubstringExpr},
@@ -95,6 +94,7 @@ use crate::{
};
use super::expressions::{create_named_struct::CreateNamedStruct, EvalMode};
+use
crate::execution::datafusion::expressions::comet_scalar_funcs::create_comet_physical_fun;
use datafusion_comet_proto::{
spark_expression::{
self, agg_expr::ExprStruct as AggExprStruct, expr::ExprStruct,
literal::Value, AggExpr,
diff --git a/native/core/src/execution/datafusion/shuffle_writer.rs
b/native/core/src/execution/datafusion/shuffle_writer.rs
index 1a1149da..f7c97bea 100644
--- a/native/core/src/execution/datafusion/shuffle_writer.rs
+++ b/native/core/src/execution/datafusion/shuffle_writer.rs
@@ -62,8 +62,8 @@ use tokio::task;
use crate::{
common::bit::ceil,
errors::{CometError, CometResult},
- execution::datafusion::spark_hash::{create_murmur3_hashes, pmod},
};
+use datafusion_comet_spark_expr::spark_hash::create_murmur3_hashes;
/// The shuffle writer operator maps each input partition to M output
partitions based on a
/// partitioning scheme. No guarantees are made about the order of the
resulting partitions.
@@ -1413,6 +1413,14 @@ impl RecordBatchStream for EmptyStream {
}
}
+fn pmod(hash: u32, n: usize) -> usize {
+ let hash = hash as i32;
+ let n = n as i32;
+ let r = hash % n;
+ let result = if r < 0 { (r + n) % n } else { r };
+ result as usize
+}
+
#[cfg(test)]
mod test {
use super::*;
@@ -1477,4 +1485,14 @@ mod test {
let rt = Runtime::new().unwrap();
rt.block_on(collect(stream)).unwrap();
}
+
+ #[test]
+ fn test_pmod() {
+ let i: Vec<u32> = vec![0x99f0149d, 0x9c67b85d, 0xc8008529, 0xa05b5d7b,
0xcd1e64fb];
+ let result = i.into_iter().map(|i| pmod(i,
200)).collect::<Vec<usize>>();
+
+ // expected partition from Spark with n=200
+ let expected = vec![69, 5, 193, 171, 115];
+ assert_eq!(result, expected);
+ }
}
diff --git a/native/core/src/execution/datafusion/util/spark_bloom_filter.rs
b/native/core/src/execution/datafusion/util/spark_bloom_filter.rs
index 22957a14..00f71767 100644
--- a/native/core/src/execution/datafusion/util/spark_bloom_filter.rs
+++ b/native/core/src/execution/datafusion/util/spark_bloom_filter.rs
@@ -15,10 +15,9 @@
// specific language governing permissions and limitations
// under the License.
-use crate::execution::datafusion::{
- spark_hash::spark_compatible_murmur3_hash,
util::spark_bit_array::SparkBitArray,
-};
+use crate::execution::datafusion::util::spark_bit_array::SparkBitArray;
use arrow_array::{ArrowNativeTypeOp, BooleanArray, Int64Array};
+use datafusion_comet_spark_expr::spark_hash::spark_compatible_murmur3_hash;
const SPARK_BLOOM_FILTER_VERSION_1: i32 = 1;
diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml
index aa4fcfc5..a535a2b8 100644
--- a/native/spark-expr/Cargo.toml
+++ b/native/spark-expr/Cargo.toml
@@ -41,10 +41,13 @@ chrono-tz = { workspace = true }
num = { workspace = true }
regex = { workspace = true }
thiserror = { workspace = true }
+unicode-segmentation = "1.11.0"
[dev-dependencies]
+arrow-data = {workspace = true}
criterion = "0.5.1"
-rand = "0.8.5"
+rand = { workspace = true}
+twox-hash = "1.6.3"
[lib]
name = "datafusion_comet_spark_expr"
@@ -60,4 +63,4 @@ harness = false
[[bench]]
name = "conditional"
-harness = false
\ No newline at end of file
+harness = false
diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs
index 22628978..14ab080b 100644
--- a/native/spark-expr/src/lib.rs
+++ b/native/spark-expr/src/lib.rs
@@ -20,9 +20,12 @@ mod error;
mod if_expr;
mod kernels;
+pub mod scalar_funcs;
+pub mod spark_hash;
mod temporal;
pub mod timezone;
pub mod utils;
+mod xxhash64;
pub use cast::{spark_cast, Cast};
pub use error::{SparkError, SparkResult};
diff --git a/native/core/src/execution/datafusion/expressions/scalar_funcs.rs
b/native/spark-expr/src/scalar_funcs.rs
similarity index 79%
rename from native/core/src/execution/datafusion/expressions/scalar_funcs.rs
rename to native/spark-expr/src/scalar_funcs.rs
index d47172df..c50b98ba 100644
--- a/native/core/src/execution/datafusion/expressions/scalar_funcs.rs
+++ b/native/spark-expr/src/scalar_funcs.rs
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-use std::{any::Any, cmp::min, fmt::Debug, sync::Arc};
+use std::{cmp::min, sync::Arc};
use arrow::{
array::{
@@ -26,17 +26,11 @@ use arrow::{
};
use arrow_array::{Array, ArrowNativeTypeOp, BooleanArray, Decimal128Array};
use arrow_schema::DataType;
-use datafusion::{
- execution::FunctionRegistry,
- functions::math::round::round,
- logical_expr::{ScalarFunctionImplementation, ScalarUDFImpl, Signature,
Volatility},
- physical_plan::ColumnarValue,
-};
+use datafusion::{functions::math::round::round, physical_plan::ColumnarValue};
use datafusion_common::{
cast::as_generic_string_array, exec_err, internal_err, DataFusionError,
Result as DataFusionResult, ScalarValue,
};
-use datafusion_expr::ScalarUDF;
use num::{
integer::{div_ceil, div_floor},
BigInt, Signed, ToPrimitive,
@@ -44,113 +38,18 @@ use num::{
use unicode_segmentation::UnicodeSegmentation;
mod unhex;
-use unhex::spark_unhex;
+pub use unhex::spark_unhex;
mod hex;
-use hex::spark_hex;
+pub use hex::spark_hex;
mod chr;
-use chr::spark_chr;
+pub use chr::SparkChrFunc;
pub mod hash_expressions;
// exposed for benchmark only
-use hash_expressions::wrap_digest_result_as_hex_string;
pub use hash_expressions::{spark_murmur3_hash, spark_xxhash64};
-macro_rules! make_comet_scalar_udf {
- ($name:expr, $func:ident, $data_type:ident) => {{
- let scalar_func = CometScalarFunction::new(
- $name.to_string(),
- Signature::variadic_any(Volatility::Immutable),
- $data_type.clone(),
- Arc::new(move |args| $func(args, &$data_type)),
- );
- Ok(Arc::new(ScalarUDF::new_from_impl(scalar_func)))
- }};
- ($name:expr, $func:expr, without $data_type:ident) => {{
- let scalar_func = CometScalarFunction::new(
- $name.to_string(),
- Signature::variadic_any(Volatility::Immutable),
- $data_type,
- $func,
- );
- Ok(Arc::new(ScalarUDF::new_from_impl(scalar_func)))
- }};
-}
-
-/// Create a physical scalar function.
-pub fn create_comet_physical_fun(
- fun_name: &str,
- data_type: DataType,
- registry: &dyn FunctionRegistry,
-) -> Result<Arc<ScalarUDF>, DataFusionError> {
- let sha2_functions = ["sha224", "sha256", "sha384", "sha512"];
- match fun_name {
- "ceil" => {
- make_comet_scalar_udf!("ceil", spark_ceil, data_type)
- }
- "floor" => {
- make_comet_scalar_udf!("floor", spark_floor, data_type)
- }
- "rpad" => {
- let func = Arc::new(spark_rpad);
- make_comet_scalar_udf!("rpad", func, without data_type)
- }
- "round" => {
- make_comet_scalar_udf!("round", spark_round, data_type)
- }
- "unscaled_value" => {
- let func = Arc::new(spark_unscaled_value);
- make_comet_scalar_udf!("unscaled_value", func, without data_type)
- }
- "make_decimal" => {
- make_comet_scalar_udf!("make_decimal", spark_make_decimal,
data_type)
- }
- "hex" => {
- let func = Arc::new(spark_hex);
- make_comet_scalar_udf!("hex", func, without data_type)
- }
- "unhex" => {
- let func = Arc::new(spark_unhex);
- make_comet_scalar_udf!("unhex", func, without data_type)
- }
- "decimal_div" => {
- make_comet_scalar_udf!("decimal_div", spark_decimal_div, data_type)
- }
- "murmur3_hash" => {
- let func = Arc::new(spark_murmur3_hash);
- make_comet_scalar_udf!("murmur3_hash", func, without data_type)
- }
- "xxhash64" => {
- let func = Arc::new(spark_xxhash64);
- make_comet_scalar_udf!("xxhash64", func, without data_type)
- }
- "chr" => {
- let func = Arc::new(spark_chr);
- make_comet_scalar_udf!("chr", func, without data_type)
- }
- "isnan" => {
- let func = Arc::new(spark_isnan);
- make_comet_scalar_udf!("isnan", func, without data_type)
- }
- sha if sha2_functions.contains(&sha) => {
- // Spark requires hex string as the result of sha2 functions, we
have to wrap the
- // result of digest functions as hex string
- let func = registry.udf(sha)?;
- let wrapped_func = Arc::new(move |args: &[ColumnarValue]| {
- wrap_digest_result_as_hex_string(args, func.fun())
- });
- let spark_func_name = "spark".to_owned() + sha;
- make_comet_scalar_udf!(spark_func_name, wrapped_func, without
data_type)
- }
- _ => registry.udf(fun_name).map_err(|e| {
- DataFusionError::Execution(format!(
- "Function {fun_name} not found in the registry: {e}",
- ))
- }),
- }
-}
-
#[inline]
fn get_precision_scale(data_type: &DataType) -> (u8, i8) {
let DataType::Decimal128(precision, scale) = data_type else {
@@ -176,61 +75,6 @@ macro_rules! downcast_compute_op {
}};
}
-struct CometScalarFunction {
- name: String,
- signature: Signature,
- data_type: DataType,
- func: ScalarFunctionImplementation,
-}
-
-impl Debug for CometScalarFunction {
- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- f.debug_struct("CometScalarFunction")
- .field("name", &self.name)
- .field("signature", &self.signature)
- .field("data_type", &self.data_type)
- .finish()
- }
-}
-
-impl CometScalarFunction {
- fn new(
- name: String,
- signature: Signature,
- data_type: DataType,
- func: ScalarFunctionImplementation,
- ) -> Self {
- Self {
- name,
- signature,
- data_type,
- func,
- }
- }
-}
-
-impl ScalarUDFImpl for CometScalarFunction {
- fn as_any(&self) -> &dyn Any {
- self
- }
-
- fn name(&self) -> &str {
- self.name.as_str()
- }
-
- fn signature(&self) -> &Signature {
- &self.signature
- }
-
- fn return_type(&self, _: &[DataType]) -> DataFusionResult<DataType> {
- Ok(self.data_type.clone())
- }
-
- fn invoke(&self, args: &[ColumnarValue]) ->
DataFusionResult<ColumnarValue> {
- (self.func)(args)
- }
-}
-
/// `ceil` function that simulates Spark `ceil` expression
pub fn spark_ceil(
args: &[ColumnarValue],
@@ -333,6 +177,7 @@ pub fn spark_floor(
}
}
+/// Spark-compatible `UnscaledValue` expression (internal to Spark optimizer)
pub fn spark_unscaled_value(args: &[ColumnarValue]) ->
DataFusionResult<ColumnarValue> {
match &args[0] {
ColumnarValue::Scalar(v) => match v {
@@ -352,6 +197,7 @@ pub fn spark_unscaled_value(args: &[ColumnarValue]) ->
DataFusionResult<Columnar
}
}
+/// Spark-compatible `MakeDecimal` expression (internal to Spark optimizer)
pub fn spark_make_decimal(
args: &[ColumnarValue],
data_type: &DataType,
@@ -489,7 +335,7 @@ macro_rules! round_integer_scalar {
}
/// `round` function that simulates Spark `round` expression
-fn spark_round(
+pub fn spark_round(
args: &[ColumnarValue],
data_type: &DataType,
) -> Result<ColumnarValue, DataFusionError> {
@@ -541,7 +387,7 @@ fn spark_round(
}
/// Similar to DataFusion `rpad`, but not to truncate when the string is
already longer than length
-fn spark_rpad(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
+pub fn spark_rpad(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
match args {
[ColumnarValue::Array(array),
ColumnarValue::Scalar(ScalarValue::Int32(Some(length)))] => {
match args[0].data_type() {
@@ -594,7 +440,7 @@ fn spark_rpad_internal<T: OffsetSizeTrait>(
// get enough scale that matches with Spark behavior, it requires to widen s1
to s2 + s3 + 1. Since
// both s2 and s3 are 38 at max., s1 is 77 at max. DataFusion division cannot
handle such scale >
// Decimal256Type::MAX_SCALE. Therefore, we need to implement this decimal
division using BigInt.
-fn spark_decimal_div(
+pub fn spark_decimal_div(
args: &[ColumnarValue],
data_type: &DataType,
) -> Result<ColumnarValue, DataFusionError> {
@@ -639,7 +485,8 @@ fn spark_decimal_div(
Ok(ColumnarValue::Array(Arc::new(result)))
}
-fn spark_isnan(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
+/// Spark-compatible `isnan` expression
+pub fn spark_isnan(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
fn set_nulls_to_false(is_nan: BooleanArray) -> ColumnarValue {
match is_nan.nulls() {
Some(nulls) => {
diff --git
a/native/core/src/execution/datafusion/expressions/scalar_funcs/chr.rs
b/native/spark-expr/src/scalar_funcs/chr.rs
similarity index 87%
rename from native/core/src/execution/datafusion/expressions/scalar_funcs/chr.rs
rename to native/spark-expr/src/scalar_funcs/chr.rs
index 0421b491..5de59f9f 100644
--- a/native/core/src/execution/datafusion/expressions/scalar_funcs/chr.rs
+++ b/native/spark-expr/src/scalar_funcs/chr.rs
@@ -26,17 +26,9 @@ use arrow::{
};
use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature,
Volatility};
-use datafusion_common::{cast::as_int64_array, exec_err, DataFusionError,
Result, ScalarValue};
+use datafusion_common::{cast::as_int64_array, exec_err, Result, ScalarValue};
-/// Returns the ASCII character having the binary equivalent to the input
expression.
-/// E.g., chr(65) = 'A'.
-/// Compatible with Apache Spark's Chr function
-pub fn spark_chr(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
- let chr_func = ChrFunc::default();
- chr_func.invoke(args)
-}
-
-pub fn chr(args: &[ArrayRef]) -> Result<ArrayRef> {
+fn chr(args: &[ArrayRef]) -> Result<ArrayRef> {
let integer_array = as_int64_array(&args[0])?;
// first map is the iterator, second is for the `Option<_>`
@@ -62,18 +54,19 @@ pub fn chr(args: &[ArrayRef]) -> Result<ArrayRef> {
Ok(Arc::new(result) as ArrayRef)
}
+/// Spark-compatible `chr` expression
#[derive(Debug)]
-pub struct ChrFunc {
+pub struct SparkChrFunc {
signature: Signature,
}
-impl Default for ChrFunc {
+impl Default for SparkChrFunc {
fn default() -> Self {
Self::new()
}
}
-impl ChrFunc {
+impl SparkChrFunc {
pub fn new() -> Self {
Self {
signature: Signature::uniform(1, vec![Int64],
Volatility::Immutable),
@@ -81,7 +74,7 @@ impl ChrFunc {
}
}
-impl ScalarUDFImpl for ChrFunc {
+impl ScalarUDFImpl for SparkChrFunc {
fn as_any(&self) -> &dyn Any {
self
}
@@ -99,11 +92,14 @@ impl ScalarUDFImpl for ChrFunc {
}
fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
- handle_chr_fn(args)
+ spark_chr(args)
}
}
-fn handle_chr_fn(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+/// Returns the ASCII character having the binary equivalent to the input
expression.
+/// E.g., chr(65) = 'A'.
+/// Compatible with Apache Spark's Chr function
+fn spark_chr(args: &[ColumnarValue]) -> Result<ColumnarValue> {
let array = args[0].clone();
match array {
ColumnarValue::Array(array) => {
diff --git
a/native/core/src/execution/datafusion/expressions/scalar_funcs/hash_expressions.rs
b/native/spark-expr/src/scalar_funcs/hash_expressions.rs
similarity index 79%
rename from
native/core/src/execution/datafusion/expressions/scalar_funcs/hash_expressions.rs
rename to native/spark-expr/src/scalar_funcs/hash_expressions.rs
index 67d72816..1a403b9e 100644
---
a/native/core/src/execution/datafusion/expressions/scalar_funcs/hash_expressions.rs
+++ b/native/spark-expr/src/scalar_funcs/hash_expressions.rs
@@ -15,15 +15,17 @@
// specific language governing permissions and limitations
// under the License.
-use crate::execution::datafusion::expressions::scalar_funcs::hex::hex_strings;
-use crate::execution::datafusion::spark_hash::{create_murmur3_hashes,
create_xxhash64_hashes};
+use crate::scalar_funcs::hex::hex_strings;
+use crate::spark_hash::{create_murmur3_hashes, create_xxhash64_hashes};
+
use arrow_array::{ArrayRef, Int32Array, Int64Array, StringArray};
+use datafusion::functions::crypto::{sha224, sha256, sha384, sha512};
use datafusion_common::cast::as_binary_array;
use datafusion_common::{exec_err, internal_err, DataFusionError, ScalarValue};
use datafusion_expr::{ColumnarValue, ScalarFunctionImplementation};
use std::sync::Arc;
-/// Spark compatible murmur3 hash in vectorized execution fashion
+/// Spark compatible murmur3 hash (just `hash` in Spark) in vectorized
execution fashion
pub fn spark_murmur3_hash(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
let length = args.len();
let seed = &args[length - 1];
@@ -111,7 +113,29 @@ pub fn spark_xxhash64(args: &[ColumnarValue]) ->
Result<ColumnarValue, DataFusio
}
}
-pub(super) fn wrap_digest_result_as_hex_string(
+/// `sha224` function that simulates Spark's `sha2` expression with bit width
224
+pub fn spark_sha224(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
+ wrap_digest_result_as_hex_string(args, sha224().fun())
+}
+
+/// `sha256` function that simulates Spark's `sha2` expression with bit width
0 or 256
+pub fn spark_sha256(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
+ wrap_digest_result_as_hex_string(args, sha256().fun())
+}
+
+/// `sha384` function that simulates Spark's `sha2` expression with bit width
384
+pub fn spark_sha384(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
+ wrap_digest_result_as_hex_string(args, sha384().fun())
+}
+
+/// `sha512` function that simulates Spark's `sha2` expression with bit width
512
+pub fn spark_sha512(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
+ wrap_digest_result_as_hex_string(args, sha512().fun())
+}
+
+// Spark requires hex string as the result of sha2 functions, we have to wrap
the
+// result of digest functions as hex string
+fn wrap_digest_result_as_hex_string(
args: &[ColumnarValue],
digest: ScalarFunctionImplementation,
) -> Result<ColumnarValue, DataFusionError> {
diff --git
a/native/core/src/execution/datafusion/expressions/scalar_funcs/hex.rs
b/native/spark-expr/src/scalar_funcs/hex.rs
similarity index 98%
rename from native/core/src/execution/datafusion/expressions/scalar_funcs/hex.rs
rename to native/spark-expr/src/scalar_funcs/hex.rs
index e6059818..e572ba5e 100644
--- a/native/core/src/execution/datafusion/expressions/scalar_funcs/hex.rs
+++ b/native/spark-expr/src/scalar_funcs/hex.rs
@@ -62,7 +62,8 @@ fn hex_bytes<T: AsRef<[u8]>>(bytes: T) -> Result<String,
std::fmt::Error> {
Ok(hex_string)
}
-pub(super) fn spark_hex(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
+/// Spark-compatible `hex` function
+pub fn spark_hex(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
if args.len() != 1 {
return Err(DataFusionError::Internal(
"hex expects exactly one argument".to_string(),
diff --git
a/native/core/src/execution/datafusion/expressions/scalar_funcs/unhex.rs
b/native/spark-expr/src/scalar_funcs/unhex.rs
similarity index 98%
rename from
native/core/src/execution/datafusion/expressions/scalar_funcs/unhex.rs
rename to native/spark-expr/src/scalar_funcs/unhex.rs
index 38d5c047..9996392b 100644
--- a/native/core/src/execution/datafusion/expressions/scalar_funcs/unhex.rs
+++ b/native/spark-expr/src/scalar_funcs/unhex.rs
@@ -109,7 +109,8 @@ fn spark_unhex_inner<T: OffsetSizeTrait>(
}
}
-pub(super) fn spark_unhex(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
+/// Spark-compatible `unhex` expression
+pub fn spark_unhex(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
if args.len() > 2 {
return exec_err!("unhex takes at most 2 arguments, but got: {}",
args.len());
}
diff --git a/native/core/src/execution/datafusion/spark_hash.rs
b/native/spark-expr/src/spark_hash.rs
similarity index 96%
rename from native/core/src/execution/datafusion/spark_hash.rs
rename to native/spark-expr/src/spark_hash.rs
index e429f543..66a103a2 100644
--- a/native/core/src/execution/datafusion/spark_hash.rs
+++ b/native/spark-expr/src/spark_hash.rs
@@ -34,10 +34,11 @@ use datafusion::{
error::{DataFusionError, Result},
};
-use
crate::execution::datafusion::expressions::xxhash64::spark_compatible_xxhash64;
+use crate::xxhash64::spark_compatible_xxhash64;
+/// Spark-compatible murmur3 hash function
#[inline]
-pub(crate) fn spark_compatible_murmur3_hash<T: AsRef<[u8]>>(data: T, seed:
u32) -> u32 {
+pub fn spark_compatible_murmur3_hash<T: AsRef<[u8]>>(data: T, seed: u32) ->
u32 {
#[inline]
fn mix_k1(mut k1: i32) -> i32 {
k1 = k1.mul_wrapping(0xcc9e2d51u32 as i32);
@@ -458,7 +459,7 @@ macro_rules! create_hashes_internal {
///
/// The number of rows to hash is determined by `hashes_buffer.len()`.
/// `hashes_buffer` should be pre-sized appropriately
-pub(crate) fn create_murmur3_hashes<'a>(
+pub fn create_murmur3_hashes<'a>(
arrays: &[ArrayRef],
hashes_buffer: &'a mut [u32],
) -> Result<&'a mut [u32]> {
@@ -489,22 +490,12 @@ pub fn create_xxhash64_hashes<'a>(
Ok(hashes_buffer)
}
-pub(crate) fn pmod(hash: u32, n: usize) -> usize {
- let hash = hash as i32;
- let n = n as i32;
- let r = hash % n;
- let result = if r < 0 { (r + n) % n } else { r };
- result as usize
-}
-
#[cfg(test)]
mod tests {
use arrow::array::{Float32Array, Float64Array};
use std::sync::Arc;
- use crate::execution::datafusion::spark_hash::{
- create_murmur3_hashes, create_xxhash64_hashes, pmod,
- };
+ use super::{create_murmur3_hashes, create_xxhash64_hashes};
use datafusion::arrow::array::{ArrayRef, Int32Array, Int64Array,
Int8Array, StringArray};
macro_rules! test_hashes_internal {
@@ -714,14 +705,4 @@ mod tests {
],
)
}
-
- #[test]
- fn test_pmod() {
- let i: Vec<u32> = vec![0x99f0149d, 0x9c67b85d, 0xc8008529, 0xa05b5d7b,
0xcd1e64fb];
- let result = i.into_iter().map(|i| pmod(i,
200)).collect::<Vec<usize>>();
-
- // expected partition from Spark with n=200
- let expected = vec![69, 5, 193, 171, 115];
- assert_eq!(result, expected);
- }
}
diff --git a/native/core/src/execution/datafusion/expressions/xxhash64.rs
b/native/spark-expr/src/xxhash64.rs
similarity index 100%
rename from native/core/src/execution/datafusion/expressions/xxhash64.rs
rename to native/spark-expr/src/xxhash64.rs
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
