YQL-20039: Make Datetime::Format work with basic resources for old MKQL versions

commit_hash:cc278d11897e8f36249bd3b285613b61c853d034

(cherry picked from commit 06674e6)

Author: igormunkin

yql/essentials/minikql/comp_nodes/mkql_udf.cpp

@@ -8,6 +8,7 @@
 #include <yql/essentials/minikql/mkql_node_printer.h>
 #include <yql/essentials/minikql/mkql_type_builder.h>
 #include <yql/essentials/minikql/mkql_utils.h>
+#include <yql/essentials/minikql/datetime/datetime64.h>
 #include <yql/essentials/utils/yql_panic.h>
 
 #include <library/cpp/containers/stack_array/stack_array.h>
@@ -27,6 +28,46 @@ TString TruncateTypeDiff(const TString& s) {
     return s.substr(0,TypeDiffLimit) + "...";
 }
 
+static const char TMResourceName[] = "DateTime2.TM";
+static const char TM64ResourceName[] = "DateTime2.TM64";
+// XXX: This class implements the wrapper to properly handle the
+// case when the signature of the emitted callable (i.e. callable
+// type) requires the extended datetime resource as an argument,
+// but the basic one is given. It wraps the unboxed value with
+// the closure to add the bridge with the implicit datetime
+// resource conversion.
+class TDateTimeConvertWrapper: public NUdf::TBoxedValue {
+public:
+    TDateTimeConvertWrapper(NUdf::TUnboxedValue&& callable)
+        : Callable_(callable)
+    {};
+
+private:
+    NUdf::TUnboxedValue Run(const NUdf::IValueBuilder* valueBuilder, const NUdf::TUnboxedValuePod* args) const final {
+        return NUdf::TUnboxedValuePod(new TDateTimeConverter(Callable_.Run(valueBuilder, args)));
+    }
+
+    class TDateTimeConverter: public NUdf::TBoxedValue {
+    public:
+        TDateTimeConverter(NUdf::TUnboxedValue&& closure)
+            : Closure_(closure)
+        {}
+    private:
+        NUdf::TUnboxedValue Run(const NUdf::IValueBuilder* valueBuilder, const NUdf::TUnboxedValuePod* args) const final {
+            NUdf::TUnboxedValuePod newArg;
+            const auto arg = args[0];
+            const auto& narrow = *reinterpret_cast<const NYql::DateTime::TTMStorage*>(arg.GetRawPtr());
+            auto& extended = *reinterpret_cast<NYql::DateTime::TTM64Storage*>(newArg.GetRawPtr());
+            extended.From(narrow);
+            return Closure_.Run(valueBuilder, &newArg);
+        }
+
+        const NUdf::TUnboxedValue Closure_;
+    };
+
+    const NUdf::TUnboxedValue Callable_;
+};
+
 template<class TValidatePolicy, class TValidateMode>
 class TSimpleUdfWrapper: public TMutableComputationNode<TSimpleUdfWrapper<TValidatePolicy,TValidateMode>> {
 using TBaseComputation = TMutableComputationNode<TSimpleUdfWrapper<TValidatePolicy,TValidateMode>>;
@@ -38,14 +79,16 @@ using TBaseComputation = TMutableComputationNode<TSimpleUdfWrapper<TValidatePoli
             NUdf::TSourcePosition pos,
             const TCallableType* callableType,
             const TCallableType* functionType,
-            TType* userType)
+            TType* userType,
+            bool wrapDateTimeConvert)
         : TBaseComputation(mutables, EValueRepresentation::Boxed)
         , FunctionName(std::move(functionName))
         , TypeConfig(std::move(typeConfig))
         , Pos(pos)
         , CallableType(callableType)
         , FunctionType(functionType)
         , UserType(userType)
+        , WrapDateTimeConvert(wrapDateTimeConvert)
     {
         this->Stateless = false;
     }
