Imported PG clickbench queries

11 files changed, 89 insertions, 16 deletions

diff --git a/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp b/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
index 9747c537e64..7b6dbd9dfaf 100644
--- a/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
+++ b/ydb/library/yql/core/common_opt/yql_co_pgselect.cpp
@@ -1707,6 +1707,16 @@ TExprNode::TPtr BuildAggregationTraits(TPositionHandle pos, bool onWindow, const
         extractor = ctx.NewLambda(pos, std::move(arguments), std::move(aggFuncArgs));
     }
 
+    if (optCtx.Types->PgEmitAggApply && !onWindow) {
+        return ctx.Builder(pos)
+            .Callable("AggApply")
+                .Atom(0, TString("pg_") + func)
+                .Add(1, type)
+                .Add(2, extractor)
+            .Seal()
+            .Build();
+    }
+
     return ctx.Builder(pos)
         .Callable(TString(onWindow ? "PgWindowTraits" : "PgAggregationTraits") + (distinctColumnName ? "Tuple" : ""))
             .Atom(0, func)
diff --git a/ydb/library/yql/core/type_ann/type_ann_list.cpp b/ydb/library/yql/core/type_ann/type_ann_list.cpp
index 6a411fe8080..de61c4d41aa 100644
--- a/ydb/library/yql/core/type_ann/type_ann_list.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_list.cpp
@@ -5337,6 +5337,11 @@ namespace {
         } else if (name == "some") {
             input->SetTypeAnn(lambda->GetTypeAnn());
         } else if (name.StartsWith("pg_")) {
+            if (lambda->GetTypeAnn()->GetKind() == ETypeAnnotationKind::Tuple) {
+                lambda = ConvertToMultiLambda(lambda, ctx.Expr);
+                return IGraphTransformer::TStatus::Repeat;
+            }
+
             auto func = name;
             func.SkipPrefix("pg_");
             TVector<ui32> argTypes;
diff --git a/ydb/library/yql/core/type_ann/type_ann_pg.cpp b/ydb/library/yql/core/type_ann/type_ann_pg.cpp
index 77acf647e43..1aa7c0591df 100644
--- a/ydb/library/yql/core/type_ann/type_ann_pg.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_pg.cpp
@@ -1056,21 +1056,7 @@ IGraphTransformer::TStatus PgAggregationTraitsWrapper(const TExprNode::TPtr& inp
         // convert lambda with tuple type into multi lambda
         auto args = input->ChildrenList();
         if (lambda->GetTypeAnn()->GetKind() == ETypeAnnotationKind::Tuple) {
-            Y_ENSURE(lambda->ChildrenSize() == 2);
-            auto tupleTypeSize = lambda->GetTypeAnn()->Cast<TTupleExprType>()->GetSize();
-            auto newArg = ctx.Expr.NewArgument(lambda->Pos(), "row");
-            auto newBody = ctx.Expr.ReplaceNode(lambda->TailPtr(), lambda->Head().Head(), newArg);
-            TExprNode::TListType bodies;
-            for (ui32 i = 0; i < tupleTypeSize; ++i) {
-                bodies.push_back(ctx.Expr.Builder(lambda->Pos())
-                    .Callable("Nth")
-                        .Add(0, newBody)
-                        .Atom(1, ToString(i))
-                    .Seal()
-                    .Build());
-            }
-
-            args[2] = ctx.Expr.NewLambda(lambda->Pos(), ctx.Expr.NewArguments(lambda->Pos(), { newArg }), std::move(bodies));
+            args[2] = ConvertToMultiLambda(lambda, ctx.Expr);
         }
 
         output = ctx.Expr.NewCallable(input->Pos(), input->Content().substr(0, input->Content().Size() - 5), std::move(args));
diff --git a/ydb/library/yql/core/yql_expr_type_annotation.cpp b/ydb/library/yql/core/yql_expr_type_annotation.cpp
index dd834c6b94b..da3f156efc9 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.cpp
+++ b/ydb/library/yql/core/yql_expr_type_annotation.cpp
@@ -6087,5 +6087,23 @@ bool ApplyOriginalType(TExprNode::TPtr input, bool isMany, const TTypeAnnotation
     return true;
 }
 
+TExprNode::TPtr ConvertToMultiLambda(const TExprNode::TPtr& lambda, TExprContext& ctx) {
+    Y_ENSURE(lambda->ChildrenSize() == 2);
+    auto tupleTypeSize = lambda->GetTypeAnn()->Cast<TTupleExprType>()->GetSize();
+    auto newArg = ctx.NewArgument(lambda->Pos(), "row");
+    auto newBody = ctx.ReplaceNode(lambda->TailPtr(), lambda->Head().Head(), newArg);
+    TExprNode::TListType bodies;
+    for (ui32 i = 0; i < tupleTypeSize; ++i) {
+        bodies.push_back(ctx.Builder(lambda->Pos())
+            .Callable("Nth")
+                .Add(0, newBody)
+                .Atom(1, ToString(i))
+            .Seal()
+            .Build());
+    }
+
+    return ctx.NewLambda(lambda->Pos(), ctx.NewArguments(lambda->Pos(), { newArg }), std::move(bodies));
+}
+
 
 } // NYql
diff --git a/ydb/library/yql/core/yql_expr_type_annotation.h b/ydb/library/yql/core/yql_expr_type_annotation.h
index a2803b7c57d..3c935c7f9b7 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.h
+++ b/ydb/library/yql/core/yql_expr_type_annotation.h
@@ -326,6 +326,7 @@ TExprNode::TPtr ExpandPgAggregationTraits(TPositionHandle pos, const NPg::TAggre
 
 const TTypeAnnotationNode* GetOriginalResultType(TPositionHandle pos, bool isMany, const TTypeAnnotationNode* originalExtractorType, TExprContext& ctx);
 bool ApplyOriginalType(TExprNode::TPtr input, bool isMany, const TTypeAnnotationNode* originalExtractorType, TExprContext& ctx);
+TExprNode::TPtr ConvertToMultiLambda(const TExprNode::TPtr& lambda, TExprContext& ctx);
 
 const TStringBuf BlockLengthColumnName = "_yql_block_length";
 }
diff --git a/ydb/library/yql/core/yql_type_annotation.h b/ydb/library/yql/core/yql_type_annotation.h
index 863f8e2bd82..01c388d1b3d 100644
--- a/ydb/library/yql/core/yql_type_annotation.h
+++ b/ydb/library/yql/core/yql_type_annotation.h
@@ -230,6 +230,7 @@ struct TTypeAnnotationContext: public TThrRefBase {
     bool YsonCastToString = true;
     ui32 FolderSubDirsLimit = 1000;
     bool UseBlocks = false;
+    bool PgEmitAggApply = false;
     IArrowResolver::TPtr ArrowResolver;
 
     // compatibility with v0 or raw s-expression code
diff --git a/ydb/library/yql/minikql/mkql_type_builder.cpp b/ydb/library/yql/minikql/mkql_type_builder.cpp
index a723a5b4208..e8e7af4bf25 100644
--- a/ydb/library/yql/minikql/mkql_type_builder.cpp
+++ b/ydb/library/yql/minikql/mkql_type_builder.cpp
@@ -1501,6 +1501,11 @@ bool ConvertArrowType(TType* itemType, std::shared_ptr<arrow::DataType>& type) {
             currentType = AS_TYPE(TOptionalType, currentType)->GetItemType();
         } while (currentType->IsOptional());
 
+        if (currentType->IsPg()) {
+            previousType = currentType;
+            ++nestLevel;
+        }
+
         // previousType is always Optional
         std::shared_ptr<arrow::DataType> innerArrowType;
         if (!ConvertArrowType(previousType, innerArrowType)) {
diff --git a/ydb/library/yql/providers/config/yql_config_provider.cpp b/ydb/library/yql/providers/config/yql_config_provider.cpp
index aa81d82bbf7..7e4dab2cc11 100644
--- a/ydb/library/yql/providers/config/yql_config_provider.cpp
+++ b/ydb/library/yql/providers/config/yql_config_provider.cpp
@@ -773,6 +773,14 @@ namespace {
                     Types.OptLLVM = "OFF";
                 }
             }
+            else if (name == "PgEmitAggApply" || name == "DisablePgEmitAggApply") {
+                if (args.size() != 0) {
+                    ctx.AddError(TIssue(pos, TStringBuilder() << "Expected no arguments, but got " << args.size()));
+                    return false;
+                }
+
+                Types.PgEmitAggApply = (name == "PgEmitAggApply");
+            }
             else {
                 ctx.AddError(TIssue(pos, TStringBuilder() << "Unsupported command: " << name));
                 return false;
diff --git a/ydb/library/yql/public/udf/arrow/block_builder.h b/ydb/library/yql/public/udf/arrow/block_builder.h
index acbeff59905..c82f60b7eb0 100644
--- a/ydb/library/yql/public/udf/arrow/block_builder.h
+++ b/ydb/library/yql/public/udf/arrow/block_builder.h
@@ -1178,6 +1178,11 @@ inline std::unique_ptr<TArrayBuilderBase> MakeArrayBuilderBase(
             }
         }
 
+        if (TPgTypeInspector(typeInfoHelper, currentType)) {
+            previousType = currentType;
+            ++nestLevel;
+        }
+
         auto builder = MakeArrayBuilderBase(typeInfoHelper, previousType, pool, maxBlockLength, pgBuilder);
         for (ui32 i = 1; i < nestLevel; ++i) {
             builder = std::make_unique<TExternalOptionalArrayBuilder>(typeInfoHelper, types[nestLevel - 1 - i], pool, maxBlockLength, std::move(builder));
diff --git a/ydb/library/yql/public/udf/arrow/block_reader.h b/ydb/library/yql/public/udf/arrow/block_reader.h
index e79f11ee206..63fb643c6e7 100644
--- a/ydb/library/yql/public/udf/arrow/block_reader.h
+++ b/ydb/library/yql/public/udf/arrow/block_reader.h
@@ -328,6 +328,11 @@ std::unique_ptr<typename TTraits::TResult> MakeBlockReaderImpl(const ITypeInfoHe
             }
         }
 
+        if (TPgTypeInspector(typeInfoHelper, currentType)) {
+            previousType = currentType;
+            ++nestLevel;
+        }
+
         auto reader = MakeBlockReaderImpl<TTraits>(typeInfoHelper, previousType, pgBuilder);
         for (ui32 i = 1; i < nestLevel; ++i) {
             reader = std::make_unique<typename TTraits::TExtOptional>(std::move(reader));
@@ -437,6 +442,19 @@ inline void UpdateBlockItemSerializeProps(const ITypeInfoHelper& typeInfoHelper,
         return;
     }
 
+    TPgTypeInspector typePg(typeInfoHelper, type);
+    if (typePg) {
+        auto desc = typeInfoHelper.FindPgTypeDescription(typePg.GetTypeId());
+        if (desc->PassByValue) {
+            *props.MaxSize += desc->Typelen;
+        } else {
+            props.MaxSize = {};
+            props.IsFixed = false;
+        }
+
+        return;
+    }
+
     Y_ENSURE(false, "Unsupported type");
 }
 
diff --git a/ydb/library/yql/sql/pg/pg_sql.cpp b/ydb/library/yql/sql/pg/pg_sql.cpp
index 70888b6f885..0f8ca060d6b 100644
--- a/ydb/library/yql/sql/pg/pg_sql.cpp
+++ b/ydb/library/yql/sql/pg/pg_sql.cpp
@@ -1372,7 +1372,23 @@ public:
         }
 
         auto name = to_lower(TString(value->name));
-        if (name == "dqengine") {
+        if (name == "useblocks" || name == "emitaggapply") {
+            if (ListLength(value->args) != 1) {
+                AddError(TStringBuilder() << "VariableSetStmt, expected 1 arg, but got: " << ListLength(value->args));
+                return nullptr;
+            }
+
+            auto arg = ListNodeNth(value->args, 0);
+            if (NodeTag(arg) == T_A_Const && (NodeTag(CAST_NODE(A_Const, arg)->val) == T_String)) {
+                auto rawStr = StrVal(CAST_NODE(A_Const, arg)->val);
+                auto configSource = L(A("DataSource"), QA(TString(NYql::ConfigProviderName)));
+                Statements.push_back(L(A("let"), A("world"), L(A(TString(NYql::ConfigureName)), A("world"), configSource,
+                    QA(TString(rawStr == "true" ? "" : "Disable") + TString((name == "useblocks") ? "UseBlocks" : "PgEmitAggApply")))));
+            } else {
+                AddError(TStringBuilder() << "VariableSetStmt, expected string literal for " << value->name << " option");
+                return nullptr;
+            }
+        } else if (name == "dqengine") {
             if (ListLength(value->args) != 1) {
                 AddError(TStringBuilder() << "VariableSetStmt, expected 1 arg, but got: " << ListLength(value->args));
                 return nullptr;