From f262ba1ed15c11aac67de8cf8642741b34796bfb Mon Sep 17 00:00:00 2001
From: Denis Khalikov <deniskhalikov@ydb.tech>
Date: Thu, 26 Jun 2025 21:14:13 +0300
Subject: [PATCH] =?UTF-8?q?[KQP=20Constant=20folding]=20Do=20not=20fold=20?=
 =?UTF-8?q?SafeCast=20because=20it=20could=20produce=20=E2=80=A6=20(#20223?=
 =?UTF-8?q?)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../opt/kqp_constant_folding_transformer.cpp  | 19 +++++++++++++++++++
 .../join/data/queries/udf_constant_fold.sql   |  3 +++
 2 files changed, 22 insertions(+)

diff --git a/ydb/core/kqp/opt/kqp_constant_folding_transformer.cpp b/ydb/core/kqp/opt/kqp_constant_folding_transformer.cpp
index 369fa82fec10..8f1e89278433 100644
--- a/ydb/core/kqp/opt/kqp_constant_folding_transformer.cpp
+++ b/ydb/core/kqp/opt/kqp_constant_folding_transformer.cpp
@@ -10,11 +10,30 @@ using namespace NKikimr::NKqp;
 using namespace NYql::NDq;
 
 namespace {
+    THashSet<TString> notAllowedDataTypeForSafeCast{"JsonDocument"};
+
+    bool IsSuitableToExtractExpr(const TExprNode::TPtr &input) {
+        if (auto maybeSafeCast = TExprBase(input).Maybe<TCoSafeCast>()) {
+            auto maybeDataType = maybeSafeCast.Cast().Type().Maybe<TCoDataType>();
+            if (!maybeDataType) {
+                if (const auto maybeOptionalType = maybeSafeCast.Cast().Type().Maybe<TCoOptionalType>()) {
+                    maybeDataType = maybeOptionalType.Cast().ItemType().Maybe<TCoDataType>();
+                }
+            }
+            return (maybeDataType && !notAllowedDataTypeForSafeCast.contains(maybeDataType.Cast().Type().Value()));
+        }
+        return true;
+    }
+
     /**
      * Traverse a lambda and create a mapping from nodes to nodes wrapped in EvaluateExpr callable
      * We check for literals specifically, since they shouldn't be evaluated
      */
     void ExtractConstantExprs(const TExprNode::TPtr& input, TNodeOnNodeOwnedMap& replaces, TExprContext& ctx, bool foldUdfs = true) {
+        if (!IsSuitableToExtractExpr(input)) {
+            return;
+        }
+
         if (TCoLambda::Match(input.Get())) {
             auto lambda = TExprBase(input).Cast<TCoLambda>();
             return ExtractConstantExprs(lambda.Body().Ptr(), replaces, ctx);
diff --git a/ydb/core/kqp/ut/join/data/queries/udf_constant_fold.sql b/ydb/core/kqp/ut/join/data/queries/udf_constant_fold.sql
index 7755475e8784..b9eb9dc57e91 100644
--- a/ydb/core/kqp/ut/join/data/queries/udf_constant_fold.sql
+++ b/ydb/core/kqp/ut/join/data/queries/udf_constant_fold.sql
@@ -1,3 +1,6 @@
 SELECT *
 FROM `/Root/S` as S
 WHERE S.payload2 = String::HexDecode("54");
+
+SELECT CAST("[6]" AS JsonDocument) as doc
+FROM `/Root/R` as R;