@@ -62,6 +105,7 @@ using TBaseComputation = TMutableComputationNode<TSimpleUdfWrapper<TValidatePoli
         NUdf::TUnboxedValue udf(NUdf::TUnboxedValuePod(funcInfo.Implementation.Release()));
         TValidate<TValidatePolicy,TValidateMode>::WrapCallable(FunctionType, udf, TStringBuilder() << "FunctionWithConfig<" << FunctionName << ">");
         ExtendArgs(udf, CallableType, funcInfo.FunctionType);
+        ConvertDateTimeArg(udf);
         return udf.Release();
     }
 private:
@@ -102,6 +146,12 @@ using TBaseComputation = TMutableComputationNode<TSimpleUdfWrapper<TValidatePoli
         }
     }
 
+    void ConvertDateTimeArg(NUdf::TUnboxedValue& callable) const {
+        if (WrapDateTimeConvert) {
+            callable = NUdf::TUnboxedValuePod(new TDateTimeConvertWrapper(std::move(callable)));
+        }
+    }
+
     void RegisterDependencies() const final {}
 
     const TString FunctionName;
@@ -110,6 +160,7 @@ using TBaseComputation = TMutableComputationNode<TSimpleUdfWrapper<TValidatePoli
     const TCallableType *const CallableType;
     const TCallableType *const FunctionType;
     TType *const UserType;
+    bool WrapDateTimeConvert;
 };
 
 class TUdfRunCodegeneratorNode: public TSimpleUdfWrapper<TValidateErrorPolicyNone, TValidateModeLazy<TValidateErrorPolicyNone>>
@@ -126,11 +177,12 @@ class TUdfRunCodegeneratorNode: public TSimpleUdfWrapper<TValidateErrorPolicyNon
             const TCallableType* callableType,
             const TCallableType* functionType,
             TType* userType,
+            bool wrapDateTimeConvert,
             TString&& moduleIRUniqID,
             TString&& moduleIR,
             TString&& fuctioNameIR,
             NUdf::TUniquePtr<NUdf::IBoxedValue>&& impl)
-        : TSimpleUdfWrapper(mutables, std::move(functionName), std::move(typeConfig), pos, callableType, functionType, userType)
+        : TSimpleUdfWrapper(mutables, std::move(functionName), std::move(typeConfig), pos, callableType, functionType, userType, wrapDateTimeConvert)
         , ModuleIRUniqID(std::move(moduleIRUniqID))
         , ModuleIR(std::move(moduleIR))
         , IRFunctionName(std::move(fuctioNameIR))
@@ -174,7 +226,8 @@ using TBaseComputation = TMutableCodegeneratorPtrNode<TUdfWrapper<TValidatePolic
             IComputationNode* runConfigNode,
             ui32 runConfigArgs,
             const TCallableType* callableType,
-            TType* userType)
+            TType* userType,
+            bool wrapDateTimeConvert)
         : TBaseComputation(mutables, EValueRepresentation::Boxed)
         , FunctionName(std::move(functionName))
         , TypeConfig(std::move(typeConfig))
@@ -183,6 +236,7 @@ using TBaseComputation = TMutableCodegeneratorPtrNode<TUdfWrapper<TValidatePolic
         , RunConfigArgs(runConfigArgs)
         , CallableType(callableType)
         , UserType(userType)
