Block input for YT MapReduce map stage

commit_hash:36f9bfe1cdd4832a2f7365fc082796dca67d8b01

11 files changed, 113 insertions, 49 deletions

diff --git a/yt/yql/providers/yt/gateway/file/yql_yt_file_mkql_compiler.cpp b/yt/yql/providers/yt/gateway/file/yql_yt_file_mkql_compiler.cpp
index 89525fcbc56..a3dea494c84 100644
--- a/yt/yql/providers/yt/gateway/file/yql_yt_file_mkql_compiler.cpp
+++ b/yt/yql/providers/yt/gateway/file/yql_yt_file_mkql_compiler.cpp
@@ -916,8 +916,16 @@ void RegisterYtFileMkqlCompilers(NCommon::TMkqlCallableCompilerBase& compiler) {
                 values = arg->GetTypeAnn()->GetKind() == ETypeAnnotationKind::Flow ?
                     ctx.ProgramBuilder.ToFlow(values) : ctx.ProgramBuilder.Iterator(values, {});
 
-                if (ETypeAnnotationKind::Multi == GetSeqItemType(*ytMapReduce.Mapper().Cast<TCoLambda>().Args().Arg(0).Ref().GetTypeAnn()).GetKind())
+                auto& lambdaInputType = GetSeqItemType(*ytMapReduce.Mapper().Cast<TCoLambda>().Args().Arg(0).Ref().GetTypeAnn());
+                if (lambdaInputType.GetKind() == ETypeAnnotationKind::Multi) {
                     values = ExpandFlow(values, ctx);
+                }
+
+                if (IsWideBlockType(lambdaInputType)) {
+                    values = ctx.ProgramBuilder.ToFlow(
+                        ctx.ProgramBuilder.WideToBlocks(
+                            ctx.ProgramBuilder.FromFlow(values)));
+                }
 
                 NCommon::TMkqlBuildContext innerCtx(ctx, {{arg, values}}, ytMapReduce.Mapper().Ref().UniqueId());
 
diff --git a/yt/yql/providers/yt/gateway/native/yql_yt_native.cpp b/yt/yql/providers/yt/gateway/native/yql_yt_native.cpp
index d0d38b994e9..be8ce859a7e 100644
--- a/yt/yql/providers/yt/gateway/native/yql_yt_native.cpp
+++ b/yt/yql/providers/yt/gateway/native/yql_yt_native.cpp
@@ -4131,6 +4131,7 @@ private:
         const TString& mapInputType,
         size_t mapDirectOutputs,
         const TExpressionResorceUsage& mapExtraUsage,
+        bool mapBlockInput,
         TString reduceLambda,
         const TString& reduceInputType,
         const TExpressionResorceUsage& reduceExtraUsage,
@@ -4143,7 +4144,7 @@ private:
         const bool testRun = execCtx->Config_->GetLocalChainTest();
         TFuture<bool> ret = testRun ? MakeFuture<bool>(false) : execCtx->LookupQueryCacheAsync();
         return ret.Apply([reduceBy, sortBy, limit, sortLimitBy, mapLambda, mapInputType, mapDirectOutputs,
-                          mapExtraUsage, reduceLambda, reduceInputType, reduceExtraUsage,
+                          mapExtraUsage, mapBlockInput, reduceLambda, reduceInputType, reduceExtraUsage,
                           intermediateMeta, intermediateSchema, intermediateStreams, inputQueryExpr, execCtx, testRun]
                          (const auto& f) mutable
         {
@@ -4250,6 +4251,7 @@ private:
             mapJob->SetTableNames(tables);
             mapJob->SetRowOffsets(rowOffsets);
             mapJob->SetUseSkiff(useSkiff, TMkqlIOSpecs::ESystemField::RowIndex);
+            mapJob->SetUseBlockInput(mapBlockInput);
             mapJob->SetYamrInput(execCtx->YamrInput);
 
             auto reduceJob = MakeIntrusive<TYqlUserJob>();
@@ -4532,6 +4534,7 @@ private:
         const bool useNativeTypes = execCtx->Options_.Config()->UseNativeYtTypes.Get().GetOrElse(DEFAULT_USE_NATIVE_YT_TYPES);
         const auto nativeTypeCompat = execCtx->Options_.Config()->NativeYtTypeCompatibility.Get(execCtx->Cluster_).GetOrElse(NTCF_LEGACY);
         const auto useIntermediateStreams = execCtx->Options_.Config()->UseIntermediateStreams.Get().GetOrElse(DEFAULT_USE_INTERMEDIATE_STREAMS);
+        const bool mapBlockInput = NYql::HasSetting(mapReduce.Settings().Ref(), EYtSettingType::BlockInputApplied);
 
         NYT::TNode intermediateMeta;
         NYT::TNode intermediateSchema;
@@ -4637,13 +4640,13 @@ private:
 
         const TString inputQueryExpr = GenerateInputQueryWhereExpression(mapReduce.Settings().Ref());
 
-        return execCtx->Session_->Queue_->Async([reduceBy, sortBy, limit, sortLimitBy, mapLambda, mapInputType, mapDirectOutputs, mapExtraUsage,
+        return execCtx->Session_->Queue_->Async([reduceBy, sortBy, limit, sortLimitBy, mapLambda, mapInputType, mapDirectOutputs, mapExtraUsage, mapBlockInput,
             reduceLambda, reduceInputType, reduceExtraUsage, intermediateMeta, intermediateSchema, intermediateStreams, useIntermediateStreams, inputQueryExpr, execCtx]()
         {
             YQL_LOG_CTX_ROOT_SESSION_SCOPE(execCtx->LogCtx_);
             execCtx->MakeUserFiles();
             if (mapLambda) {
-                return ExecMapReduce(reduceBy, sortBy, limit, sortLimitBy, mapLambda, mapInputType, mapDirectOutputs, mapExtraUsage,
+                return ExecMapReduce(reduceBy, sortBy, limit, sortLimitBy, mapLambda, mapInputType, mapDirectOutputs, mapExtraUsage, mapBlockInput,
                     reduceLambda, reduceInputType, reduceExtraUsage, intermediateMeta, intermediateSchema, intermediateStreams, inputQueryExpr, execCtx);
             } else {
                 return ExecMapReduce(reduceBy, sortBy, limit, sortLimitBy, reduceLambda, reduceInputType, reduceExtraUsage, intermediateSchema,
diff --git a/yt/yql/providers/yt/provider/yql_yt_block_input.cpp b/yt/yql/providers/yt/provider/yql_yt_block_input.cpp
index bda09645b41..5f0f3e07396 100644
--- a/yt/yql/providers/yt/provider/yql_yt_block_input.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_block_input.cpp
@@ -2,8 +2,10 @@
 
 #include <yql/essentials/core/yql_opt_utils.h>
 #include <yql/essentials/providers/common/transform/yql_optimize.h>
-#include <yt/yql/providers/yt/common/yql_names.h>
 #include <yql/essentials/utils/log/log.h>
+#include <yt/yql/providers/yt/common/yql_names.h>
+#include <yt/yql/providers/yt/provider/yql_yt_block_io_utils.h>
+#include <yt/yql/providers/yt/provider/yql_yt_helpers.h>
 
 namespace NYql {
 
@@ -22,54 +24,48 @@ public:
         , State_(std::move(state))
     {
 #define HNDL(name) "YtBlockInput-"#name, Hndl(&TYtBlockInputTransformer::name)
-        AddHandler(0, &TYtMap::Match, HNDL(TryTransformMap));
+        AddHandler(0, &TYtMap::Match, HNDL(TryTransformMap<TYtMap>));
+        AddHandler(0, &TYtMapReduce::Match, HNDL(TryTransformMap<TYtMapReduce>));
         AddHandler(0, &TYtTableContent::Match, HNDL(TryTransformTableContent));
 #undef HNDL
     }
 
 private:
+    template<typename TYtOpWithMap>
     TMaybeNode<TExprBase> TryTransformMap(TExprBase node, TExprContext& ctx) const {
-        auto map = node.Cast<TYtMap>();
+        auto op = node.Cast<TYtOpWithMap>();
         if (
-            NYql::HasSetting(map.Settings().Ref(), EYtSettingType::BlockInputApplied)
-            || !NYql::HasSetting(map.Settings().Ref(), EYtSettingType::BlockInputReady)
-            || !CanRewriteMap(map, ctx)
+            NYql::HasSetting(op.Settings().Ref(), EYtSettingType::BlockInputApplied)
+            || !NYql::HasSetting(op.Settings().Ref(), EYtSettingType::BlockInputReady)
+            || !CanRewriteMap(op, ctx)
         ) {
-            return map;
+            return op;
         }
 
-        YQL_CLOG(INFO, ProviderYt) << "Rewrite YtMap with block input";
+        YQL_CLOG(INFO, ProviderYt) << "Rewrite " << TYtOpWithMap::CallableName() << " with block input";
 
-        auto settings = RemoveSetting(map.Settings().Ref(), EYtSettingType::BlockInputReady, ctx);
+        auto settings = RemoveSetting(op.Settings().Ref(), EYtSettingType::BlockInputReady, ctx);
         settings = AddSetting(*settings, EYtSettingType::BlockInputApplied, TExprNode::TPtr(), ctx);
-        auto mapperLambda = Build<TCoLambda>(ctx, map.Mapper().Pos())
-            .Args({"flow"})
-            .Body<TExprApplier>()
-                .Apply(map.Mapper())
-                .With<TCoToFlow>(0)
-                    .Input<TCoWideFromBlocks>()
-                        .Input<TCoFromFlow>()
-                            .Input("flow")
-                        .Build()
-                    .Build()
-                .Build()
-            .Build()
-            .Done()
-            .Ptr();
 
-        return Build<TYtMap>(ctx, node.Pos())
-            .InitFrom(map)
+        auto mapperLambda = WrapLambdaWithBlockInput(op.Mapper().template Cast<TCoLambda>(), ctx);
+        return Build<TYtOpWithMap>(ctx, node.Pos())
+            .InitFrom(op)
             .Settings(settings)
             .Mapper(mapperLambda)
             .Done();
     }
 
-    bool CanRewriteMap(const TYtMap& map, TExprContext& ctx) const {
-        if (auto flowSetting = NYql::GetSetting(map.Settings().Ref(), EYtSettingType::Flow); !flowSetting || flowSetting->ChildrenSize() < 2) {
+    bool CanRewriteMap(const TYtWithUserJobsOpBase& op, TExprContext& ctx) const {
+        auto mapLambda = GetMapLambda(op);
+        if (!mapLambda) {
+            return false;
+        }
+
+        if (auto flowSetting = NYql::GetSetting(op.Settings().Ref(), EYtSettingType::Flow); !flowSetting || flowSetting->ChildrenSize() < 2) {
             return false;
         }
 
-        return EnsureWideFlowType(map.Mapper().Args().Arg(0).Ref(), ctx);
+        return EnsureWideFlowType(mapLambda.Cast().Args().Arg(0).Ref(), ctx);
     }
 
     TMaybeNode<TExprBase> TryTransformTableContent(TExprBase node, TExprContext& ctx, const TGetParents& getParents) const {
diff --git a/yt/yql/providers/yt/provider/yql_yt_block_io_filter.cpp b/yt/yql/providers/yt/provider/yql_yt_block_io_filter.cpp
index 3a1e05d9aa8..6d1cc92b4d7 100644
--- a/yt/yql/providers/yt/provider/yql_yt_block_io_filter.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_block_io_filter.cpp
@@ -21,8 +21,9 @@ public:
         , Finalizer_(std::move(finalizer))
     {
 #define HNDL(name) "YtBlockIOFilter-"#name, Hndl(&YtBlockIOFilterTransformer::name)
-        AddHandler(0, &TYtMap::Match, HNDL(HandleMapInput));
+        AddHandler(0, &TYtMap::Match, HNDL(HandleMapInput<TYtMap>));
         AddHandler(0, &TYtMap::Match, HNDL(HandleMapOutput));
+        AddHandler(0, &TYtMapReduce::Match, HNDL(HandleMapInput<TYtMapReduce>));
         AddHandler(0, &TYtTableContent::Match, HNDL(HandleTableContent));
 #undef HNDL
     }
@@ -40,28 +41,29 @@ private:
         TOptimizeTransformerBase::Rewind();
     }
 
+    template<typename TYtOpWithMap>
     TMaybeNode<TExprBase> HandleMapInput(TExprBase node, TExprContext& ctx) const {
-        auto map = node.Cast<TYtMap>();
-        if (NYql::HasSetting(map.Settings().Ref(), EYtSettingType::BlockInputApplied)) {
-            return map;
+        auto op = node.Cast<TYtOpWithMap>();
+        if (NYql::HasSetting(op.Settings().Ref(), EYtSettingType::BlockInputApplied)) {
+            return op;
         }
 
         if (!State_->Configuration->JobBlockInput.Get().GetOrElse(Types->UseBlocks)) {
-            return map;
+            return op;
         }
 
-        auto settings = map.Settings().Ptr();
-        bool canUseBlockInput = CanUseBlockInputForMap(map);
+        auto settings = op.Settings().Ptr();
+        bool canUseBlockInput = CanUseBlockInputForMap(op);
         bool hasSetting = HasSetting(*settings, EYtSettingType::BlockInputReady);
         if (canUseBlockInput && !hasSetting) {
             settings = AddSetting(*settings, EYtSettingType::BlockInputReady, TExprNode::TPtr(), ctx);
         } else if (!canUseBlockInput && hasSetting) {
             settings = RemoveSetting(*settings, EYtSettingType::BlockInputReady, ctx);
         } else {
-            return map;
+            return op;
         }
-        return Build<TYtMap>(ctx, node.Pos())
-            .InitFrom(map)
+        return Build<TYtOpWithMap>(ctx, node.Pos())
+            .InitFrom(op)
             .Settings(settings)
             .Done();
     }
@@ -98,16 +100,21 @@ private:
             .Done();
     }
 
-    bool CanUseBlockInputForMap(const TYtMap& map) const {
-        if (!NYql::HasSetting(map.Settings().Ref(), EYtSettingType::Flow)) {
+    bool CanUseBlockInputForMap(const TYtWithUserJobsOpBase& op) const {
+        auto mapLambda = GetMapLambda(op);
+        if (!mapLambda) {
+            return false;
+        }
+
+        if (!NYql::HasSetting(op.Settings().Ref(), EYtSettingType::Flow)) {
             return false;
         }
 
-        if (map.Input().Size() > 1) {
+        if (op.Input().Size() > 1) {
             return false;
         }
 
-        for (auto path : map.Input().Item(0).Paths()) {
+        for (auto path : op.Input().Item(0).Paths()) {
             if (!IsYtTableSuitableForArrowInput(path.Table(), [](const TString&) {})) {
                 return false;
             }
@@ -117,7 +124,7 @@ private:
         auto supportedTypes = State_->Configuration->JobBlockInputSupportedTypes.Get().GetOrElse(DEFAULT_BLOCK_INPUT_SUPPORTED_TYPES);
         auto supportedDataTypes = State_->Configuration->JobBlockInputSupportedDataTypes.Get().GetOrElse(DEFAULT_BLOCK_INPUT_SUPPORTED_DATA_TYPES);
 
-        auto lambdaInputType = map.Mapper().Args().Arg(0).Ref().GetTypeAnn();
+        auto lambdaInputType = mapLambda.Cast().Args().Arg(0).Ref().GetTypeAnn();
         if (!CheckBlockIOSupportedTypes(*lambdaInputType, supportedTypes, supportedDataTypes, [](const TString&) {}, wideFlowLimit)) {
             return false;
         }
diff --git a/yt/yql/providers/yt/provider/yql_yt_block_io_utils.cpp b/yt/yql/providers/yt/provider/yql_yt_block_io_utils.cpp
index 639fa14ca01..17fbabe3f69 100644
--- a/yt/yql/providers/yt/provider/yql_yt_block_io_utils.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_block_io_utils.cpp
@@ -5,6 +5,8 @@
 
 namespace NYql {
 
+using namespace NNodes;
+
 bool CheckBlockIOSupportedTypes(
     const TTypeAnnotationNode& type,
     const TSet<TString>& supportedTypes,
@@ -42,4 +44,20 @@ bool CheckBlockIOSupportedTypes(
     return true;
 }
 
+TCoLambda WrapLambdaWithBlockInput(TCoLambda lambda, TExprContext& ctx) {
+    return Build<TCoLambda>(ctx, lambda.Pos())
+        .Args({"flow"})
+        .Body<TExprApplier>()
+            .Apply(lambda)
+            .With<TCoToFlow>(0)
+                .Input<TCoWideFromBlocks>()
+                    .Input<TCoFromFlow>()
+                        .Input("flow")
+                    .Build()
+                .Build()
+            .Build()
+        .Build()
+        .Done();
+}
+
 }
diff --git a/yt/yql/providers/yt/provider/yql_yt_block_io_utils.h b/yt/yql/providers/yt/provider/yql_yt_block_io_utils.h
index dc72518fe6b..0d55ffbfdea 100644
--- a/yt/yql/providers/yt/provider/yql_yt_block_io_utils.h
+++ b/yt/yql/providers/yt/provider/yql_yt_block_io_utils.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <yql/essentials/ast/yql_expr.h>
+#include <yql/essentials/core/expr_nodes/yql_expr_nodes.h>
 
 namespace NYql {
 
@@ -13,4 +14,6 @@ bool CheckBlockIOSupportedTypes(
     bool allowNestedOptionals = true
 );
 
+NNodes::TCoLambda WrapLambdaWithBlockInput(NNodes::TCoLambda lambda, TExprContext& ctx);
+
 }
diff --git a/yt/yql/providers/yt/provider/yql_yt_datasink_type_ann.cpp b/yt/yql/providers/yt/provider/yql_yt_datasink_type_ann.cpp
index 8f899ae1107..09f54d55883 100644
--- a/yt/yql/providers/yt/provider/yql_yt_datasink_type_ann.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_datasink_type_ann.cpp
@@ -1294,7 +1294,7 @@ private:
 
         auto mapReduce = TYtMapReduce(input);
 
-        const auto acceptedSettings = EYtSettingType::ReduceBy
+        auto acceptedSettings = EYtSettingType::ReduceBy
             | EYtSettingType::ReduceFilterBy
             | EYtSettingType::SortBy
             | EYtSettingType::Limit
@@ -1306,6 +1306,10 @@ private:
             | EYtSettingType::ReduceInputType
             | EYtSettingType::NoDq
             | EYtSettingType::QLFilter;
+
+        if (hasMapLambda) {
+            acceptedSettings |= EYtSettingType::BlockInputReady | EYtSettingType::BlockInputApplied;
+        }
         if (!ValidateSettings(mapReduce.Settings().Ref(), acceptedSettings, ctx)) {
             return TStatus::Error;
         }
@@ -1343,11 +1347,12 @@ private:
 
         auto itemType = GetInputItemType(mapReduce.Input(), ctx);
         const auto useFlow = NYql::GetSetting(mapReduce.Settings().Ref(), EYtSettingType::Flow);
+        const auto blockInputApplied = NYql::GetSetting(mapReduce.Settings().Ref(), EYtSettingType::BlockInputApplied);
 
         auto& mapLambda = input->ChildRef(TYtMapReduce::idx_Mapper);
         TTypeAnnotationNode::TListType mapDirectOutputTypes;
         if (hasMapLambda) {
-            const auto mapLambdaInputType = MakeInputType(itemType, useFlow, TExprNode::TPtr(), ctx);
+            const auto mapLambdaInputType = MakeInputType(itemType, useFlow, blockInputApplied, ctx);
 
             if (!UpdateLambdaAllArgumentsTypes(mapLambda, {mapLambdaInputType}, ctx)) {
                 return TStatus::Error;
diff --git a/yt/yql/providers/yt/provider/yql_yt_helpers.cpp b/yt/yql/providers/yt/provider/yql_yt_helpers.cpp
index 67f6dd127c4..4af9b954c11 100644
--- a/yt/yql/providers/yt/provider/yql_yt_helpers.cpp
+++ b/yt/yql/providers/yt/provider/yql_yt_helpers.cpp
@@ -2328,4 +2328,14 @@ bool IsYtTableSuitableForArrowInput(NNodes::TExprBase tableNode, std::function<v
     return true;
 }
 
+TMaybeNode<TCoLambda> GetMapLambda(const TYtWithUserJobsOpBase& op) {
+    if (auto map = op.Maybe<TYtMap>()) {
+        return map.Cast().Mapper();
+    } else if (auto maybeLambda = op.Maybe<TYtMapReduce>().Mapper().Maybe<TCoLambda>()) {
+        return maybeLambda.Cast();
+    }
+
+    return {};
+}
+
 } // NYql
diff --git a/yt/yql/providers/yt/provider/yql_yt_helpers.h b/yt/yql/providers/yt/provider/yql_yt_helpers.h
index 3a821b16d91..53e9a071d5e 100644
--- a/yt/yql/providers/yt/provider/yql_yt_helpers.h
+++ b/yt/yql/providers/yt/provider/yql_yt_helpers.h
@@ -143,4 +143,6 @@ bool HasYtRowNumber(const TExprNode& node);
 
 bool IsYtTableSuitableForArrowInput(NNodes::TExprBase table, std::function<void(const TString&)> unsupportedHandler);
 
+NNodes::TMaybeNode<NNodes::TCoLambda> GetMapLambda(const NNodes::TYtWithUserJobsOpBase& op);
+
 }
diff --git a/yt/yql/tests/sql/suites/blocks/block_input_mapreduce.cfg b/yt/yql/tests/sql/suites/blocks/block_input_mapreduce.cfg
new file mode 100644
index 00000000000..426917836ba
--- /dev/null
+++ b/yt/yql/tests/sql/suites/blocks/block_input_mapreduce.cfg
@@ -0,0 +1,2 @@
+in Input input_strings.txt
+providers yt
diff --git a/yt/yql/tests/sql/suites/blocks/block_input_mapreduce.sql b/yt/yql/tests/sql/suites/blocks/block_input_mapreduce.sql
new file mode 100644
index 00000000000..d938dce6fce
--- /dev/null
+++ b/yt/yql/tests/sql/suites/blocks/block_input_mapreduce.sql
@@ -0,0 +1,10 @@
+USE plato;
+
+PRAGMA yt.EnableFuseMapToMapReduce;
+PRAGMA yt.JobBlockInput;
+
+SELECT
+    key
+FROM Input
+WHERE key < "100"
+GROUP BY key;