Skip to content

Commit 25db829

Browse files
FlyangzFlyangz
committed
fix ext func caching
1 parent a603cc0 commit 25db829

1 file changed

Lines changed: 43 additions & 37 deletions

File tree

  • native-engine/datafusion-ext-functions/src

‎native-engine/datafusion-ext-functions/src/lib.rs‎

Lines changed: 43 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
// See the License for the specific language governing permissions and
1414
// limitations under the License.
1515

16-
use std::sync::Arc;
16+
use std::sync::{Arc, OnceLock};
1717

1818
use datafusion::{common::Result, logical_expr::ScalarFunctionImplementation};
1919
use datafusion_ext_commons::df_unimplemented_err;
@@ -39,51 +39,57 @@ pub fn create_auron_ext_function(
3939
name: &str,
4040
spark_partition_id: usize,
4141
) -> Result<ScalarFunctionImplementation> {
42+
macro_rules! cache {
43+
($func:path) => {{
44+
static CELL: OnceLock<ScalarFunctionImplementation> = OnceLock::new();
45+
CELL.get_or_init(|| Arc::new($func)).clone()
46+
}};
47+
}
4248
// auron ext functions, if used for spark should be start with 'Spark_',
4349
// if used for flink should be start with 'Flink_',
4450
// same to other engines.
4551
Ok(match name {
4652
"Placeholder" => Arc::new(|_| panic!("placeholder() should never be called")),
47-
"Spark_NullIf" => Arc::new(spark_null_if::spark_null_if),
48-
"Spark_NullIfZero" => Arc::new(spark_null_if::spark_null_if_zero),
49-
"Spark_UnscaledValue" => Arc::new(spark_unscaled_value::spark_unscaled_value),
50-
"Spark_MakeDecimal" => Arc::new(spark_make_decimal::spark_make_decimal),
51-
"Spark_CheckOverflow" => Arc::new(spark_check_overflow::spark_check_overflow),
52-
"Spark_Murmur3Hash" => Arc::new(spark_hash::spark_murmur3_hash),
53-
"Spark_XxHash64" => Arc::new(spark_hash::spark_xxhash64),
54-
"Spark_Sha224" => Arc::new(spark_crypto::spark_sha224),
55-
"Spark_Sha256" => Arc::new(spark_crypto::spark_sha256),
56-
"Spark_Sha384" => Arc::new(spark_crypto::spark_sha384),
57-
"Spark_Sha512" => Arc::new(spark_crypto::spark_sha512),
58-
"Spark_MD5" => Arc::new(spark_crypto::spark_md5),
59-
"Spark_GetJsonObject" => Arc::new(spark_get_json_object::spark_get_json_object),
53+
"Spark_NullIf" => cache!(spark_null_if::spark_null_if),
54+
"Spark_NullIfZero" => cache!(spark_null_if::spark_null_if_zero),
55+
"Spark_UnscaledValue" => cache!(spark_unscaled_value::spark_unscaled_value),
56+
"Spark_MakeDecimal" => cache!(spark_make_decimal::spark_make_decimal),
57+
"Spark_CheckOverflow" => cache!(spark_check_overflow::spark_check_overflow),
58+
"Spark_Murmur3Hash" => cache!(spark_hash::spark_murmur3_hash),
59+
"Spark_XxHash64" => cache!(spark_hash::spark_xxhash64),
60+
"Spark_Sha224" => cache!(spark_crypto::spark_sha224),
61+
"Spark_Sha256" => cache!(spark_crypto::spark_sha256),
62+
"Spark_Sha384" => cache!(spark_crypto::spark_sha384),
63+
"Spark_Sha512" => cache!(spark_crypto::spark_sha512),
64+
"Spark_MD5" => cache!(spark_crypto::spark_md5),
65+
"Spark_GetJsonObject" => cache!(spark_get_json_object::spark_get_json_object),
6066
"Spark_GetParsedJsonObject" => {
61-
Arc::new(spark_get_json_object::spark_get_parsed_json_object)
67+
cache!(spark_get_json_object::spark_get_parsed_json_object)
6268
}
63-
"Spark_ParseJson" => Arc::new(spark_get_json_object::spark_parse_json),
64-
"Spark_MakeArray" => Arc::new(spark_make_array::array),
65-
"Spark_StringSpace" => Arc::new(spark_strings::string_space),
66-
"Spark_StringRepeat" => Arc::new(spark_strings::string_repeat),
67-
"Spark_StringSplit" => Arc::new(spark_strings::string_split),
68-
"Spark_StringConcat" => Arc::new(spark_strings::string_concat),
69-
"Spark_StringConcatWs" => Arc::new(spark_strings::string_concat_ws),
70-
"Spark_StringLower" => Arc::new(spark_strings::string_lower),
71-
"Spark_StringUpper" => Arc::new(spark_strings::string_upper),
72-
"Spark_InitCap" => Arc::new(spark_initcap::string_initcap),
73-
"Spark_Year" => Arc::new(spark_dates::spark_year),
74-
"Spark_Month" => Arc::new(spark_dates::spark_month),
75-
"Spark_Day" => Arc::new(spark_dates::spark_day),
76-
"Spark_Quarter" => Arc::new(spark_dates::spark_quarter),
77-
"Spark_Hour" => Arc::new(spark_dates::spark_hour),
78-
"Spark_Minute" => Arc::new(spark_dates::spark_minute),
79-
"Spark_Second" => Arc::new(spark_dates::spark_second),
80-
"Spark_BrickhouseArrayUnion" => Arc::new(brickhouse::array_union::array_union),
81-
"Spark_Round" => Arc::new(spark_round::spark_round),
82-
"Spark_BRound" => Arc::new(spark_bround::spark_bround),
69+
"Spark_ParseJson" => cache!(spark_get_json_object::spark_parse_json),
70+
"Spark_MakeArray" => cache!(spark_make_array::array),
71+
"Spark_StringSpace" => cache!(spark_strings::string_space),
72+
"Spark_StringRepeat" => cache!(spark_strings::string_repeat),
73+
"Spark_StringSplit" => cache!(spark_strings::string_split),
74+
"Spark_StringConcat" => cache!(spark_strings::string_concat),
75+
"Spark_StringConcatWs" => cache!(spark_strings::string_concat_ws),
76+
"Spark_StringLower" => cache!(spark_strings::string_lower),
77+
"Spark_StringUpper" => cache!(spark_strings::string_upper),
78+
"Spark_InitCap" => cache!(spark_initcap::string_initcap),
79+
"Spark_Year" => cache!(spark_dates::spark_year),
80+
"Spark_Month" => cache!(spark_dates::spark_month),
81+
"Spark_Day" => cache!(spark_dates::spark_day),
82+
"Spark_Quarter" => cache!(spark_dates::spark_quarter),
83+
"Spark_Hour" => cache!(spark_dates::spark_hour),
84+
"Spark_Minute" => cache!(spark_dates::spark_minute),
85+
"Spark_Second" => cache!(spark_dates::spark_second),
86+
"Spark_BrickhouseArrayUnion" => cache!(brickhouse::array_union::array_union),
87+
"Spark_Round" => cache!(spark_round::spark_round),
88+
"Spark_BRound" => cache!(spark_bround::spark_bround),
8389
"Spark_NormalizeNanAndZero" => {
84-
Arc::new(spark_normalize_nan_and_zero::spark_normalize_nan_and_zero)
90+
cache!(spark_normalize_nan_and_zero::spark_normalize_nan_and_zero)
8591
}
86-
"Spark_IsNaN" => Arc::new(spark_isnan::spark_isnan),
92+
"Spark_IsNaN" => cache!(spark_isnan::spark_isnan),
8793
_ => df_unimplemented_err!("spark ext function not implemented: {name}")?,
8894
})
8995
}

0 commit comments

Comments
 (0)