Implement partitions() table function

commit_hash:19bbcd26d46a4ca8d18f0dbae605b48ac823c614

Author: nepal

yql/essentials/minikql/mkql_type_ops.cpp

@@ -2445,6 +2445,29 @@ bool IsValidStringValue(NUdf::EDataSlot type, NUdf::TStringRef buf) {
     MKQL_ENSURE(false, "Incorrect data slot: " << (ui32)type);
 }
 
+TMaybe<TString> RegexMatchingValidStringValues(NUdf::EDataSlot type, ERegexFlavor flavor) {
+    if (flavor != ERegexFlavor::RE2) {
+        return {};
+    }
+
+    const auto& typeInfo = NUdf::GetDataTypeInfo(type);
+    if (typeInfo.Features & NUdf::EDataTypeFeatures::IntegralType) {
+        TStringBuilder result;
+        if (typeInfo.Features & NUdf::EDataTypeFeatures::SignedIntegralType) {
+            result << R"([-+]?)";
+        } else {
+            result << R"(\+?)";
+        }
+        result << R"(\d+)";
+        return result;
+    }
+    if (type == NUdf::EDataSlot::Date) {
+        return R"(\d{4}-\d{2}-\d{2})";
+    }
+    // TODO: more cases
+    return {};
+}
+
 NUdf::TUnboxedValuePod ValueFromString(NUdf::EDataSlot type, NUdf::TStringRef buf) {
     switch (type) {
     case NUdf::EDataSlot::Bool: {

yql/essentials/minikql/mkql_type_ops.h

@@ -21,6 +21,11 @@ ui32 GetMonthLength(ui32 month, bool isLeap);
 
 bool IsValidStringValue(NUdf::EDataSlot type, NUdf::TStringRef buf);
 
+enum class ERegexFlavor {
+    RE2,
+};
+TMaybe<TString> RegexMatchingValidStringValues(NUdf::EDataSlot type, ERegexFlavor flavor);
+
 NUdf::TUnboxedValuePod ValueFromString(NUdf::EDataSlot type, NUdf::TStringRef buf);
 NUdf::TUnboxedValuePod SimpleValueFromYson(NUdf::EDataSlot type, NUdf::TStringRef buf);
 

yql/essentials/sql/v1/complete/name/service/static/name_service.cpp

@@ -189,6 +189,7 @@ class TFunctionNameService: public IRankingNameService {
                   "WalkFolders",
                   "EACH",
                   "PARTITION_LIST",
+                  "PARTITIONS",
               }, NormalizeName))
     {
     }

yql/essentials/sql/v1/complete/sql_complete_ut.cpp

@@ -604,6 +604,7 @@ Y_UNIT_TEST(SelectFrom) {
             {FunctionName, "FILTER()", 1},
             {FunctionName, "FOLDER()", 1},
             {FunctionName, "LIKE()", 1},
+            {FunctionName, "PARTITIONS()", 1},
             {FunctionName, "PARTITION_LIST()", 1},
             {FunctionName, "RANGE()", 1},
             {FunctionName, "REGEXP()", 1},

yql/essentials/sql/v1/query.cpp

@@ -849,6 +849,48 @@ class TPrepTableKeys: public ITableKeys {
                 return nullptr;
             }
             return Y(func.EndsWith("strict") ? "MrPartitionListStrict" : "MrPartitionList", Y("EvaluateExpr", arg.Expr));
+        } else if (func == "partitions" || func == "partitionsstrict") {
+            auto requiredLangVer = MakeLangVersion(2025, 4);
+            if (!IsBackwardCompatibleFeatureAvailable(ctx.Settings.LangVer, requiredLangVer, ctx.Settings.BackportMode)) {
+                auto str = FormatLangVersion(requiredLangVer);
+                YQL_ENSURE(str);
+                ctx.Error(Pos_) << "PARTITIONS table function is not available before language version " << *str;
+                return nullptr;
+            }
+
+            if (ctx.DiscoveryMode) {
+                ctx.Error(Pos_, TIssuesIds::YQL_NOT_ALLOWED_IN_DISCOVERY) << "PARTITIONS is not allowed in Discovery mode";
+                return nullptr;
+            }
+
+            if (Args_.size() != 2) {
+                ctx.Error(Pos_) << "PARTITIONS requires 2 arguments, but got: " << Args_.size();
+                return nullptr;
+            }
+
+            if (Args_[0].HasAt || Args_[1].HasAt) {
+                ctx.Error(Pos_) << "Temporary tables are not supported here";
+                return nullptr;
+            }
+
+            if (!Args_[1].View.empty()) {
+                ctx.Error(Pos_) << "VIEW should be used only in first argument";
+                return nullptr;
+            }
+
+            ExtractTableName(ctx, Args_[0]);
+            ExtractTableName(ctx, Args_[1]);
+            auto path = ctx.GetPrefixedPath(Service_, Cluster_, Args_[0].Id);
+            if (!path) {
+                return nullptr;
+            }
+            TNodePtr key = Y("Key", Q(Y(Q("table"), Y("String", path))));
+            key = AddView(key, Args_[0].View);
+            if (!ValidateView(GetPos(), ctx, Service_, Args_[0].View)) {
+                return nullptr;
+            }
+            TDeferredAtom pattern = Args_[1].Id;
+            return Y(func.EndsWith("strict") ? "MrPartitionsStrict" : "MrPartitions", key, pattern.Build());
         }
 
         ctx.Error(Pos_) << "Unknown table name preprocessor: " << Func_;

yt/yql/providers/yt/common/yql_names.h

@@ -55,6 +55,8 @@ const TStringBuf MrWalkFoldersImplName = "MrWalkFoldersImpl";
 const TStringBuf MrRangeInputListInternal = "MrRangeInputListInternal";
 const TStringBuf MrPartitionListName = "MrPartitionList";
 const TStringBuf MrPartitionListStrictName = "MrPartitionListStrict";
+const TStringBuf MrPartitionsName = "MrPartitions";
+const TStringBuf MrPartitionsStrictName = "MrPartitionsStrict";
 
 // YT related names
 const TStringBuf READ_SCHEMA_ATTR_NAME = "_read_schema";

yt/yql/providers/yt/provider/ya.make

@@ -23,6 +23,7 @@ SRCS(
     yql_yt_helpers.cpp
     yql_yt_intent_determination.cpp
     yql_yt_io_discovery.cpp
+    yql_yt_io_discovery_partitions.cpp
     yql_yt_io_discovery_walk_folders.cpp
     yql_yt_join_impl.cpp
     yql_yt_join_reorder.cpp
@@ -72,6 +73,7 @@ SRCS(
 )
 
 PEERDIR(
+    contrib/libs/re2
     library/cpp/yson/node
     library/cpp/json/writer
     library/cpp/json

yt/yql/providers/yt/provider/yql_yt_io_discovery.cpp

@@ -3,6 +3,7 @@
 #include "yql_yt_gateway.h"
 #include "yql_yt_op_settings.h"
 #include "yql_yt_helpers.h"
+#include "yql_yt_io_discovery_partitions.h"
 #include "yql_yt_io_discovery_walk_folders.h"
 
 #include <yt/yql/providers/yt/expr_nodes/yql_yt_expr_nodes.h>
@@ -52,6 +53,7 @@ class TYtIODiscoveryTransformer : public TGraphTransformerBase {
         const bool evaluationInProgress = State_->Types->EvaluationInProgress;
         TOptimizeExprSettings settings(nullptr);
         settings.VisitChanges = true;
+        bool seenMrPartitions = false;
         auto status = OptimizeExpr(input, output, [&](const TExprNode::TPtr& node, TExprContext& ctx) -> TExprNode::TPtr {
             if (auto maybeRead = TMaybeNode<TYtRead>(node)) {
                 if (!maybeRead.DataSource()) { // Validates provider
@@ -69,6 +71,11 @@ class TYtIODiscoveryTransformer : public TGraphTransformerBase {
                     return {};
                 }
 
+                if (read.Arg(2).Ref().IsCallable({MrPartitionsName, MrPartitionsStrictName})) {
+                    seenMrPartitions = true;
+                    return ExpandMrPartitions(read, ctx, *(State_->Types));
+                }
+
                 if (read.Arg(2).Ref().IsCallable({MrPartitionListName, MrPartitionListStrictName})) {
                     return ExpandPartitionList(read, ctx);
                 }
@@ -193,6 +200,12 @@ class TYtIODiscoveryTransformer : public TGraphTransformerBase {
         }, ctx, settings);
 
         if (status.Level != TStatus::Ok) {
+            if (seenMrPartitions && status.Level == TStatus::Repeat) {
+                ctx.Step
+                    .Repeat(TExprStep::ExpandApplyForLambdas)
+                    .Repeat(TExprStep::ExprEval);
+                return TStatus(TStatus::Repeat, true);
+            }
             return status;
         }