YQL-19123 More SQL features for Variant type

More SQL features for Variant type
commit_hash:51c15343e2d24190ec59085888dfa3fd008cafcc

1 files changed, 49 insertions, 0 deletions

diff --git a/yql/essentials/sql/v1/builtin.cpp b/yql/essentials/sql/v1/builtin.cpp
index 5d33b2061fa..a4bec630660 100644
--- a/yql/essentials/sql/v1/builtin.cpp
+++ b/yql/essentials/sql/v1/builtin.cpp
@@ -3010,6 +3010,7 @@ struct TBuiltinFuncData {
             {"callableresulttype", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("CallableResultType", 1, 1) },
             {"callableargumenttype", BuildSimpleBuiltinFactoryCallback<TYqlCallableArgumentType>() },
             {"variantunderlyingtype", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("VariantUnderlyingType", 1, 1) },
+            {"variantitem", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("SqlVariantItem", 1, 1) },
             {"fromysonsimpletype", BuildNamedArgcBuiltinFactoryCallback<TCallNodeImpl>("FromYsonSimpleType", 2, 2) },
             {"currentutcdate", BuildNamedDepsArgcBuiltinFactoryCallback<TCallNodeDepArgs>(0, "CurrentUtcDate", 0, -1) },
             {"currentutcdatetime", BuildNamedDepsArgcBuiltinFactoryCallback<TCallNodeDepArgs>(0, "CurrentUtcDatetime", 0, -1) },
@@ -3705,6 +3706,54 @@ TNodePtr BuildBuiltinFunc(TContext& ctx, TPosition pos, TString name, const TVec
                 BuildTuple(pos, {BuildQuotedAtom(pos, ""), args[1]}),
             };
             return new TCallNodeImpl(pos, "FlattenMembers", 2, 2, flattenMembersArgs);
+        } else if (normalizedName == "visit" || normalizedName == "visitordefault") {
+            bool withDefault = normalizedName == "visitordefault";
+            TNodePtr variant;
+            TVector<TNodePtr> labels, handlers;
+            TMaybe<TNodePtr> dflt;
+            if (mustUseNamed && *mustUseNamed) {
+                *mustUseNamed = false;
+                auto &positional = *args[0]->GetTupleNode();
+                if (positional.GetTupleSize() != (withDefault ? 2 : 1)) {
+                    return new TInvalidBuiltin(pos, TStringBuilder() << name
+                        << " requires exactly " << (withDefault ? 2 : 1) << " positional arguments when named args are used");
+                }
+                auto &named = *args[1]->GetStructNode();
+                variant = positional.GetTupleElement(0);
+                auto &namedExprs = named.GetExprs();
+                labels.reserve(namedExprs.size());
+                handlers.reserve(namedExprs.size());
+                for (size_t idx = 0; idx < namedExprs.size(); idx++) {
+                    labels.push_back(BuildQuotedAtom(pos, namedExprs[idx]->GetLabel()));
+                    handlers.push_back(namedExprs[idx]);
+                }
+                if (withDefault) {
+                    dflt = positional.GetTupleElement(positional.GetTupleSize() - 1);
+                }
+            } else {
+                variant = args[0];
+                size_t defaultSuffix = withDefault ? 1 : 0;
+                labels.reserve(args.size() - 1 - defaultSuffix);
+                handlers.reserve(args.size() - 1 - defaultSuffix);
+                for (size_t idx = 0; idx + 1 < args.size() - defaultSuffix; idx++) {
+                    labels.push_back(BuildQuotedAtom(pos, ToString(idx)));
+                    handlers.push_back(args[idx + 1]);
+                }
+                if (withDefault) {
+                    dflt = args.back();
+                }
+            }
+            TVector<TNodePtr> resultArgs;
+            resultArgs.reserve(1 + labels.size() + handlers.size());
+            resultArgs.emplace_back(std::move(variant));
+            for (size_t idx = 0; idx < labels.size(); idx++) {
+                resultArgs.emplace_back(std::move(labels[idx]));
+                resultArgs.emplace_back(std::move(handlers[idx]));
+            }
+            if (dflt.Defined()) {
+                resultArgs.emplace_back(std::move(dflt->Get()));
+            }
+            return new TCallNodeImpl(pos, "SqlVisit", 1, -1, resultArgs);
         } else if (normalizedName == "sqlexternalfunction") {
             return new TCallNodeImpl(pos, "SqlExternalFunction", args);
         } else {