+        , WrapDateTimeConvert(wrapDateTimeConvert)
         , UdfIndex(mutables.CurValueIndex++)
     {
         this->Stateless = false;
@@ -193,6 +247,7 @@ using TBaseComputation = TMutableCodegeneratorPtrNode<TUdfWrapper<TValidatePolic
         if (!udf.HasValue()) {
             MakeUdf(ctx, udf);
         }
+        ConvertDateTimeArg(udf);
         NStackArray::TStackArray<NUdf::TUnboxedValue> args(ALLOC_ON_STACK(NUdf::TUnboxedValue, RunConfigArgs));
         args[0] = RunConfigNode->GetValue(ctx);
         auto callable = udf.Run(ctx.Builder, args.data());
@@ -226,6 +281,11 @@ using TBaseComputation = TMutableCodegeneratorPtrNode<TUdfWrapper<TValidatePolic
 
         block = main;
 
+        const auto convertFunc = ConstantInt::get(Type::getInt64Ty(context), GetMethodPtr(&TUdfWrapper::ConvertDateTimeArg));
+        const auto convertType = FunctionType::get(Type::getVoidTy(context), {self->getType(), udfPtr->getType()}, false);
+        const auto convertFuncPtr = CastInst::Create(Instruction::IntToPtr, convertFunc, PointerType::getUnqual(convertType), "convert", block);
+        CallInst::Create(convertType, convertFuncPtr, {self, udfPtr}, "", block);
+
         const auto argsType = ArrayType::get(valueType, RunConfigArgs);
         const auto args = new AllocaInst(argsType, 0U, "args", block);
         const auto zero = ConstantInt::get(indexType, 0);
@@ -270,6 +330,12 @@ using TBaseComputation = TMutableCodegeneratorPtrNode<TUdfWrapper<TValidatePolic
         TValidate<TValidatePolicy,TValidateMode>::WrapCallable(CallableType, callable, TStringBuilder() << "FunctionWithConfig<" << FunctionName << ">");
     }
 
+    void ConvertDateTimeArg(NUdf::TUnboxedValue& callable) const {
+        if (WrapDateTimeConvert) {
+            callable = NUdf::TUnboxedValuePod(new TDateTimeConvertWrapper(std::move(callable)));
+        }
+    }
+
     void RegisterDependencies() const final {
         this->DependsOn(RunConfigNode);
     }
@@ -281,6 +347,7 @@ using TBaseComputation = TMutableCodegeneratorPtrNode<TUdfWrapper<TValidatePolic
     const ui32 RunConfigArgs;
     const TCallableType* CallableType;
     TType* const UserType;
+    bool WrapDateTimeConvert;
     const ui32 UdfIndex;
 };
 
@@ -310,6 +377,63 @@ inline IComputationNode* CreateUdfWrapper(const TComputationNodeFactoryContext&
     };
 }
 
