Block mapjoin refactor

- Transition to block lists for block map join's right input (BlockMapJoinCore/BlockMapJoinIndex/BlockStorage nodes affected)

  * Right key columns/key drops are now addressed by name

- Optimizers which fuses ListToBlocks over ListFromBlocks and vice versa
commit_hash:bdcee24edd1e5298c038716d4d205858a199d0db

Author: loochek

yql/essentials/core/peephole_opt/yql_opt_peephole_physical.cpp

@@ -191,6 +191,25 @@ TExprNode::TPtr OptimizeWideToBlocks(const TExprNode::TPtr& node, TExprContext&
     return node;
 }
 
+TExprNode::TPtr OptimizeListToBlocks(const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& types) {
+    Y_UNUSED(ctx);
+    Y_UNUSED(types);
+    const auto& input = node->Head();
+    if (input.IsCallable("ListFromBlocks")) {
+        // TODO: actually (ListToBlocks (ListFromBlocks a)) is (ReplicateScalars a),
+        // but there is no sources of scalars in block lists at the moment, so just ensure it
+        auto inputItemType = input.Head().GetTypeAnn()->Cast<TListExprType>()->GetItemType()->Cast<TStructExprType>();
+        for (auto item : inputItemType->GetItems()) {
+            YQL_ENSURE(item->GetItemType()->IsBlock() || (item->GetItemType()->IsScalar() && item->GetName() == BlockLengthColumnName));
+        }
+
+        YQL_CLOG(DEBUG, CorePeepHole) << "Drop " << node->Content() << " over " << input.Content();
+        return input.HeadPtr();
+    }
+
+    return node;
+}
+
 TExprNode::TPtr OptimizeWideFromBlocks(const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& types) {
     Y_UNUSED(types);
     const auto& input = node->Head();
@@ -207,6 +226,18 @@ TExprNode::TPtr OptimizeWideFromBlocks(const TExprNode::TPtr& node, TExprContext
     return node;
 }
 
+TExprNode::TPtr OptimizeListFromBlocks(const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& types) {
+    Y_UNUSED(ctx);
+    Y_UNUSED(types);
+    const auto& input = node->Head();
+    if (input.IsCallable("ListToBlocks")) {
+        YQL_CLOG(DEBUG, CorePeepHole) << "Drop " << node->Content() << " over " << input.Content();
+        return input.HeadPtr();
+    }
+
+    return node;
+}
+
 TExprNode::TPtr OptimizeWideTakeSkipBlocks(const TExprNode::TPtr& node, TExprContext& ctx, TTypeAnnotationContext& types) {
     Y_UNUSED(types);
     const auto& input = node->HeadPtr();
@@ -8995,7 +9026,9 @@ struct TPeepHoleRules {
         {"NarrowMap", &OptimizeWideMapBlocks},
         {"WideFilter", &OptimizeWideFilterBlocks},
         {"WideToBlocks", &OptimizeWideToBlocks},
+        {"ListToBlocks", &OptimizeListToBlocks},
         {"WideFromBlocks", &OptimizeWideFromBlocks},
+        {"ListFromBlocks", &OptimizeListFromBlocks},
         {"WideTakeBlocks", &OptimizeWideTakeSkipBlocks},
         {"WideSkipBlocks", &OptimizeWideTakeSkipBlocks},
         {"BlockCompress", &OptimizeBlockCompress},

yql/essentials/core/type_ann/type_ann_join.cpp

@@ -315,57 +315,66 @@ class TBlockStorage : public TComputationValue<TBlockStorage> {
 
     TBlockStorage(
         TMemoryUsageInfo* memInfo,
-        const TVector<TType*>& itemTypes,
-        NUdf::TUnboxedValue stream,
+        const TVector<TType*>& types,
+        size_t blockLengthIndex,
+        NUdf::TUnboxedValue listIter,
         TStringBuf resourceTag,
         arrow::MemoryPool* pool
     )
         : TBase(memInfo)
-        , InputsDescr_(ToValueDescr(itemTypes))
-        , Stream_(std::move(stream))
-        , Inputs_(itemTypes.size())
+        , InputsDescr_(ToValueDescr(types))
+        , Readers_(types.size())
+        , Hashers_(types.size())
+        , Comparators_(types.size())
+        , Trimmers_(types.size())
+        , ListIter_(std::move(listIter))
+        , BlockLengthIndex_(blockLengthIndex)
         , ResourceTag_(std::move(resourceTag))
     {
         TBlockTypeHelper helper;
-        for (size_t i = 0; i < itemTypes.size(); i++) {
-            TType* blockItemType = AS_TYPE(TBlockType, itemTypes[i])->GetItemType();
-            Readers_.push_back(MakeBlockReader(TTypeInfoHelper(), blockItemType));
-            Hashers_.push_back(helper.MakeHasher(blockItemType));
-            Comparators_.push_back(helper.MakeComparator(blockItemType));
-            Trimmers_.push_back(MakeBlockTrimmer(TTypeInfoHelper(), blockItemType, pool));
+        for (size_t i = 0; i < types.size(); i++) {
+            if (i == BlockLengthIndex_) {
+                continue;
+            }
+
+            TType* blockItemType = AS_TYPE(TBlockType, types[i])->GetItemType();
+            Readers_[i] = MakeBlockReader(TTypeInfoHelper(), blockItemType);
+            Hashers_[i] = helper.MakeHasher(blockItemType);
+            Comparators_[i] = helper.MakeComparator(blockItemType);
+            Trimmers_[i] = MakeBlockTrimmer(TTypeInfoHelper(), blockItemType, pool);
         }
     }
 
-    NUdf::EFetchStatus FetchStream() {
-        switch (Stream_.WideFetch(Inputs_.data(), Inputs_.size())) {
-        case NUdf::EFetchStatus::Yield:
-            return NUdf::EFetchStatus::Yield;
-        case NUdf::EFetchStatus::Finish:
+    bool FetchNextBlock() {
+        if (!ListIter_.Next(Block_)) {
             IsFinished_ = true;
-            return NUdf::EFetchStatus::Finish;
-        case NUdf::EFetchStatus::Ok:
-            break;
+            return false;
         }
+        BlockItems_ = Block_.GetElements();
 
         Y_ENSURE(!IsFinished_, "Got data on finished stream");
 
-        std::vector<arrow::Datum> blockColumns;
-        for (size_t i = 0; i < Inputs_.size() - 1; i++) {
-            auto& datum = TArrowBlock::From(Inputs_[i]).GetDatum();
+        std::vector<arrow::Datum> blockColumns(Readers_.size());
+        for (size_t i = 0; i < Readers_.size(); i++) {
+            if (i == BlockLengthIndex_) {
+                continue;
+            }
+
+            auto& datum = TArrowBlock::From(BlockItems_[i]).GetDatum();
             ARROW_DEBUG_CHECK_DATUM_TYPES(InputsDescr_[i], datum.descr());
             if (datum.is_scalar()) {
-                blockColumns.push_back(datum);
+                blockColumns[i] = datum;
             } else {
                 MKQL_ENSURE(datum.is_array(), "Expecting array");
-                blockColumns.push_back(Trimmers_[i]->Trim(datum.array()));
+                blockColumns[i] = Trimmers_[i]->Trim(datum.array());
             }
         }
 
-        auto blockSize = ::GetBlockCount(Inputs_[Inputs_.size() - 1]);
+        auto blockSize = ::GetBlockCount(BlockItems_[BlockLengthIndex_]);
         Data_.emplace_back(blockSize, std::move(blockColumns));
         RowCount_ += blockSize;
 
-        return NUdf::EFetchStatus::Ok;
+        return true;
     }
 
     const TBlock& GetBlock(size_t blockOffset) const {
@@ -392,11 +401,11 @@ class TBlockStorage : public TComputationValue<TBlockStorage> {
     }
 
     TBlockItem GetItem(TRowEntry entry, ui32 columnIdx) const {
-        Y_ENSURE(columnIdx < Inputs_.size() - 1);
         return GetItemFromBlock(GetBlock(entry.BlockOffset), columnIdx, entry.ItemOffset);
     }
 
     TBlockItem GetItemFromBlock(const TBlock& block, ui32 columnIdx, size_t offset) const {
+        Y_ENSURE(columnIdx < Readers_.size() && columnIdx != BlockLengthIndex_);
         Y_ENSURE(offset < block.Size);
         const auto& datum = block.Columns[columnIdx];
         if (datum.is_scalar()) {
@@ -447,8 +456,11 @@ class TBlockStorage : public TComputationValue<TBlockStorage> {
     size_t RowCount_ = 0;
     bool IsFinished_ = false;
 
-    NUdf::TUnboxedValue Stream_;
-    TUnboxedValueVector Inputs_;
+    NUdf::TUnboxedValue ListIter_;
+    NUdf::TUnboxedValue Block_;
+    const NUdf::TUnboxedValue* BlockItems_ = nullptr;
+
+    size_t BlockLengthIndex_ = 0;
 
     const TStringBuf ResourceTag_;
 };
@@ -459,33 +471,44 @@ class TBlockStorageWrapper : public TMutableComputationNode<TBlockStorageWrapper
 public:
     TBlockStorageWrapper(
         TComputationMutables& mutables,
-        TVector<TType*>&& itemTypes,
-        IComputationNode* stream,
+        TStructType* structType,
+        IComputationNode* list,
         const TStringBuf& resourceTag
     )
         : TBaseComputation(mutables, EValueRepresentation::Boxed)
-        , ItemTypes_(std::move(itemTypes))
-        , Stream_(stream)
+        , List_(list)
         , ResourceTag_(resourceTag)
-    {}
+    {
+        for (size_t i = 0; i < structType->GetMembersCount(); i++) {
+            if (structType->GetMemberName(i) == NYql::BlockLengthColumnName) {
+                BlockLengthIndex_ = i;
+                Types_.push_back(nullptr);
+                continue;
+            }
+            Types_.push_back(structType->GetMemberType(i));
+        }
+    }
 
     NUdf::TUnboxedValuePod DoCalculate(TComputationContext& ctx) const {
         return ctx.HolderFactory.Create<TBlockStorage>(
-            ItemTypes_,
-            std::move(Stream_->GetValue(ctx)),
+            Types_,
+            BlockLengthIndex_,
+            List_->GetValue(ctx).GetListIterator(),
             ResourceTag_,
             &ctx.ArrowMemoryPool
         );
     }
 
 private:
     void RegisterDependencies() const final {
-        DependsOn(Stream_);
+        DependsOn(List_);
     }
 
 private:
-    const TVector<TType*> ItemTypes_;
-    IComputationNode* const Stream_;
+    TVector<TType*> Types_;
+    size_t BlockLengthIndex_ = 0;
+
+    IComputationNode* const List_;
 
     const TString ResourceTag_;
 };
@@ -992,17 +1015,13 @@ class TBlockMapJoinCoreWraper : public TMutableComputationNode<TBlockMapJoinCore
             auto& indexState = *static_cast<TIndexState*>(RightBlockIndex_.GetResource());
             auto& storageState = *static_cast<TStorageState*>(indexState.GetBlockStorage().GetResource());
 
-            if (!RightStreamConsumed_) {
-                auto fetchStatus = NUdf::EFetchStatus::Ok;
-                while (fetchStatus != NUdf::EFetchStatus::Finish) {
-                    fetchStatus = storageState.FetchStream();
-                    if (fetchStatus == NUdf::EFetchStatus::Yield) {
-                        return NUdf::EFetchStatus::Yield;
-                    }
+            if (!RightInputConsumed_) {
+                while (storageState.FetchNextBlock()) {
+                    // Fetch entire data from the right input
                 }
 
                 indexState.BuildIndex();
-                RightStreamConsumed_ = true;
+                RightInputConsumed_ = true;
             }
 
             auto* inputFields = joinState.GetRawInputFields();
@@ -1104,7 +1123,7 @@ class TBlockMapJoinCoreWraper : public TMutableComputationNode<TBlockMapJoinCore
         const TVector<ui32>& LeftKeyColumns_;
 
         const TVector<ui32>& RightIOMap_;
-        bool RightStreamConsumed_ = false;
+        bool RightInputConsumed_ = false;
 
         std::array<typename TIndexState::TIterator, PrefetchBatchSize> LookupBatchIterators_;
         ui32 LookupBatchCurrent_ = 0;
@@ -1204,16 +1223,12 @@ class TBlockCrossJoinCoreWraper : public TMutableComputationNode<TBlockCrossJoin
             auto& joinState = *static_cast<TJoinState*>(JoinState_.AsBoxed().Get());
             auto& storageState = *static_cast<TStorageState*>(RightBlockStorage_.GetResource());
 
-            if (!RightStreamConsumed_) {
-                auto fetchStatus = NUdf::EFetchStatus::Ok;
-                while (fetchStatus != NUdf::EFetchStatus::Finish) {
-                    fetchStatus = storageState.FetchStream();
-                    if (fetchStatus == NUdf::EFetchStatus::Yield) {
-                        return NUdf::EFetchStatus::Yield;
-                    }
+            if (!RightInputConsumed_) {
+                while (storageState.FetchNextBlock()) {
+                    // Fetch entire data from the right input
                 }
 
-                RightStreamConsumed_ = true;
+                RightInputConsumed_ = true;
                 RightRowIterator_ = storageState.GetRowIterator();
             }
 
@@ -1270,7 +1285,7 @@ class TBlockCrossJoinCoreWraper : public TMutableComputationNode<TBlockCrossJoin
         NUdf::TUnboxedValue JoinState_;
 
         const TVector<ui32>& RightIOMap_;
-        bool RightStreamConsumed_ = false;
+        bool RightInputConsumed_ = false;
 
         TStorageState::TRowIterator RightRowIterator_;
 
@@ -1310,19 +1325,15 @@ IComputationNode* WrapBlockStorage(TCallable& callable, const TComputationNodeFa
     MKQL_ENSURE(resultResourceType->GetTag().StartsWith(BlockStorageResourcePrefix), "Expected block storage resource");
 
     const auto inputType = callable.GetInput(0).GetStaticType();
-    MKQL_ENSURE(inputType->IsStream(), "Expected WideStream as an input stream");
-    const auto inputStreamType = AS_TYPE(TStreamType, inputType);
-    MKQL_ENSURE(inputStreamType->GetItemType()->IsMulti(),
-                "Expected Multi as a left stream item type");
-    const auto inputStreamComponents = GetWideComponents(inputStreamType);
-    MKQL_ENSURE(inputStreamComponents.size() > 0, "Expected at least one column");
-    TVector<TType*> inputStreamItems(inputStreamComponents.cbegin(), inputStreamComponents.cend());
+    MKQL_ENSURE(inputType->IsList(), "Expected List as an input stream");
+    const auto inputItemType = AS_TYPE(TListType, inputType)->GetItemType();;
+    MKQL_ENSURE(inputItemType->IsStruct(), "Expected Struct as a list item type");
 
-    const auto inputStream = LocateNode(ctx.NodeLocator, callable, 0);
+    const auto list = LocateNode(ctx.NodeLocator, callable, 0);
     return new TBlockStorageWrapper(
         ctx.Mutables,
-        std::move(inputStreamItems),
-        inputStream,
+        AS_TYPE(TStructType, inputItemType),
+        list,
         resultResourceType->GetTag()
     );
 }
@@ -1341,9 +1352,9 @@ IComputationNode* WrapBlockMapJoinIndex(TCallable& callable, const TComputationN
     MKQL_ENSURE(inputResourceType->GetTag().StartsWith(BlockStorageResourcePrefix), "Expected block storage resource");
 
     auto origInputItemType = AS_VALUE(TTypeType, callable.GetInput(1));
-    MKQL_ENSURE(origInputItemType->IsMulti(), "Expected Multi as an input item type");
-    const auto streamComponents = AS_TYPE(TMultiType, origInputItemType)->GetElements();
-    MKQL_ENSURE(streamComponents.size() > 0, "Expected at least one column");
+    MKQL_ENSURE(origInputItemType->IsStruct(), "Expected Struct as an input item type");
+    const auto origInputItemStructType = AS_TYPE(TStructType, origInputItemType);
+    MKQL_ENSURE(origInputItemStructType->GetMembersCount() > 0, "Expected at least one column");
 
     const auto keyColumnsLiteral = callable.GetInput(2);
     const auto keyColumnsTuple = AS_VALUE(TTupleLiteral, keyColumnsLiteral);
@@ -1355,7 +1366,7 @@ IComputationNode* WrapBlockMapJoinIndex(TCallable& callable, const TComputationN
     }
 
     for (ui32 keyColumn : keyColumns) {
-        MKQL_ENSURE(keyColumn < streamComponents.size() - 1, "Key column out of range");
+        MKQL_ENSURE(keyColumn < origInputItemStructType->GetMembersCount(), "Key column out of range");
     }
 
     const auto anyNode = callable.GetInput(3);
@@ -1408,10 +1419,9 @@ IComputationNode* WrapBlockMapJoinCore(TCallable& callable, const TComputationNo
     }
 
     auto origRightItemType = AS_VALUE(TTypeType, callable.GetInput(2));
-    MKQL_ENSURE(origRightItemType->IsMulti(), "Expected Multi as a right stream item type");
-    const auto rightStreamComponents = AS_TYPE(TMultiType, origRightItemType)->GetElements();
-    MKQL_ENSURE(rightStreamComponents.size() > 0, "Expected at least one column");
-    const TVector<TType*> rightStreamItems(rightStreamComponents.cbegin(), rightStreamComponents.cend());
+    MKQL_ENSURE(origRightItemType->IsStruct(), "Expected Struct as a right stream item type");
+    const auto origRightItemStructType = AS_TYPE(TStructType, origRightItemType);
+    MKQL_ENSURE(origRightItemStructType->GetMembersCount() > 0, "Expected at least one column");
 
     const auto leftKeyColumnsLiteral = callable.GetInput(4);
     const auto leftKeyColumnsTuple = AS_VALUE(TTupleLiteral, leftKeyColumnsLiteral);
@@ -1479,8 +1489,8 @@ IComputationNode* WrapBlockMapJoinCore(TCallable& callable, const TComputationNo
     // XXX: Mind the last wide item, containing block length.
     TVector<ui32> rightIOMap;
     if (joinKind == EJoinKind::Inner || joinKind == EJoinKind::Left || joinKind == EJoinKind::Cross) {
-        for (size_t i = 0; i < rightStreamItems.size() - 1; i++) {
-            if (rightKeyDrops.contains(i)) {
+        for (size_t i = 0; i < origRightItemStructType->GetMembersCount(); i++) {
+            if (rightKeyDrops.contains(i) || origRightItemStructType->GetMemberName(i) == NYql::BlockLengthColumnName) {
                 continue;
             }
             rightIOMap.push_back(i);