DX-126900: [C++][Gandiva] Add TimeIR for unit-aware time64[us/ns] support

cpp/src/gandiva/CMakeLists.txt

@@ -62,6 +62,7 @@ set(SRC_FILES
     decimal_ir.cc
     decimal_type_util.cc
     timestamp_ir.cc
+    time_ir.cc
     decimal_xlarge.cc
     engine.cc
     date_utils.cc

cpp/src/gandiva/function_registry_arithmetic.cc

@@ -221,6 +221,13 @@ std::vector<NativeFunction> GetArithmeticFunctionRegistry() {
       BINARY_RELATIONAL_BOOL_DATE_FN(less_than_or_equal_to, {}),
       BINARY_RELATIONAL_BOOL_DATE_FN(greater_than, {}),
       BINARY_RELATIONAL_BOOL_DATE_FN(greater_than_or_equal_to, {}),
+      // time64 compare functions (unit-agnostic, see TimeIR)
+      BINARY_RELATIONAL_SAFE_NULL_IF_NULL(equal, ({"eq", "same"}), time64),
+      BINARY_RELATIONAL_SAFE_NULL_IF_NULL(not_equal, {}, time64),
+      BINARY_RELATIONAL_SAFE_NULL_IF_NULL(less_than, {}, time64),
+      BINARY_RELATIONAL_SAFE_NULL_IF_NULL(less_than_or_equal_to, {}, time64),
+      BINARY_RELATIONAL_SAFE_NULL_IF_NULL(greater_than, {}, time64),
+      BINARY_RELATIONAL_SAFE_NULL_IF_NULL(greater_than_or_equal_to, {}, time64),
       BASE_NUMERIC_TYPES(MULTIPLE_SAFE_NULL_IF_NULL, greatest, {}),
       BASE_NUMERIC_TYPES(MULTIPLE_SAFE_NULL_IF_NULL, least, {}),
 

cpp/src/gandiva/function_registry_datetime.cc

@@ -45,6 +45,13 @@ namespace gandiva {
       TIME_TYPES(EXTRACT_SAFE_NULL_IF_NULL, name##Minute, {"minute"}), \
       TIME_TYPES(EXTRACT_SAFE_NULL_IF_NULL, name##Second, {"second"})
 
+// time64 functions are registered for time64[us]; time64[ns] calls match the same
+// signature and are remapped to the _ns variants by TimeIR at codegen time.
+#define TIME64_EXTRACTION_FNS(name)                                \
+  EXTRACT_SAFE_NULL_IF_NULL(name##Hour, {"hour"}, time64),         \
+      EXTRACT_SAFE_NULL_IF_NULL(name##Minute, {"minute"}, time64), \
+      EXTRACT_SAFE_NULL_IF_NULL(name##Second, {"second"}, time64)
+
 #define NEXT_DAY_FNS(name) DATE_TYPES(NEXT_DAY_SAFE_NULL_IF_NULL, name, {})
 
 std::vector<NativeFunction> GetDateTimeFunctionRegistry() {
@@ -59,6 +66,7 @@ std::vector<NativeFunction> GetDateTimeFunctionRegistry() {
       DATE_TYPES(EXTRACT_SAFE_NULL_IF_NULL, extractEpoch, {}),
 
       TIME_EXTRACTION_FNS(extract),
+      TIME64_EXTRACTION_FNS(extract),
 
       NEXT_DAY_FNS(next_day),
 

cpp/src/gandiva/function_registry_math_ops.cc

@@ -70,11 +70,15 @@ std::vector<NativeFunction> GetMathOpsFunctionRegistry() {
 
       UNARY_SAFE_NULL_NEVER_BOOL_FN(isnull, {}),
       UNARY_SAFE_NULL_NEVER_BOOL_FN(isnotnull, {}),
+      UNARY_SAFE_NULL_NEVER_BOOL(isnull, {}, time64),
+      UNARY_SAFE_NULL_NEVER_BOOL(isnotnull, {}, time64),
 
       NUMERIC_TYPES(UNARY_SAFE_NULL_NEVER_BOOL, isnumeric, {}),
 
       BINARY_SAFE_NULL_NEVER_BOOL_FN(is_distinct_from, {}),
       BINARY_SAFE_NULL_NEVER_BOOL_FN(is_not_distinct_from, {}),
+      BINARY_SAFE_NULL_NEVER_BOOL(is_distinct_from, {}, time64),
+      BINARY_SAFE_NULL_NEVER_BOOL(is_not_distinct_from, {}, time64),
 
       UNARY_UNSAFE_NULL_IF_NULL(factorial, {}, int32, int64),
       UNARY_UNSAFE_NULL_IF_NULL(factorial, {}, int64, int64),

cpp/src/gandiva/function_signature.cc

@@ -56,6 +56,11 @@ bool DataTypeEquals(const DataTypePtr& left, const DataTypePtr& right) {
         auto tright = checked_cast<arrow::TimestampType*>(right.get());
         return tleft->timezone() == tright->timezone();
       }
+      case arrow::Type::TIME64: {
+        // For time64 types, the TimeUnit isn't part of the signature
+        // (conversion is handled at codegen time by TimeIR).
+        return true;
+      }
       default:
         return left->Equals(right);
     }

cpp/src/gandiva/function_signature_test.cc

@@ -110,4 +110,20 @@ TEST_F(TestFunctionSignature, TestHash) {
   EXPECT_EQ(f3.Hash(), f4.Hash());
 }
 
+TEST_F(TestFunctionSignature, TestTime64UnitNotPartOfSignature) {
+  // time64 units are resolved at codegen time by TimeIR, so time64[us] and time64[ns]
+  // match the same registered signature.
+  FunctionSignature micro("extractHour", {arrow::time64(arrow::TimeUnit::MICRO)},
+                          arrow::int64());
+  FunctionSignature nano("extractHour", {arrow::time64(arrow::TimeUnit::NANO)},
+                         arrow::int64());
+  EXPECT_EQ(micro, nano);
+  EXPECT_EQ(micro.Hash(), nano.Hash());
+
+  // time32 and time64 remain distinct types.
+  EXPECT_FALSE(micro == FunctionSignature("extractHour",
+                                          {arrow::time32(arrow::TimeUnit::MILLI)},
+                                          arrow::int64()));
+}
+
 }  // namespace gandiva

cpp/src/gandiva/llvm_generator.cc

@@ -30,6 +30,7 @@
 #include "gandiva/expression.h"
 #include "gandiva/llvm_types.h"
 #include "gandiva/lvalue.h"
+#include "gandiva/time_ir.h"
 #include "gandiva/timestamp_ir.h"
 
 namespace gandiva {
@@ -1296,6 +1297,40 @@ Result<std::string> LLVMGenerator::ResolveTimestampPcName(const std::string& pc_
   return pc_name;
 }
 
+Result<std::string> LLVMGenerator::ResolveTimePcName(const std::string& pc_name,
+                                                     const DataTypeVector& params) {
+  arrow::TimeUnit::type time_unit = arrow::TimeUnit::MICRO;
+  bool found_time = false;
+  for (const auto& param : params) {
+    if (param->id() == arrow::Type::TIME64) {
+      auto unit = arrow::internal::checked_cast<const arrow::Time64Type&>(*param).unit();
+      if (!found_time) {
+        time_unit = unit;
+        found_time = true;
+      } else if (unit != time_unit) {
+        return Status::Invalid(
+            "Gandiva cannot compile expression: mixed time64 units in function '",
+            pc_name, "'. All time64 arguments must have the same TimeUnit.");
+      }
+    }
+  }
+  // time64 signatures are registered with MICRO, so the unsuffixed function already
+  // handles MICRO arguments.
+  if (!found_time || time_unit == arrow::TimeUnit::MICRO ||
+      TimeIR::IsUnitAgnosticFunction(pc_name)) {
+    return pc_name;
+  }
+  std::string remapped = pc_name + TimeIR::UnitSuffix(time_unit);
+  if (TimeIR::IsTimeIRFunction(remapped)) {
+    ARROW_LOG(DEBUG) << "TimeIR remap: " << pc_name << " -> " << remapped;
+    return remapped;
+  }
+  // Never fall through to the MICRO function: it would misinterpret the data.
+  return Status::Invalid("Gandiva cannot compile expression: function '", pc_name,
+                         "' has no variant for time64 unit ",
+                         arrow::internal::ToString(time_unit), ".");
+}
+
 LValuePtr LLVMGenerator::Visitor::BuildFunctionCall(const NativeFunction* func,
                                                     DataTypePtr arrow_return_type,
                                                     std::vector<llvm::Value*>* params,
@@ -1315,6 +1350,12 @@ LValuePtr LLVMGenerator::Visitor::BuildFunctionCall(const NativeFunction* func,
       return nullptr;
     }
     pc_name = resolve_result.MoveValueUnsafe();
+    auto time_resolve_result = ResolveTimePcName(pc_name, descriptor->params());
+    if (!time_resolve_result.ok()) {
+      status_ = time_resolve_result.status();
+      return nullptr;
+    }
+    pc_name = time_resolve_result.MoveValueUnsafe();
   }
 
   if (arrow_return_type_id == arrow::Type::DECIMAL) {

cpp/src/gandiva/llvm_generator.h

@@ -67,6 +67,13 @@ class GANDIVA_EXPORT LLVMGenerator {
   static Result<std::string> ResolveTimestampPcName(const std::string& pc_name,
                                                      const DataTypeVector& params);
 
+  /// \brief Resolve the precompiled function name, remapping to a TimeIR variant
+  /// when the function's params include a time64[ns] argument.
+  /// Returns an error if params contain mixed time64 TimeUnits, or if the function
+  /// has no variant for the time64 unit and is not unit-agnostic.
+  static Result<std::string> ResolveTimePcName(const std::string& pc_name,
+                                               const DataTypeVector& params);
+
   /// \brief Build the code for the expression trees for default mode with a LLVM
   /// ObjectCache. Each element in the vector represents an expression tree
   Status Build(const ExpressionVector& exprs, SelectionVector::Mode mode);
@@ -96,6 +103,7 @@ class GANDIVA_EXPORT LLVMGenerator {
                          std::shared_ptr<FunctionRegistry> function_registry);
 
   FRIEND_TEST(TestLLVMGenerator, VerifyPCFunctions);
+  FRIEND_TEST(TestLLVMGenerator, VerifyTime64Functions);
   FRIEND_TEST(TestLLVMGenerator, TestAdd);
   FRIEND_TEST(TestLLVMGenerator, TestNullInternal);
   friend class TestLLVMGenerator;

cpp/src/gandiva/llvm_generator_test.cc

@@ -28,6 +28,7 @@
 #include "gandiva/func_descriptor.h"
 #include "gandiva/function_registry.h"
 #include "gandiva/tests/test_util.h"
+#include "gandiva/time_ir.h"
 #include "gandiva/timestamp_ir.h"
 
 namespace gandiva {
@@ -234,4 +235,87 @@ TEST_F(TestLLVMGenerator, ResolveTimestampPcName) {
   }
 }
 
+TEST_F(TestLLVMGenerator, ResolveTimePcName) {
+  // Non-time64 params — name unchanged
+  {
+    DataTypeVector params{arrow::time32(arrow::TimeUnit::MILLI)};
+    ASSERT_OK_AND_ASSIGN(auto name,
+                         LLVMGenerator::ResolveTimePcName("extractHour_time32", params));
+    EXPECT_EQ(name, "extractHour_time32");
+  }
+
+  // time64[us] — no remap (us is the registered time64 unit)
+  {
+    DataTypeVector params{arrow::time64(arrow::TimeUnit::MICRO)};
+    ASSERT_OK_AND_ASSIGN(auto name,
+                         LLVMGenerator::ResolveTimePcName("extractHour_time64", params));
+    EXPECT_EQ(name, "extractHour_time64");
+  }
+
+  // time64[ns] with a unit-specific function — gets _ns suffix
+  {
+    DataTypeVector params{arrow::time64(arrow::TimeUnit::NANO)};
+    ASSERT_OK_AND_ASSIGN(auto name,
+                         LLVMGenerator::ResolveTimePcName("extractHour_time64", params));
+    EXPECT_EQ(name, "extractHour_time64_ns");
+  }
+
+  // time64[ns] with a unit-agnostic function — name unchanged
+  {
+    DataTypeVector params{arrow::time64(arrow::TimeUnit::NANO),
+                          arrow::time64(arrow::TimeUnit::NANO)};
+    ASSERT_OK_AND_ASSIGN(
+        auto name, LLVMGenerator::ResolveTimePcName("less_than_time64_time64", params));
+    EXPECT_EQ(name, "less_than_time64_time64");
+  }
+
+  // Mixed time64 units — returns Invalid status
+  {
+    DataTypeVector params{arrow::time64(arrow::TimeUnit::MICRO),
+                          arrow::time64(arrow::TimeUnit::NANO)};
+    auto result = LLVMGenerator::ResolveTimePcName("less_than_time64_time64", params);
+    EXPECT_FALSE(result.ok());
+    EXPECT_THAT(result.status().message(), testing::HasSubstr("mixed time64 units"));
+  }
+
+  // time64[ns] without an _ns variant must not fall back to the microsecond function
+  {
+    DataTypeVector params{arrow::time64(arrow::TimeUnit::NANO)};
+    auto result = LLVMGenerator::ResolveTimePcName("no_such_function", params);
+    EXPECT_FALSE(result.ok());
+    EXPECT_THAT(result.status().message(), testing::HasSubstr("no variant for time64"));
+  }
+}
+
+// Every registered time64 function must resolve, for both time64 units, to a function
+// that exists in the precompiled module.
+TEST_F(TestLLVMGenerator, VerifyTime64Functions) {
+  ASSERT_OK_AND_ASSIGN(auto generator, LLVMGenerator::Make(TestConfiguration(), false));
+
+  llvm::Module* module = generator->module();
+  ASSERT_OK(generator->engine_->LoadFunctionIRs());
+  int num_time64_functions = 0;
+  for (auto& native_function : *registry_) {
+    for (auto& signature : native_function.signatures()) {
+      bool has_time64 = false;
+      for (auto& param : signature.param_types()) {
+        has_time64 |= param->id() == arrow::Type::TIME64;
+      }
+      if (!has_time64) continue;
+      ++num_time64_functions;
+      for (auto unit : {arrow::TimeUnit::MICRO, arrow::TimeUnit::NANO}) {
+        DataTypeVector params;
+        for (auto& param : signature.param_types()) {
+          params.push_back(param->id() == arrow::Type::TIME64 ? arrow::time64(unit)
+                                                              : param);
+        }
+        ASSERT_OK_AND_ASSIGN(auto name, LLVMGenerator::ResolveTimePcName(
+                                            native_function.pc_name(), params));
+        EXPECT_NE(module->getFunction(name), nullptr) << name;
+      }
+    }
+  }
+  EXPECT_GT(num_time64_functions, 0);
+}
+
 }  // namespace gandiva

cpp/src/gandiva/precompiled/CMakeLists.txt

@@ -27,6 +27,7 @@ set(PRECOMPILED_SRCS
     time.cc
     timestamp_arithmetic.cc
     timestamp_unit_ops.cc
+    time_unit_ops.cc
     ../../arrow/util/basic_decimal.cc)
 set(GANDIVA_PRECOMPILED_BC_FILES)
 foreach(SOURCE ${PRECOMPILED_SRCS})
@@ -63,6 +64,7 @@ add_gandiva_test(precompiled-test
                  time_test.cc
                  time.cc
                  timestamp_arithmetic.cc
+                 time_unit_ops.cc
                  ../cast_time.cc
                  ../../arrow/vendored/datetime.cpp
                  hash_test.cc

cpp/src/gandiva/precompiled/arithmetic_ops.cc

@@ -131,6 +131,15 @@ NUMERIC_DATE_TYPES(BINARY_RELATIONAL, less_than_or_equal_to, <=)
 NUMERIC_DATE_TYPES(BINARY_RELATIONAL, greater_than, >)
 NUMERIC_DATE_TYPES(BINARY_RELATIONAL, greater_than_or_equal_to, >=)
 
+// time64 values of a single unit (micro or nano) compare directly, so these are
+// unit-agnostic. Mixed units are rejected at codegen time (see TimeIR).
+BINARY_RELATIONAL(equal, time64, ==)
+BINARY_RELATIONAL(not_equal, time64, !=)
+BINARY_RELATIONAL(less_than, time64, <)
+BINARY_RELATIONAL(less_than_or_equal_to, time64, <=)
+BINARY_RELATIONAL(greater_than, time64, >)
+BINARY_RELATIONAL(greater_than_or_equal_to, time64, >=)
+
 #undef BINARY_RELATIONAL
 
 // Returns the greatest or least value from a list of values
@@ -236,6 +245,8 @@ CAST_INT_FLOAT(castINT, float64, int32)
 
 NUMERIC_BOOL_DATE_TYPES(VALIDITY_OP, isnull, !)
 NUMERIC_BOOL_DATE_TYPES(VALIDITY_OP, isnotnull, +)
+VALIDITY_OP(isnull, time64, !)
+VALIDITY_OP(isnotnull, time64, +)
 NUMERIC_TYPES(VALIDITY_OP, isnumeric, +)
 
 #undef VALIDITY_OP
@@ -346,6 +357,8 @@ gdv_boolean not_boolean(gdv_boolean in) { return !in; }
 
 NUMERIC_BOOL_DATE_FUNCTION(IS_DISTINCT_FROM)
 NUMERIC_BOOL_DATE_FUNCTION(IS_NOT_DISTINCT_FROM)
+IS_DISTINCT_FROM(time64)
+IS_NOT_DISTINCT_FROM(time64)
 
 #undef IS_DISTINCT_FROM
 #undef IS_NOT_DISTINCT_FROM

cpp/src/gandiva/precompiled/time_test.cc

@@ -249,6 +249,28 @@ TEST(TestTime, TestExtractTime) {
   EXPECT_EQ(extractSecond_time32(time_as_millis_in_day), 33);
 }
 
+TEST(TestTime, TestExtractTime64) {
+  // 10:20:33.123456 in micros
+  gdv_time64 time_as_micros_in_day = 37233123456LL;
+  EXPECT_EQ(extractHour_time64(time_as_micros_in_day), 10);
+  EXPECT_EQ(extractMinute_time64(time_as_micros_in_day), 20);
+  EXPECT_EQ(extractSecond_time64(time_as_micros_in_day), 33);
+
+  // 10:20:33.123456789 in nanos
+  gdv_time64 time_as_nanos_in_day = 37233123456789LL;
+  EXPECT_EQ(extractHour_time64_ns(time_as_nanos_in_day), 10);
+  EXPECT_EQ(extractMinute_time64_ns(time_as_nanos_in_day), 20);
+  EXPECT_EQ(extractSecond_time64_ns(time_as_nanos_in_day), 33);
+
+  // 23:59:59.999999 / 23:59:59.999999999
+  EXPECT_EQ(extractHour_time64(86399999999LL), 23);
+  EXPECT_EQ(extractMinute_time64(86399999999LL), 59);
+  EXPECT_EQ(extractSecond_time64(86399999999LL), 59);
+  EXPECT_EQ(extractHour_time64_ns(86399999999999LL), 23);
+  EXPECT_EQ(extractMinute_time64_ns(86399999999999LL), 59);
+  EXPECT_EQ(extractSecond_time64_ns(86399999999999LL), 59);
+}
+
 TEST(TestTime, TestDateDiff) {
   gdv_timestamp ts1 = StringToTimestamp("2019-06-30 00:00:00");
   gdv_timestamp ts2 = StringToTimestamp("2019-05-31 00:00:00");