+// XXX: The helper below allows to make a stitchless upgrade
+// of MKQL runtime, regarding the incompatible changes made for
+// DateTime::Format UDF.
+template<bool Extended>
+static bool IsDateTimeResource(const TType* type) {
+    if (!type->IsResource()) {
+        return false;
+    }
+    const auto resourceName = AS_TYPE(TResourceType, type)->GetTag();
+
+    if constexpr (Extended) {
+        return resourceName == NUdf::TStringRef::Of(TM64ResourceName);
+    } else {
+        return resourceName == NUdf::TStringRef::Of(TMResourceName);
+    }
+}
+
+static bool IsDateTimeConvertible(const NUdf::TStringRef& funcName,
+                                  const TCallableType* nodeType,
+                                  const TCallableType* funcType,
+                                  bool& needConvert)
+{
+    Y_DEBUG_ABORT_UNLESS(!needConvert);
+    if (funcName == NUdf::TStringRef::Of("DateTime2.Format")) {
+        Y_DEBUG_ABORT_UNLESS(nodeType->GetArgumentsCount());
+        Y_DEBUG_ABORT_UNLESS(funcType->GetArgumentsCount());
+        // XXX: In general, DateTime resources are not convertible
+        // in runtime, but for the stitchless upgrade of MKQL
+        // runtime, we consider the basic resource, being
+        // convertible to the extended one for DateTime2.Format...
+        if (IsDateTimeResource<false>(nodeType->GetArgumentType(0)) &&
+            IsDateTimeResource<true>(funcType->GetArgumentType(0)))
+        {
+            // XXX: ... and to implicitly convert the basic resource
+            // to the extended one, the closure has to be wrapped by
+            // DateTimeConvertWrapper.
+            needConvert = true;
+            return true;
+        }
+    }
+    if (funcName == NUdf::TStringRef::Of("DateTime2.Convert")) {
+        // XXX: Vice versa convertion is forbidden as well, but
+        // for the stitchless upgrade of MKQL runtime, we consider
+        // the extended resource, being compatible with the basic
+        // one for the return type of DateTime2.Convert.
+        if (IsDateTimeResource<true>(nodeType->GetReturnType()) &&
+            IsDateTimeResource<false>(funcType->GetReturnType()))
+        {
+            // XXX: However, DateTime2.Convert has to convert the
+            // basic resource to the extended one.
+            needConvert = false;
+            return true;
+        }
+    }
+    return false;
+}
+
 }
 
 IComputationNode* WrapUdf(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
@@ -343,6 +467,7 @@ IComputationNode* WrapUdf(TCallable& callable, const TComputationNodeFactoryCont
 
     MKQL_ENSURE(status.IsOk(), status.GetError());
 
+    bool wrapDateTimeConvert = false;
     const auto callableFuncType = AS_TYPE(TCallableType, funcInfo.FunctionType);
     const auto callableNodeType = AS_TYPE(TCallableType, callable.GetType()->GetReturnType());
     const auto runConfigFuncType = funcInfo.RunConfigType;
@@ -401,46 +526,59 @@ IComputationNode* WrapUdf(TCallable& callable, const TComputationNodeFactoryCont
         }
         const auto closureFuncType = runConfigNodeType->IsVoid()
                                    ? callableFuncType
-                                   : AS_TYPE(TCallableType, callableFuncType)->GetReturnType();
+                                   : AS_TYPE(TCallableType, callableFuncType->GetReturnType());
         const auto closureNodeType = runConfigNodeType->IsVoid()
-                                   ? AS_TYPE(TCallableType, callableNodeType)->GetReturnType()
+                                   ? AS_TYPE(TCallableType, callableNodeType->GetReturnType())
                                    : callableNodeType;
         if (!closureNodeType->IsConvertableTo(*closureFuncType)) {
-            TString diff = TStringBuilder()
-                << "type mismatch, expected return type: "
-                << PrintNode(closureNodeType, true)
-                << ", actual: "
-                << PrintNode(closureFuncType, true);
-            UdfTerminate((TStringBuilder() << pos
-                                           << " Udf Function '"
-                                           << funcName
-                                           << "' "
-                                           << TruncateTypeDiff(diff)).c_str());
+            if (!IsDateTimeConvertible(funcName, closureNodeType, closureFuncType, wrapDateTimeConvert)) {
+                TString diff = TStringBuilder()
+                    << "type mismatch, expected return type: "
+                    << PrintNode(closureNodeType, true)
+                    << ", actual: "
+                    << PrintNode(closureFuncType, true);
+                UdfTerminate((TStringBuilder() << pos
+                                               << " Udf Function '"
+                                               << funcName
+                                               << "' "
+                                               << TruncateTypeDiff(diff)).c_str());
+            }
+            MKQL_ENSURE(funcName == NUdf::TStringRef::Of("DateTime2.Format") ||
+                        funcName == NUdf::TStringRef::Of("DateTime2.Convert"),
+                        "Unexpected function violates the convertible invariants");
         }
 
         const auto runConfigCompNode = LocateNode(ctx.NodeLocator, *runCfgNode.GetNode());
         const auto runConfigArgs = funcInfo.FunctionType->GetArgumentsCount();
         return runConfigNodeType->IsVoid()
-            ? CreateUdfWrapper<true>(ctx, std::move(funcName), std::move(typeConfig), pos, callableNodeType, callableFuncType, userType)
-            : CreateUdfWrapper<false>(ctx, std::move(funcName), std::move(typeConfig), pos, runConfigCompNode, runConfigArgs, callableNodeType, userType);
+            ? CreateUdfWrapper<true>(ctx, std::move(funcName), std::move(typeConfig), pos, callableNodeType, callableFuncType, userType, wrapDateTimeConvert)
+            : CreateUdfWrapper<false>(ctx, std::move(funcName), std::move(typeConfig), pos, runConfigCompNode, runConfigArgs, callableNodeType, userType, wrapDateTimeConvert);
+    }
+
+    if (!callableFuncType->IsConvertableTo(*callableNodeType, true)) {
+        if (!IsDateTimeConvertible(funcName, callableNodeType, callableFuncType, wrapDateTimeConvert)) {
+            TString diff = TStringBuilder() << "type mismatch, expected return type: " << PrintNode(callableNodeType, true) <<
+                    ", actual:" << PrintNode(callableFuncType, true);
+            UdfTerminate((TStringBuilder() << pos << " UDF Function '" << funcName << "' " << TruncateTypeDiff(diff)).c_str());
+        }
+        MKQL_ENSURE(funcName == NUdf::TStringRef::Of("DateTime2.Format") ||
+                    funcName == NUdf::TStringRef::Of("DateTime2.Convert"),
+                    "Unexpected function violates the convertible invariants");
     }
-    MKQL_ENSURE(callableFuncType->IsConvertableTo(*callableNodeType, true),
-                "Function '" << funcName << "' type mismatch, expected return type: " << PrintNode(callableNodeType, true) <<
-                ", actual:" << PrintNode(callableFuncType, true));
     MKQL_ENSURE(funcInfo.Implementation, "UDF implementation is not set for function " << funcName);
 
     if (runConfigFuncType->IsVoid()) {
         if (ctx.ValidateMode == NUdf::EValidateMode::None && funcInfo.ModuleIR && funcInfo.IRFunctionName) {
             return new TUdfRunCodegeneratorNode(
-                ctx.Mutables, std::move(funcName), std::move(typeConfig), pos, callableNodeType, callableFuncType, userType,
+                ctx.Mutables, std::move(funcName), std::move(typeConfig), pos, callableNodeType, callableFuncType, userType, wrapDateTimeConvert,
                 std::move(funcInfo.ModuleIRUniqID), std::move(funcInfo.ModuleIR), std::move(funcInfo.IRFunctionName), std::move(funcInfo.Implementation)
             );
         }
-        return CreateUdfWrapper<true>(ctx, std::move(funcName), std::move(typeConfig), pos, callableNodeType, callableFuncType, userType);
+        return CreateUdfWrapper<true>(ctx, std::move(funcName), std::move(typeConfig), pos, callableNodeType, callableFuncType, userType, wrapDateTimeConvert);
     }
 
     const auto runCfgCompNode = LocateNode(ctx.NodeLocator, *runCfgNode.GetNode());
-    return CreateUdfWrapper<false>(ctx, std::move(funcName), std::move(typeConfig), pos, runCfgCompNode, 1U, callableNodeType, userType);
+    return CreateUdfWrapper<false>(ctx, std::move(funcName), std::move(typeConfig), pos, runCfgCompNode, 1U, callableNodeType, userType, wrapDateTimeConvert);
 }
 
 IComputationNode* WrapScriptUdf(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
@@ -479,7 +617,7 @@ IComputationNode* WrapScriptUdf(TCallable& callable, const TComputationNodeFacto
     const auto funcTypeInfo = static_cast<TCallableType*>(callableResultType);
 
     const auto programCompNode = LocateNode(ctx.NodeLocator, *programNode.GetNode());
-    return CreateUdfWrapper<false>(ctx, std::move(funcName), std::move(typeConfig), pos, programCompNode, 1U, funcTypeInfo, userType);
+    return CreateUdfWrapper<false>(ctx, std::move(funcName), std::move(typeConfig), pos, programCompNode, 1U, funcTypeInfo, userType, false);
 }
 
 }