chore: wire time functions + slice, shift

comphead · comphead · commit c575989da7db · 2026-05-20T11:45:01.000-07:00
diff --git a/docs/source/contributor-guide/spark_expressions_support.md b/docs/source/contributor-guide/spark_expressions_support.md
@@ -151,7 +151,7 @@
 - [x] get
 - [ ] sequence
 - [ ] shuffle
-- [ ] slice
+- [x] slice
 - [x] sort_array
 
 ### bitwise_funcs
@@ -165,7 +165,7 @@
 - [x] bit_get
 - [x] getbit
 - [x] shiftright
-- [ ] shiftrightunsigned
+- [x] shiftrightunsigned
 - [x] `|`
 - [x] `~`
 
@@ -214,7 +214,7 @@
 
 ### datetime_funcs
 
-- [ ] add_months
+- [x] add_months
 - [ ] convert_timezone
 - [ ] curdate
 - [ ] current_date
@@ -286,9 +286,9 @@
 - [ ] try_to_time
 - [ ] try_to_timestamp
 - [x] unix_date
-- [ ] unix_micros
-- [ ] unix_millis
-- [ ] unix_seconds
+- [x] unix_micros
+- [x] unix_millis
+- [x] unix_seconds
 - [x] unix_timestamp
 - [x] weekday
 - [x] weekofyear
diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md
@@ -101,6 +101,7 @@ of expressions that be disabled.
 
 | Expression       | SQL                          |
 | ---------------- | ---------------------------- |
+| AddMonths        | `add_months`                 |
 | CurrentTimeZone  | `current_timezone`           |
 | DateAdd          | `date_add`                   |
 | DateDiff         | `datediff`                   |
@@ -122,6 +123,9 @@ of expressions that be disabled.
 | TruncDate        | `trunc`                      |
 | TruncTimestamp   | `date_trunc`                 |
 | UnixDate         | `unix_date`                  |
+| UnixMicros       | `unix_micros`                |
+| UnixMillis       | `unix_millis`                |
+| UnixSeconds      | `unix_seconds`               |
 | UnixTimestamp    | `unix_timestamp`             |
 | Year             | `year`                       |
 | Month            | `month`                      |
@@ -201,16 +205,17 @@ of expressions that be disabled.
 
 ## Bitwise Expressions
 
-| Expression   | SQL  |
-| ------------ | ---- |
-| BitwiseAnd   | `&`  |
-| BitwiseCount |      |
-| BitwiseGet   |      |
-| BitwiseOr    | `\|` |
-| BitwiseNot   | `~`  |
-| BitwiseXor   | `^`  |
-| ShiftLeft    | `<<` |
-| ShiftRight   | `>>` |
+| Expression         | SQL   |
+| ------------------ | ----- |
+| BitwiseAnd         | `&`   |
+| BitwiseCount       |       |
+| BitwiseGet         |       |
+| BitwiseOr          | `\|`  |
+| BitwiseNot         | `~`   |
+| BitwiseXor         | `^`   |
+| ShiftLeft          | `<<`  |
+| ShiftRight         | `>>`  |
+| ShiftRightUnsigned | `>>>` |
 
 ## Aggregate Expressions
 
@@ -281,6 +286,7 @@ Comet supports using the following aggregate functions within window contexts wi
 | Flatten        |
 | GetArrayItem   |
 | Size           |
+| Slice          |
 | SortArray      |
 
 ## Map Expressions
diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs
@@ -43,15 +43,19 @@ use datafusion::{
 };
 use datafusion_comet_proto::spark_operator::Operator;
 use datafusion_spark::function::array::array_contains::SparkArrayContains;
+use datafusion_spark::function::array::slice::SparkSlice;
 use datafusion_spark::function::bitwise::bit_count::SparkBitCount;
 use datafusion_spark::function::bitwise::bit_get::SparkBitGet;
+use datafusion_spark::function::bitwise::bit_shift::SparkBitShift;
 use datafusion_spark::function::bitwise::bitwise_not::SparkBitwiseNot;
+use datafusion_spark::function::datetime::add_months::SparkAddMonths;
 use datafusion_spark::function::datetime::date_add::SparkDateAdd;
 use datafusion_spark::function::datetime::date_sub::SparkDateSub;
 use datafusion_spark::function::datetime::from_utc_timestamp::SparkFromUtcTimestamp;
 use datafusion_spark::function::datetime::last_day::SparkLastDay;
 use datafusion_spark::function::datetime::next_day::SparkNextDay;
 use datafusion_spark::function::datetime::to_utc_timestamp::SparkToUtcTimestamp;
