YT block input for table content

commit_hash:6ece06798fc8cef2e4d8e62cf5b9634b3162aa45

Author: loochek

yt/yql/providers/yt/codec/yt_codec.h

@@ -133,6 +133,10 @@ class TMkqlIOSpecs {
         UseBlockOutput_ = true;
     }
 
+    void SetIsTableContent() {
+        IsTableContent_ = true;
+    }
+
     void SetTableOffsets(const TVector<ui64>& offsets);
 
     void Clear();
@@ -148,6 +152,7 @@ class TMkqlIOSpecs {
     bool UseSkiff_ = false;
     bool UseBlockInput_ = false;
     bool UseBlockOutput_ = false;
+    bool IsTableContent_ = false;
     TString OptLLVM_;
     TSystemFields SystemFields_;
 

yt/yql/providers/yt/codec/yt_codec_io.cpp

@@ -1536,7 +1536,7 @@ class TArrowDecoder : public TMkqlReaderImpl::TDecoder {
         YQL_ENSURE(inputFields.size() == ColumnConverters_.size());
 
         auto rowIndices = batch->GetColumnByName("$row_index");
-        YQL_ENSURE(rowIndices || decoder.Dynamic);
+        YQL_ENSURE(rowIndices || decoder.Dynamic || Specs_.IsTableContent_);
 
         arrow::compute::ExecContext execContext(Pool_);
         std::vector<arrow::Datum> convertedBatch;

yt/yql/providers/yt/common/yql_yt_settings.cpp

@@ -524,6 +524,7 @@ TYtConfiguration::TYtConfiguration(TTypeAnnotationContext& typeCtx)
     REGISTER_SETTING(*this, MaxColumnGroups);
     REGISTER_SETTING(*this, ExtendedStatsMaxChunkCount);
     REGISTER_SETTING(*this, JobBlockInput);
+    REGISTER_SETTING(*this, JobBlockTableContent);
     REGISTER_SETTING(*this, JobBlockOutput).Parser([](const TString& v) { return FromString<EBlockOutputMode>(v); });
     REGISTER_SETTING(*this, _EnableYtDqProcessWriteConstraints);
     REGISTER_SETTING(*this, CompactForDistinct);

yt/yql/providers/yt/common/yql_yt_settings.h

@@ -294,6 +294,7 @@ struct TYtSettings {
     NCommon::TConfSetting<ui16, false> MaxColumnGroups;
     NCommon::TConfSetting<ui64, false> ExtendedStatsMaxChunkCount;
     NCommon::TConfSetting<bool, false> JobBlockInput;
+    NCommon::TConfSetting<bool, false> JobBlockTableContent;
     NCommon::TConfSetting<TSet<TString>, false> JobBlockInputSupportedTypes;
     NCommon::TConfSetting<TSet<NUdf::EDataSlot>, false> JobBlockInputSupportedDataTypes;
     NCommon::TConfSetting<EBlockOutputMode, false> JobBlockOutput;

yt/yql/providers/yt/comp_nodes/ya.make.inc

@@ -4,11 +4,13 @@ INCLUDE(${ARCADIA_ROOT}/yql/essentials/minikql/invoke_builtins/header.ya.make.in
 SET(ORIG_SRC_DIR ${ARCADIA_ROOT}/yt/yql/providers/yt/comp_nodes)
 
 SET(ORIG_SOURCES
+    yql_mkql_file_block_stream.cpp
     yql_mkql_file_input_state.cpp
     yql_mkql_file_list.cpp
     yql_mkql_input_stream.cpp
     yql_mkql_input.cpp
     yql_mkql_output.cpp
+    yql_mkql_block_table_content.cpp
     yql_mkql_table_content.cpp
     yql_mkql_table.cpp
     yql_mkql_ungrouping_list.cpp

yt/yql/providers/yt/comp_nodes/yql_mkql_block_table_content.cpp

@@ -0,0 +1,73 @@
+#include "yql_mkql_block_table_content.h"
+#include "yql_mkql_file_block_stream.h"
+
+#include <yql/essentials/minikql/computation/mkql_computation_node_impl.h>
+#include <yql/essentials/minikql/mkql_node_cast.h>
+#include <yql/essentials/minikql/defs.h>
+
+#include <yql/essentials/public/udf/udf_value.h>
+
+#include <util/generic/vector.h>
+#include <util/generic/string.h>
+#include <util/generic/size_literals.h>
+
+namespace NYql {
+
+using namespace NKikimr;
+using namespace NKikimr::NMiniKQL;
+
+class TYtBlockTableContentWrapper : public TMutableComputationNode<TYtBlockTableContentWrapper> {
+    typedef TMutableComputationNode<TYtBlockTableContentWrapper> TBaseComputation;
+public:
+    TYtBlockTableContentWrapper(TComputationMutables& mutables, NCommon::TCodecContext& codecCtx,
+        TVector<TString>&& files, const TString& inputSpec, TStructType* origStructType, bool decompress, std::optional<ui64> expectedRowCount)
+        : TBaseComputation(mutables)
+        , Files_(std::move(files))
+        , Decompress_(decompress)
+        , ExpectedRowCount_(std::move(expectedRowCount))
+    {
+        Spec_.SetUseBlockInput();
+        Spec_.SetIsTableContent();
+        Spec_.Init(codecCtx, inputSpec, {}, {}, origStructType, {}, TString());
+    }
+
+    NUdf::TUnboxedValuePod DoCalculate(TComputationContext& ctx) const {
+        return ctx.HolderFactory.Create<TFileWideBlockStreamValue>(Spec_, ctx.HolderFactory, Files_, Decompress_, 4, 1_MB, ExpectedRowCount_);
+    }
+
+private:
+    void RegisterDependencies() const final {}
+
+    TMkqlIOSpecs Spec_;
+    TVector<TString> Files_;
+    const bool Decompress_;
+    const std::optional<ui64> ExpectedRowCount_;
+};
+
+IComputationNode* WrapYtBlockTableContent(NCommon::TCodecContext& codecCtx,
+    TComputationMutables& mutables, TCallable& callable, TStringBuf pathPrefix)
+{
+    MKQL_ENSURE(callable.GetInputsCount() == 6, "Expected 6 arguments");
+    TString uniqueId(AS_VALUE(TDataLiteral, callable.GetInput(0))->AsValue().AsStringRef());
+    auto origStructType = AS_TYPE(TStructType, AS_VALUE(TTypeType, callable.GetInput(1)));
+    const ui32 tablesCount = AS_VALUE(TDataLiteral, callable.GetInput(2))->AsValue().Get<ui32>();
+    TString inputSpec(AS_VALUE(TDataLiteral, callable.GetInput(3))->AsValue().AsStringRef());
+    const bool decompress = AS_VALUE(TDataLiteral, callable.GetInput(4))->AsValue().Get<bool>();
+
+    std::optional<ui64> length;
+    TTupleLiteral* lengthTuple = AS_VALUE(TTupleLiteral, callable.GetInput(5));
+    if (lengthTuple->GetValuesCount() > 0) {
+        MKQL_ENSURE(lengthTuple->GetValuesCount() == 1, "Expect 1 element in the length tuple");
+        length = AS_VALUE(TDataLiteral, lengthTuple->GetValue(0))->AsValue().Get<ui64>();
+    }
+
+    TVector<TString> files;
+    for (ui32 index = 0; index < tablesCount; ++index) {
+        files.push_back(TStringBuilder() << pathPrefix << uniqueId << '_' << index);
+    }
+
+    return new TYtBlockTableContentWrapper(mutables, codecCtx, std::move(files), inputSpec,
+        origStructType, decompress, length);
+}
+
+} // NYql

yt/yql/providers/yt/comp_nodes/yql_mkql_block_table_content.h

@@ -0,0 +1,18 @@
+#pragma once
+
+#include <yql/essentials/providers/common/codec/yql_codec.h>
+
+#include <yql/essentials/minikql/mkql_node.h>
+#include <yql/essentials/minikql/computation/mkql_computation_node.h>
+
+#include <util/generic/string.h>
+#include <util/generic/strbuf.h>
+
+namespace NYql {
+
+NKikimr::NMiniKQL::IComputationNode* WrapYtBlockTableContent(
+    NYql::NCommon::TCodecContext& codecCtx,
+    NKikimr::NMiniKQL::TComputationMutables& mutables,
+    NKikimr::NMiniKQL::TCallable& callable, TStringBuf pathPrefix);
+
+} // NYql

yt/yql/providers/yt/comp_nodes/yql_mkql_file_block_stream.cpp

@@ -0,0 +1,50 @@
+#include "yql_mkql_file_block_stream.h"
+#include "yql_mkql_file_input_state.h"
+
+namespace NYql {
+
+using namespace NKikimr::NMiniKQL;
+
+TFileWideBlockStreamValue::TFileWideBlockStreamValue(
+    TMemoryUsageInfo* memInfo,
+    const TMkqlIOSpecs& spec,
+    const THolderFactory& holderFactory,
+    const TVector<TString>& filePaths,
+    bool decompress,
+    size_t blockCount,
+    size_t blockSize,
+    std::optional<ui64> expectedRowCount
+)
+    : TComputationValue(memInfo)
+    , Spec_(spec)
+    , HolderFactory_(holderFactory)
+    , FilePaths_(filePaths)
+    , Decompress_(decompress)
+    , BlockCount_(blockCount)
+    , BlockSize_(blockSize)
+    , ExpectedRowCount_(expectedRowCount)
+{
+    State_ = MakeHolder<TFileInputState>(Spec_, HolderFactory_, MakeMkqlFileInputs(FilePaths_, Decompress_), BlockCount_, BlockSize_);
+}
+
+NUdf::EFetchStatus TFileWideBlockStreamValue::WideFetch(NUdf::TUnboxedValue* output, ui32 width) {
+    if (!AtStart_) {
+        State_->Next();
+    }
+    AtStart_ = false;
+    if (!State_->IsValid()) {
+        MKQL_ENSURE(!ExpectedRowCount_ || *ExpectedRowCount_ == State_->GetRecordIndex(), "Invalid file row count");
+        return NUdf::EFetchStatus::Finish;
+    }
+
+    auto elements = State_->GetCurrent().GetElements();
+    for (ui32 i = 0; i < width; i++) {
+        if (auto out = output++) {
+            *out = elements[i];
+        }
+    }
+
+    return NUdf::EFetchStatus::Ok;
+}
+
+}

yt/yql/providers/yt/comp_nodes/yql_mkql_file_block_stream.h

@@ -0,0 +1,43 @@
+#pragma once
+
+#include "yql_mkql_file_input_state.h"
+
+#include <yt/yql/providers/yt/codec/yt_codec.h>
+#include <yql/essentials/minikql/computation/mkql_computation_node.h>
+
+#include <util/generic/ptr.h>
+#include <util/generic/vector.h>
+#include <util/generic/string.h>
+
+namespace NYql {
+
+class TFileWideBlockStreamValue : public NKikimr::NMiniKQL::TComputationValue<TFileWideBlockStreamValue> {
+public:
+    TFileWideBlockStreamValue(
+        NKikimr::NMiniKQL::TMemoryUsageInfo* memInfo,
+        const TMkqlIOSpecs& spec,
+        const NKikimr::NMiniKQL::THolderFactory& holderFactory,
+        const TVector<TString>& filePaths,
+        bool decompress,
+        size_t blockCount,
+        size_t blockSize,
+        std::optional<ui64> expectedRowCount
+    );
+
+private:
+    NUdf::EFetchStatus WideFetch(NUdf::TUnboxedValue* output, ui32 width);
+
+private:
+    const TMkqlIOSpecs& Spec_;
+    const NKikimr::NMiniKQL::THolderFactory& HolderFactory_;
+    const TVector<TString> FilePaths_;
+    const bool Decompress_;
+    const size_t BlockCount_;
+    const size_t BlockSize_;
+    const std::optional<ui64> ExpectedRowCount_;
+
+    bool AtStart_ = true;
+    THolder<TFileInputState> State_;
+};
+
+}

yt/yql/providers/yt/comp_nodes/yql_mkql_file_input_state.cpp

@@ -1,5 +1,6 @@
 #include "yql_mkql_file_input_state.h"
 
+#include <yql/essentials/minikql/computation/mkql_block_impl.h>
 #include <yql/essentials/utils/yql_panic.h>
 
 #include <util/system/fs.h>
@@ -55,7 +56,13 @@ bool TFileInputState::NextValue() {
         }
 
         MkqlReader_.Next();
-        ++CurrentRecord_;
+        if (Spec_->UseBlockInput_) {
+            auto blockCountValue = CurrentValue_.GetElement(Spec_->Inputs[CurrentInput_]->StructSize);
+            CurrentRecord_ += GetBlockCount(blockCountValue);
+        } else {
+            ++CurrentRecord_;
+        }
+
         return true;
     }
 }