feat: support factorial, pmod, and rint expressions

native/core/src/execution/jni_api.rs

@@ -52,7 +52,10 @@ use datafusion_spark::function::hash::sha1::SparkSha1;
 use datafusion_spark::function::hash::sha2::SparkSha2;
 use datafusion_spark::function::map::map_from_entries::MapFromEntries;
 use datafusion_spark::function::math::expm1::SparkExpm1;
+use datafusion_spark::function::math::factorial::SparkFactorial;
 use datafusion_spark::function::math::hex::SparkHex;
+use datafusion_spark::function::math::modulus::SparkPmod;
+use datafusion_spark::function::math::rint::SparkRint;
 use datafusion_spark::function::math::width_bucket::SparkWidthBucket;
 use datafusion_spark::function::string::char::CharFunc;
 use datafusion_spark::function::string::concat::SparkConcat;
@@ -404,6 +407,9 @@ fn register_datafusion_spark_function(session_ctx: &SessionContext) {
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkCrc32::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSpace::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitCount::default()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkFactorial::default()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkPmod::default()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkRint::default()));
 }
 
 /// Prepares arrow arrays for output.

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -96,6 +96,7 @@ object QueryPlanSerde extends Logging with CometExprShim {
     classOf[Divide] -> CometDivide,
     classOf[Exp] -> CometScalarFunction("exp"),
     classOf[Expm1] -> CometScalarFunction("expm1"),
+    classOf[Factorial] -> CometScalarFunction("factorial"),
     classOf[Floor] -> CometFloor,
     classOf[Hex] -> CometHex,
     classOf[IntegralDivide] -> CometIntegralDivide,
@@ -104,10 +105,12 @@ object QueryPlanSerde extends Logging with CometExprShim {
     classOf[Log2] -> CometLog2,
     classOf[Log10] -> CometLog10,
     classOf[Multiply] -> CometMultiply,
+    classOf[Pmod] -> CometScalarFunction("pmod"),
     classOf[Pow] -> CometScalarFunction("pow"),
     classOf[Rand] -> CometRand,
     classOf[Randn] -> CometRandn,
     classOf[Remainder] -> CometRemainder,
+    classOf[Rint] -> CometScalarFunction("rint"),
     classOf[Round] -> CometRound,
     classOf[Signum] -> CometScalarFunction("signum"),
     classOf[Sin] -> CometScalarFunction("sin"),

spark/src/test/resources/sql-tests/expressions/math/factorial.sql

@@ -0,0 +1,32 @@
+-- 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.
+
+-- ConfigMatrix: parquet.enable.dictionary=false,true
+
+statement
+CREATE TABLE test_factorial(i int) USING parquet
+
+statement
+INSERT INTO test_factorial VALUES (0), (1), (5), (10), (20), (21), (-1), (NULL)
+
+-- column input
+query
+SELECT factorial(i) FROM test_factorial
+
+-- literal arguments
+query
+SELECT factorial(0), factorial(1), factorial(5), factorial(20), factorial(21), factorial(-1), factorial(NULL)

spark/src/test/resources/sql-tests/expressions/math/pmod.sql

@@ -0,0 +1,46 @@
+-- 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.
+
+-- ConfigMatrix: parquet.enable.dictionary=false,true
+
+statement
+CREATE TABLE test_pmod_int(a int, b int) USING parquet
+
+statement
+INSERT INTO test_pmod_int VALUES (10, 3), (7, -2), (-7, 2), (-7, -2), (0, 5), (10, 0), (NULL, 3), (10, NULL)
+
+-- integer column input
+query
+SELECT pmod(a, b) FROM test_pmod_int
+
+-- integer literal arguments
+query
+SELECT pmod(10, 3), pmod(7, -2), pmod(-7, 2), pmod(-7, -2), pmod(10, 0), pmod(NULL, 3)
+
+statement
+CREATE TABLE test_pmod_double(a double, b double) USING parquet
+
+statement
+INSERT INTO test_pmod_double VALUES (10.5, 3.0), (-7.5, 2.0), (10.0, 0.0), (NULL, 1.0)
+
+-- floating-point column input
+query
+SELECT pmod(a, b) FROM test_pmod_double
+
+-- floating-point literal arguments
+query
+SELECT pmod(10.5, 3.0), pmod(-7.5, 2.0), pmod(10.0, 0.0)

spark/src/test/resources/sql-tests/expressions/math/rint.sql

@@ -0,0 +1,32 @@
+-- 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.
+
+-- ConfigMatrix: parquet.enable.dictionary=false,true
+
+statement
+CREATE TABLE test_rint(d double) USING parquet
+
+statement
+INSERT INTO test_rint VALUES (1.5), (2.5), (-1.5), (0.0), (3.7), (NULL), (cast('NaN' as double)), (cast('Infinity' as double))
+
+-- column input
+query
+SELECT rint(d) FROM test_rint
+
+-- literal arguments
+query
+SELECT rint(1.5), rint(2.5), rint(-1.5), rint(0.0), rint(3.7), rint(NULL)