+use datafusion_spark::function::datetime::unix::SparkUnixTimestamp;
 use datafusion_spark::function::hash::crc32::SparkCrc32;
 use datafusion_spark::function::hash::sha1::SparkSha1;
 use datafusion_spark::function::hash::sha2::SparkSha2;
@@ -599,6 +603,12 @@ fn register_datafusion_spark_function(session_ctx: &SessionContext) {
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkTryUrlDecode::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkCsc::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkFactorial::default()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkAddMonths::default()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSlice::default()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitShift::right_unsigned()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkUnixTimestamp::microseconds()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkUnixTimestamp::milliseconds()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkUnixTimestamp::seconds()));
 }
 
 /// Prepares arrow arrays for output.
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -63,6 +63,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
     classOf[ArrayPosition] -> CometArrayPosition,
     classOf[ArrayRemove] -> CometArrayRemove,
     classOf[ArrayRepeat] -> CometArrayRepeat,
+    classOf[Slice] -> CometScalarFunction("slice"),
     classOf[SortArray] -> CometSortArray,
     classOf[ArraysOverlap] -> CometArraysOverlap,
     classOf[ArrayUnion] -> CometArrayUnion,
@@ -213,10 +214,12 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
     classOf[BitwiseNot] -> CometBitwiseNot,
     classOf[BitwiseXor] -> CometBitwiseXor,
     classOf[ShiftLeft] -> CometShiftLeft,
