Fix BlockExpandChunked optimizers

4 files changed, 40 insertions, 25 deletions

diff --git a/ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp b/ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp
index d1741d53c79..46871c921ae 100644
--- a/ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp
+++ b/ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp
@@ -112,17 +112,6 @@ TExprNode::TPtr RebuildArgumentsOnlyLambdaForBlocks(const TExprNode& lambda, TEx
     return ctx.NewLambda(lambda.Pos(), ctx.NewArguments(lambda.Pos(), std::move(newArgs)), std::move(newRoots));
 }
 
-TExprNode::TPtr OptimizeWideFromBlocks(const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& types) {
-    Y_UNUSED(types);
-    if (node->Head().IsCallable("BlockExpandChunked")) {
-        // WideFromBlocks accepts chunked input
-        YQL_CLOG(DEBUG, Core) << "Drop " << node->Head().Content() << " under " << node->Content();
-        return ctx.ChangeChild(*node, 0, node->Head().HeadPtr());
-    }
-
-    return node;
-}
-
 TExprNode::TPtr OptimizeWideToBlocks(const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& types) {
     Y_UNUSED(types);
     if (node->Head().IsCallable("WideFromBlocks")) {
@@ -5254,13 +5243,11 @@ TExprNode::TPtr OptimizeWideMapBlocks(const TExprNode::TPtr& node, TExprContext&
 
     auto ret = ctx.Builder(node->Pos())
         .Callable("WideFromBlocks")
-            .Callable(0, "BlockExpandChunked")
-                .Callable(0, "WideMap")
-                    .Callable(0, "WideToBlocks")
-                        .Add(0, node->HeadPtr())
-                    .Seal()
-                    .Add(1, blockLambda)
+            .Callable(0, "WideMap")
+                .Callable(0, "WideToBlocks")
+                    .Add(0, node->HeadPtr())
                 .Seal()
+                .Add(1, blockLambda)
             .Seal()
         .Seal()
         .Build();
@@ -5288,13 +5275,11 @@ TExprNode::TPtr OptimizeWideFilterBlocks(const TExprNode::TPtr& node, TExprConte
     }
 
     auto blockMapped = ctx.Builder(node->Pos())
-        .Callable("BlockExpandChunked")
-            .Callable(0, "WideMap")
-                .Callable(0, "WideToBlocks")
-                    .Add(0, node->HeadPtr())
-                .Seal()
-                .Add(1, blockLambda)
+        .Callable("WideMap")
+            .Callable(0, "WideToBlocks")
+                .Add(0, node->HeadPtr())
             .Seal()
+            .Add(1, blockLambda)
         .Seal()
         .Build();
 
@@ -7309,7 +7294,6 @@ struct TPeepHoleRules {
         {"WideMap", &OptimizeWideMapBlocks},
         {"NarrowMap", &OptimizeWideMapBlocks},
         {"WideFilter", &OptimizeWideFilterBlocks},
-        {"WideFromBlocks", &OptimizeWideFromBlocks},
         {"WideToBlocks", &OptimizeWideToBlocks},
         {"Skip", &OptimizeSkipTakeToBlocks},
         {"Take", &OptimizeSkipTakeToBlocks},
diff --git a/ydb/library/yql/core/yql_expr_type_annotation.cpp b/ydb/library/yql/core/yql_expr_type_annotation.cpp
index bc3c46f512a..4ea50581b24 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.cpp
+++ b/ydb/library/yql/core/yql_expr_type_annotation.cpp
@@ -2792,6 +2792,32 @@ bool EnsureWideStreamType(TPositionHandle position, const TTypeAnnotationNode& t
     return true;
 }
 
+bool IsWideBlockType(const TTypeAnnotationNode& type) {
+    if (type.GetKind() != ETypeAnnotationKind::Multi) {
+        return false;
+    }
+
+    const auto& items = type.Cast<TMultiExprType>()->GetItems();
+    if (items.empty()) {
+        return false;
+    }
+
+    if (!AllOf(items, [](const auto& item){ return item->IsBlockOrScalar(); })) {
+        return false;
+    }
+
+    if (items.back()->GetKind() != ETypeAnnotationKind::Scalar) {
+        return false;
+    }
+
+    auto blockLenType = items.back()->Cast<TScalarExprType>()->GetItemType();
+    if (blockLenType->GetKind() != ETypeAnnotationKind::Data) {
+        return false;
+    }
+
+    return blockLenType->Cast<TDataExprType>()->GetSlot() == EDataSlot::Uint64;
+}
+
 bool EnsureWideBlockType(TPositionHandle position, const TTypeAnnotationNode& type, TTypeAnnotationNode::TListType& blockItemTypes, TExprContext& ctx, bool allowScalar) {
     if (HasError(&type, ctx)) {
         return false;
diff --git a/ydb/library/yql/core/yql_expr_type_annotation.h b/ydb/library/yql/core/yql_expr_type_annotation.h
index f2b6430ff55..453cb694af8 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.h
+++ b/ydb/library/yql/core/yql_expr_type_annotation.h
@@ -124,6 +124,7 @@ bool EnsureWideFlowType(const TExprNode& node, TExprContext& ctx);
 bool EnsureWideFlowType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx);
 bool EnsureWideStreamType(const TExprNode& node, TExprContext& ctx);
 bool EnsureWideStreamType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx);
+bool IsWideBlockType(const TTypeAnnotationNode& type);
 bool EnsureWideBlockType(TPositionHandle position, const TTypeAnnotationNode& type, TTypeAnnotationNode::TListType& blockItemTypes, TExprContext& ctx, bool allowScalar = true);
 bool EnsureWideFlowBlockType(const TExprNode& node, TTypeAnnotationNode::TListType& blockItemTypes, TExprContext& ctx, bool allowScalar = true);
 bool EnsureWideStreamBlockType(const TExprNode& node, TTypeAnnotationNode::TListType& blockItemTypes, TExprContext& ctx, bool allowScalar = true);
diff --git a/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp b/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp
index 31d0518f861..acc08c3542a 100644
--- a/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp
+++ b/ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp
@@ -675,7 +675,11 @@ TMkqlCommonCallableCompiler::TShared::TShared() {
     AddCallable("WideMap", [](const TExprNode& node, TMkqlBuildContext& ctx) {
         const auto arg = MkqlBuildExpr(node.Head(), ctx);
         const auto lambda = [&](TRuntimeNode::TList items) { return MkqlBuildWideLambda(node.Tail(), ctx, items); };
-        return ctx.ProgramBuilder.WideMap(arg, lambda);
+        TRuntimeNode result = ctx.ProgramBuilder.WideMap(arg, lambda);
+        if (IsWideBlockType(*node.GetTypeAnn()->Cast<TFlowExprType>()->GetItemType())) {
+            result = ctx.ProgramBuilder.BlockExpandChunked(result);
+        }
+        return result;
     });
 
     AddCallable("WideChain1Map", [](const TExprNode& node, TMkqlBuildContext& ctx) {