Support block IO in purecalc (#7016)

Author: igormunkin

ydb/library/yql/public/purecalc/common/compile_mkql.cpp

@@ -98,6 +98,7 @@ NKikimr::NMiniKQL::TRuntimeNode CompileMkql(const TExprNode::TPtr& exprRoot, TEx
     NCommon::TMkqlCommonCallableCompiler compiler;
 
     compiler.AddCallable(PurecalcInputCallableName, MakeSelfCallableCompiler());
+    compiler.AddCallable(PurecalcBlockInputCallableName, MakeSelfCallableCompiler());
     compiler.OverrideCallable("FileContent", MakeFileContentCallableCompiler(userData));
     compiler.OverrideCallable("FilePath", MakeFilePathCallableCompiler(userData));
     compiler.OverrideCallable("FolderPath", MakeFolderPathCallableCompiler(userData));

ydb/library/yql/public/purecalc/common/interface.h

@@ -591,12 +591,28 @@ namespace NYql {
              */
             virtual const NKikimr::NMiniKQL::TStructType* GetInputType(bool original = false) const = 0;
 
+            /**
+             * MiniKQL input struct type of the specified input for this program.
+             * The returned type is the actual type of the specified input node.
+             */
+            virtual const NKikimr::NMiniKQL::TStructType* GetRawInputType(ui32) const = 0;
+            /**
+             * Overload for single-input programs.
+             */
+            virtual const NKikimr::NMiniKQL::TStructType* GetRawInputType() const = 0;
+
             /**
              * MiniKQL output struct type for this program. The returned type is equivalent to the deduced output
              * schema (see IWorker::MakeFullOutputSchema()).
              */
             virtual const NKikimr::NMiniKQL::TType* GetOutputType() const = 0;
 
+            /**
+             * MiniKQL output struct type for this program. The returned type is
+             * the actual type of the root node.
+             */
+            virtual const NKikimr::NMiniKQL::TType* GetRawOutputType() const = 0;
+
             /**
              * Make input type schema for specified input as deduced by program optimizer. This schema is equivalent
              * to one provided by input spec up to the order of the fields in structures.
@@ -795,7 +811,7 @@ namespace NYql {
                 return AllVirtualColumns_;
             }
 
-            static constexpr bool ProvidesBlocks = false;
+            virtual bool ProvidesBlocks() const { return false; }
         };
 
         /**
@@ -840,7 +856,7 @@ namespace NYql {
                 OutputColumnsFilter_ = outputColumnsFilter;
             }
 
-            static constexpr bool AcceptsBlocks = false;
+            virtual bool AcceptsBlocks() const { return false; }
         };
 
         ////////////////////////////////////////////////////////////////////////////////////////////////////

ydb/library/yql/public/purecalc/common/names.cpp

@@ -5,12 +5,15 @@
 namespace NYql::NPureCalc {
     const TStringBuf PurecalcSysColumnsPrefix = "_yql_sys_";
     const TStringBuf PurecalcSysColumnTablePath = "_yql_sys_tablepath";
+    const TStringBuf PurecalcBlockColumnLength = "_yql_block_length";
 
     const TStringBuf PurecalcDefaultCluster = "view";
     const TStringBuf PurecalcDefaultService = "data";
 
     const TStringBuf PurecalcInputCallableName = "Self";
     const TStringBuf PurecalcInputTablePrefix = "Input";
 
+    const TStringBuf PurecalcBlockInputCallableName = "BlockSelf";
+
     const TStringBuf PurecalcUdfModulePrefix = "<purecalc>::";
 }

ydb/library/yql/public/purecalc/common/names.h

@@ -5,12 +5,15 @@
 namespace NYql::NPureCalc {
     extern const TStringBuf PurecalcSysColumnsPrefix;
     extern const TStringBuf PurecalcSysColumnTablePath;
+    extern const TStringBuf PurecalcBlockColumnLength;
 
     extern const TStringBuf PurecalcDefaultCluster;
     extern const TStringBuf PurecalcDefaultService;
 
     extern const TStringBuf PurecalcInputCallableName;
     extern const TStringBuf PurecalcInputTablePrefix;
 
+    extern const TStringBuf PurecalcBlockInputCallableName;
+
     extern const TStringBuf PurecalcUdfModulePrefix;
 }

ydb/library/yql/public/purecalc/common/transformations/align_output_schema.cpp

@@ -1,6 +1,8 @@
 #include "align_output_schema.h"
 
+#include <ydb/library/yql/public/purecalc/common/names.h>
 #include <ydb/library/yql/public/purecalc/common/type_from_schema.h>
+#include <ydb/library/yql/public/purecalc/common/transformations/utils.h>
 
 #include <ydb/library/yql/core/yql_expr_type_annotation.h>
 
@@ -11,11 +13,17 @@ namespace {
     class TOutputAligner : public TSyncTransformerBase {
     private:
         const TTypeAnnotationNode* OutputStruct_;
+        bool AcceptsBlocks_;
         EProcessorMode ProcessorMode_;
 
     public:
-        explicit TOutputAligner(const TTypeAnnotationNode* outputStruct, EProcessorMode processorMode)
+        explicit TOutputAligner(
+            const TTypeAnnotationNode* outputStruct,
+            bool acceptsBlocks,
+            EProcessorMode processorMode
+        )
             : OutputStruct_(outputStruct)
+            , AcceptsBlocks_(acceptsBlocks)
             , ProcessorMode_(processorMode)
         {
         }
@@ -29,6 +37,19 @@ namespace {
             const auto* actualType = MakeActualType(input);
             const auto* actualItemType = MakeActualItemType(input);
 
+            // XXX: Tweak the obtained expression type, is the spec supports blocks:
+            // 1. Remove "_yql_block_length" attribute, since it's for internal usage.
+            // 2. Strip block container from the type to store its internal type.
+            if (AcceptsBlocks_) {
+                Y_ENSURE(actualItemType->GetKind() == ETypeAnnotationKind::Struct);
+                actualItemType = UnwrapBlockStruct(actualItemType->Cast<TStructExprType>(), ctx);
+                if (ProcessorMode_ == EProcessorMode::PullList) {
+                    actualType = ctx.MakeType<TListExprType>(actualItemType);
+                } else {
+                    actualType = ctx.MakeType<TStreamExprType>(actualItemType);
+                }
+            }
+
             if (!ValidateOutputType(actualItemType, expectedItemType, ctx)) {
                 return TStatus::Error;
             }
@@ -78,8 +99,12 @@ namespace {
             auto actualType = MakeActualType(input);
             switch (actualType->GetKind()) {
                 case ETypeAnnotationKind::Stream:
+                    Y_ENSURE(ProcessorMode_ != EProcessorMode::PullList,
+                             "processor mode mismatches the actual container type");
                     return actualType->Cast<TStreamExprType>()->GetItemType();
                 case ETypeAnnotationKind::List:
+                    Y_ENSURE(ProcessorMode_ == EProcessorMode::PullList,
+                             "processor mode mismatches the actual container type");
                     return actualType->Cast<TListExprType>()->GetItemType();
                 default:
                     Y_ABORT("unexpected return type");
@@ -88,6 +113,10 @@ namespace {
     };
 }
 
-TAutoPtr<IGraphTransformer> NYql::NPureCalc::MakeOutputAligner(const TTypeAnnotationNode* outputStruct, EProcessorMode processorMode) {
-    return new TOutputAligner(outputStruct, processorMode);
+TAutoPtr<IGraphTransformer> NYql::NPureCalc::MakeOutputAligner(
+    const TTypeAnnotationNode* outputStruct,
+    bool acceptsBlocks,
+    EProcessorMode processorMode
+) {
+    return new TOutputAligner(outputStruct, acceptsBlocks, processorMode);
 }

ydb/library/yql/public/purecalc/common/transformations/align_output_schema.h

@@ -11,10 +11,14 @@ namespace NYql {
          * A transformer which converts an output type of the expression to the given type or reports an error.
          *
          * @param outputStruct destination output struct type.
+         * @param acceptsBlocks indicates, whether the output type need to be
+         *        preprocessed.
+         * @param processorMode specifies the top-most container of the result.
          * @return a graph transformer for type alignment.
          */
         TAutoPtr<IGraphTransformer> MakeOutputAligner(
             const TTypeAnnotationNode* outputStruct,
+            bool acceptsBlocks,
             EProcessorMode processorMode
         );
     }

ydb/library/yql/public/purecalc/common/transformations/replace_table_reads.cpp

@@ -1,6 +1,7 @@
 #include "replace_table_reads.h"
 
 #include <ydb/library/yql/public/purecalc/common/names.h>
+#include <ydb/library/yql/public/purecalc/common/transformations/utils.h>
 
 #include <ydb/library/yql/core/yql_expr_optimize.h>
 #include <ydb/library/yql/core/yql_expr_type_annotation.h>
@@ -11,23 +12,26 @@ using namespace NYql::NPureCalc;
 namespace {
     class TTableReadsReplacer: public TSyncTransformerBase {
     private:
-        ui32 InputsNumber_;
+        const TVector<const TStructExprType*>& InputStructs_;
         bool UseSystemColumns_;
-        TString TablePrefix_;
+        EProcessorMode ProcessorMode_;
         TString CallableName_;
+        TString TablePrefix_;
         bool Complete_ = false;
 
     public:
         explicit TTableReadsReplacer(
-            ui32 inputsNumber,
+            const TVector<const TStructExprType*>& inputStructs,
             bool useSystemColumns,
-            TString tablePrefix,
-            TString inputNodeName
+            EProcessorMode processorMode,
+            TString inputNodeName,
+            TString tablePrefix
         )
-            : InputsNumber_(inputsNumber)
+            : InputStructs_(inputStructs)
             , UseSystemColumns_(useSystemColumns)
-            , TablePrefix_(std::move(tablePrefix))
+            , ProcessorMode_(processorMode)
             , CallableName_(std::move(inputNodeName))
+            , TablePrefix_(std::move(tablePrefix))
         {
         }
 
@@ -135,6 +139,13 @@ namespace {
                 .Seal()
                 .Build();
 
+            if (inputNode->IsCallable(PurecalcBlockInputCallableName)) {
+                const auto inputStruct = InputStructs_[inputIndex]->Cast<TStructExprType>();
+                const auto blocksLambda = NodeFromBlocks(replacePos, inputStruct, ctx);
+                bool wrapLMap = ProcessorMode_ == EProcessorMode::PullList;
+                inputNode = ApplyToIterable(replacePos, inputNode, blocksLambda, wrapLMap, ctx);
+            }
+
             if (UseSystemColumns_) {
                 auto mapLambda = ctx.Builder(replacePos)
                     .Lambda()
@@ -163,7 +174,7 @@ namespace {
                 return MakeIntrusive<TIssue>(ctx.GetPosition(node->Pos()), TStringBuilder() << "At function: " << node->Content());
             });
 
-            if (!InputsNumber_) {
+            if (InputStructs_.empty()) {
                 ctx.AddError(TIssue(ctx.GetPosition(node->Pos()), "No inputs provided by input spec"));
                 return nullptr;
             }
@@ -174,10 +185,10 @@ namespace {
 
             auto builder = ctx.Builder(replacePos);
 
-            if (InputsNumber_ > 1) {
+            if (InputStructs_.size() > 1) {
                 auto listBuilder = builder.List();
 
-                for (ui32 i = 0; i < InputsNumber_; ++i) {
+                for (ui32 i = 0; i < InputStructs_.size(); ++i) {
                     listBuilder.Callable(i, CallableName_).Atom(0, ToString(i)).Seal();
                 }
 
@@ -226,10 +237,11 @@ namespace {
 }
 
 TAutoPtr<IGraphTransformer> NYql::NPureCalc::MakeTableReadsReplacer(
-    ui32 inputsNumber,
+    const TVector<const TStructExprType*>& inputStructs,
     bool useSystemColumns,
-    TString tablePrefix,
-    TString callableName
+    EProcessorMode processorMode,
+    TString callableName,
+    TString tablePrefix
 ) {
-    return new TTableReadsReplacer(inputsNumber, useSystemColumns, std::move(tablePrefix), std::move(callableName));
+    return new TTableReadsReplacer(inputStructs, useSystemColumns, processorMode, std::move(callableName), std::move(tablePrefix));
 }

ydb/library/yql/public/purecalc/common/transformations/replace_table_reads.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <ydb/library/yql/public/purecalc/common/names.h>
+#include <ydb/library/yql/public/purecalc/common/processor_mode.h>
 
 #include <ydb/library/yql/core/yql_graph_transformer.h>
 
@@ -15,14 +16,15 @@ namespace NYql::NPureCalc {
      *
      * @param inputStructs types of each input.
      * @param useSystemColumns whether to allow special system columns in input structs.
-     * @param tablePrefix required prefix for all table names (e.g. `Input`).
      * @param callableName name of the special callable used to get input data (e.g. `Self`).
+     * @param tablePrefix required prefix for all table names (e.g. `Input`).
      * @param return a graph transformer for replacing table reads.
      */
     TAutoPtr<IGraphTransformer> MakeTableReadsReplacer(
-        ui32 inputsNumber,
+        const TVector<const TStructExprType*>& inputStructs,
         bool useSystemColumns,
-        TString tablePrefix = TString{PurecalcInputTablePrefix},
-        TString callableName = TString{PurecalcInputCallableName}
+        EProcessorMode processorMode,
+        TString callableName = TString{PurecalcInputCallableName},
+        TString tablePrefix = TString{PurecalcInputTablePrefix}
     );
 }

ydb/library/yql/public/purecalc/common/transformations/root_to_blocks.cpp

@@ -0,0 +1,65 @@
+#include "root_to_blocks.h"
+
+#include <ydb/library/yql/public/purecalc/common/transformations/utils.h>
+
+#include <ydb/library/yql/core/yql_expr_type_annotation.h>
+
+using namespace NYql;
+using namespace NYql::NPureCalc;
+
+namespace {
+
+class TRootToBlocks: public TSyncTransformerBase {
+private:
+    bool AcceptsBlocks_;
+    EProcessorMode ProcessorMode_;
+    bool Wrapped_;
+
+public:
+    explicit TRootToBlocks(bool acceptsBlocks, EProcessorMode processorMode)
+        : AcceptsBlocks_(acceptsBlocks)
+        , ProcessorMode_(processorMode)
+        , Wrapped_(false)
+    {
+    }
+
+public:
+    void Rewind() override {
+        Wrapped_ = false;
+    }
+
+    TStatus DoTransform(TExprNode::TPtr input, TExprNode::TPtr& output, TExprContext& ctx) final {
+        if (Wrapped_ || !AcceptsBlocks_) {
+            return IGraphTransformer::TStatus::Ok;
+        }
+
+        const TTypeAnnotationNode* returnItemType;
+        const TTypeAnnotationNode* returnType = input->GetTypeAnn();
+        if (ProcessorMode_ == EProcessorMode::PullList) {
+            Y_ENSURE(returnType->GetKind() == ETypeAnnotationKind::List);
+            returnItemType = returnType->Cast<TListExprType>()->GetItemType();
+        } else {
+            Y_ENSURE(returnType->GetKind() == ETypeAnnotationKind::Stream);
+            returnItemType = returnType->Cast<TStreamExprType>()->GetItemType();
+        }
+
+        Y_ENSURE(returnItemType->GetKind() == ETypeAnnotationKind::Struct);
+        const TStructExprType* structType = returnItemType->Cast<TStructExprType>();
+        const auto blocksLambda = NodeToBlocks(input->Pos(), structType, ctx);
+        bool wrapLMap = ProcessorMode_ == EProcessorMode::PullList;
+        output = ApplyToIterable(input->Pos(), input, blocksLambda, wrapLMap, ctx);
+
+        Wrapped_ = true;
+
+        return IGraphTransformer::TStatus(IGraphTransformer::TStatus::Repeat, true);
+    }
+};
+
+} // namespace
+
+TAutoPtr<IGraphTransformer> NYql::NPureCalc::MakeRootToBlocks(
+    bool acceptsBlocks,
+    EProcessorMode processorMode
+) {
+    return new TRootToBlocks(acceptsBlocks, processorMode);
+}

ydb/library/yql/public/purecalc/common/transformations/root_to_blocks.h

@@ -0,0 +1,22 @@
+#pragma once
+
+#include <ydb/library/yql/public/purecalc/common/processor_mode.h>
+
+#include <ydb/library/yql/core/yql_graph_transformer.h>
+
+namespace NYql {
+    namespace NPureCalc {
+        /**
+         * A transformer which rewrite the root to respect block types.
+         *
+         * @param acceptsBlock allows using this transformer in pipeline and
+         *        skip this phase if no block output is required.
+         * @param processorMode specifies the top-most container of the result.
+         * @return a graph transformer for rewriting the root node.
+         */
+        TAutoPtr<IGraphTransformer> MakeRootToBlocks(
+            bool acceptsBlocks,
+            EProcessorMode processorMode
+        );
+    }
+}