-    classOf[ShiftRight] -> CometShiftRight)
+    classOf[ShiftRight] -> CometShiftRight,
+    classOf[ShiftRightUnsigned] -> CometScalarFunction("shiftrightunsigned"))
 
   private[comet] val temporalExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] =
     Map(
+      classOf[AddMonths] -> CometScalarFunction("add_months"),
       classOf[DateAdd] -> CometDateAdd,
       classOf[DateDiff] -> CometDateDiff,
       classOf[DateFormatClass] -> CometDateFormat,
@@ -225,6 +228,9 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
       classOf[Hours] -> CometHours,
       classOf[DateSub] -> CometDateSub,
       classOf[UnixDate] -> CometUnixDate,
+      classOf[UnixMicros] -> CometScalarFunction("unix_micros"),
+      classOf[UnixMillis] -> CometScalarFunction("unix_millis"),
+      classOf[UnixSeconds] -> CometScalarFunction("unix_seconds"),
       classOf[FromUnixTime] -> CometFromUnixTime,
       classOf[FromUTCTimestamp] -> CometFromUTCTimestamp,
       classOf[ToUTCTimestamp] -> CometToUTCTimestamp,
diff --git a/spark/src/test/resources/sql-tests/expressions/array/slice.sql b/spark/src/test/resources/sql-tests/expressions/array/slice.sql
@@ -0,0 +1,62 @@
+-- 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_slice(arr array<int>, s int, l int) USING parquet
+
+statement
+INSERT INTO test_slice VALUES
+  (array(1, 2, 3, 4, 5), 2, 3),
+  (array(1, 2, 3, 4, 5), 1, 5),
+  (array(1, 2, 3, 4, 5), 1, 10),
+  (array(1, 2, 3, 4, 5), 3, 0),
+  (array(1, 2, 3, 4, 5), -2, 2),
+  (array(1, 2, 3, 4, 5), -10, 2),
+  (array(1, 2, 3, 4, 5), 10, 2),
+  (array(1, NULL, 3, NULL, 5), 1, 5),
+  (array(), 1, 3),
+  (NULL, 1, 3),
+  (array(1, 2, 3), NULL, 2),
+  (array(1, 2, 3), 1, NULL)
+
+-- column array, column start, column length
+query
+SELECT slice(arr, s, l) FROM test_slice
+
+-- column array, literal start and length
+query
+SELECT slice(arr, 2, 2) FROM test_slice
+
+-- column array, negative literal start
+query
+SELECT slice(arr, -1, 1) FROM test_slice
+
+-- string element type
+statement
+CREATE TABLE test_slice_string(arr array<string>) USING parquet
+
+statement
+INSERT INTO test_slice_string VALUES
+  (array('a', 'b', 'c', 'd')),
+  (array('é', '日本', '', 'x')),
+  (array('a', NULL, 'c')),
+  (NULL)
+
+query
+SELECT slice(arr, 1, 2) FROM test_slice_string
diff --git a/spark/src/test/resources/sql-tests/expressions/bitwise/shiftrightunsigned.sql b/spark/src/test/resources/sql-tests/expressions/bitwise/shiftrightunsigned.sql
@@ -0,0 +1,63 @@
+-- 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
+
+-- Spark's ShiftRightUnsigned: first arg is Int or Long, second is Int.
+-- Returns the same integer type as the first argument. Shift amount is
+-- normalized to the bit width (Java semantics) for negative/large shifts.
+
+statement
+CREATE TABLE test_shiftrightunsigned_int(v int, s int) USING parquet
+
+statement
+INSERT INTO test_shiftrightunsigned_int VALUES
+  (1, 1),
+  (-1, 1),
+  (8, 2),
+  (2147483647, 1),
+  (-2147483648, 1),
+  (0, 0),
+  (1, 0),
+  (1, 31),
+  (1, 32),
+  (1, 33),
+  (1, -1),
+  (NULL, 1),
+  (1, NULL)
+
+query
+SELECT shiftrightunsigned(v, s) FROM test_shiftrightunsigned_int
+
+statement
+CREATE TABLE test_shiftrightunsigned_long(v bigint, s int) USING parquet
+
+statement
+INSERT INTO test_shiftrightunsigned_long VALUES
+  (1, 1),
+  (-1, 1),
+  (9223372036854775807, 1),
+  (-9223372036854775808, 1),
+  (0, 0),
+  (1, 63),
+  (1, 64),
+  (1, -1),
+  (NULL, 1),
+  (1, NULL)
+
+query
+SELECT shiftrightunsigned(v, s) FROM test_shiftrightunsigned_long
diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/unix_micros.sql b/spark/src/test/resources/sql-tests/expressions/datetime/unix_micros.sql
@@ -0,0 +1,41 @@
+-- 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.
+
+-- The result is microseconds since epoch in UTC, so it must not depend on the
+-- session timezone.
+-- ConfigMatrix: spark.sql.session.timeZone=UTC,America/Los_Angeles
+
+statement
+CREATE TABLE test_unix_micros(ts timestamp) USING parquet
+
+statement
+INSERT INTO test_unix_micros VALUES
+  (timestamp('1970-01-01 00:00:00')),
+  (timestamp('2024-01-15 12:34:56.123456')),
+  (timestamp('1969-12-31 23:59:59.999999')),
+  (timestamp('9999-12-31 23:59:59.999999')),
+  (timestamp('0001-01-01 00:00:00')),
+  (NULL)
+
+query
+SELECT unix_micros(ts) FROM test_unix_micros
+
+-- literal arguments
+query
+SELECT unix_micros(timestamp('1970-01-01 00:00:00')),
+       unix_micros(timestamp('2024-01-15 12:34:56.123456')),
+       unix_micros(NULL)
diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/unix_millis.sql b/spark/src/test/resources/sql-tests/expressions/datetime/unix_millis.sql
@@ -0,0 +1,41 @@
+-- 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.
+
+-- The result is milliseconds since epoch in UTC, so it must not depend on the
+-- session timezone.
+-- ConfigMatrix: spark.sql.session.timeZone=UTC,America/Los_Angeles
+
+statement
+CREATE TABLE test_unix_millis(ts timestamp) USING parquet
+
+statement
+INSERT INTO test_unix_millis VALUES
+  (timestamp('1970-01-01 00:00:00')),
+  (timestamp('2024-01-15 12:34:56.123456')),
+  (timestamp('1969-12-31 23:59:59.999999')),
+  (timestamp('9999-12-31 23:59:59.999999')),
+  (timestamp('0001-01-01 00:00:00')),
+  (NULL)
+
+query
+SELECT unix_millis(ts) FROM test_unix_millis
+
+-- literal arguments
+query
+SELECT unix_millis(timestamp('1970-01-01 00:00:00')),
+       unix_millis(timestamp('2024-01-15 12:34:56.123456')),
+       unix_millis(NULL)
diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/unix_seconds.sql b/spark/src/test/resources/sql-tests/expressions/datetime/unix_seconds.sql
