block join with neumann ht and packed tuple format (#27450)

Co-authored-by: Mikhail Filitov <filitovme@gmail.com>

Author: nfrmtk

ydb/core/kqp/query_compiler/kqp_mkql_compiler.cpp

@@ -475,12 +475,12 @@ TIntrusivePtr<IMkqlCallableCompiler> CreateKqlCompiler(const TKqlCompileContext&
                 auto wideStreamComponentsSize = [](TRuntimeNode node)->int {
                     return AS_TYPE(TMultiType, AS_TYPE(TStreamType,node.GetStaticType())->GetItemType())->GetElementsCount();
                 };
-                TDqRenames renames{};
+                TDqUserRenames renames{};
                 for(int index = 0; index < wideStreamComponentsSize(leftInput) - 1; ++index) {
-                    renames.emplace_back(index, JoinSide::kLeft);       
+                    renames.emplace_back(index, EJoinSide::kLeft);       
                 }
                 for(int index = 0; index < wideStreamComponentsSize(rightInput) - 1; ++index) {
-                    renames.emplace_back(index, JoinSide::kRight);       
+                    renames.emplace_back(index, EJoinSide::kRight);       
                 }
                 return TGraceJoinRenames::FromDq(renames);    
             }();

ydb/core/kqp/tools/join_perf/construct_join_graph.cpp

@@ -88,9 +88,9 @@ THolder<IComputationGraph> ConstructJoinGraphStream(EJoinKind joinKind, ETestedJ
 
     const TVector<TType*> resultTypesArr = [&] {
         TVector<TType*> arr;
-        TDqRenames dqRenames = FromGraceFormat(renames);
+        TDqUserRenames dqRenames = FromGraceFormat(renames);
         for (auto rename : dqRenames) {
-            if (rename.Side == JoinSide::kLeft) {
+            if (rename.Side == EJoinSide::kLeft) {
                 auto* resType = descr.LeftSource.ColumnTypes[rename.Index];
                 arr.push_back([&] {
                     if (ForceLeftOptional(joinKind) && !resType->IsOptional()) {

ydb/core/kqp/tools/join_perf/construct_join_graph.h

@@ -17,7 +17,7 @@ struct TJoinDescription {
     TJoinSourceData LeftSource;
     TJoinSourceData RightSource;
     TDqSetup<false>* Setup;
-    std::optional<TDqRenames> CustomRenames;
+    std::optional<TDqUserRenames> CustomRenames;
 };
 
 bool IsBlockJoin(ETestedJoinAlgo algo);

ydb/library/yql/dq/comp_nodes/dq_block_hash_join.cpp

@@ -1,6 +1,8 @@
 #pragma once
 #include "type_utils.h"
 #include <util/string/printf.h>
+#include <ydb/library/yql/dq/comp_nodes/hash_join_utils/block_layout_converter.h>
+#include <ydb/library/yql/dq/comp_nodes/hash_join_utils/neumann_hash_table.h>
 #include <yql/essentials/minikql/comp_nodes/mkql_rh_hash.h>
 
 namespace NKikimr::NMiniKQL::NJoinTable {
@@ -28,7 +30,8 @@ class TStdJoinTable {
 
     void Add(TSizedTuple tuple) {
         MKQL_ENSURE(BuiltTable.empty(), "JoinTable is built already");
-        MKQL_ENSURE(std::ssize(tuple) == TupleSize, Sprintf("tuple size promise(%i) vs actual(%i) mismatch", TupleSize, std::ssize(tuple)));
+        MKQL_ENSURE(std::ssize(tuple) == TupleSize,
+                    Sprintf("tuple size promise(%i) vs actual(%i) mismatch", TupleSize, std::ssize(tuple)));
         for (int idx = 0; idx < TupleSize; ++idx) {
             Tuples.push_back(tuple[idx]);
         }
@@ -81,4 +84,38 @@ class TStdJoinTable {
         BuiltTable;
 };
 
+class TNeumannJoinTable : NNonCopyable::TMoveOnly {
+  public:
+    struct Tuple {
+        const ui8* PackedData;
+        const ui8* OverflowBegin;
+    };
+
+    TNeumannJoinTable(const NPackedTuple::TTupleLayout* layout)
+        : Table_(layout)
+    {}
+
+    void BuildWith(IBlockLayoutConverter::TPackResult data) {
+        BuildData_ = std::move(data);
+        Table_.Build(BuildData_.PackedTuples.data(), BuildData_.Overflow.data(), BuildData_.NTuples);
+        Built_ = true;
+    }
+
+    bool Empty() {
+        return Table_.Empty();
+    }
+
+    void Lookup(Tuple row, std::invocable<Tuple> auto consume) const {
+        MKQL_ENSURE(Built_, "table must be built before lookup");
+        Table_.Apply(row.PackedData, row.OverflowBegin, [consume, this](const ui8* tuplePackedData) {
+            consume(Tuple{tuplePackedData, BuildData_.Overflow.data()});
+        });
+    }
+
+  private:
+    bool Built_ = false;
+    IBlockLayoutConverter::TPackResult BuildData_;
+    NKikimr::NMiniKQL::NPackedTuple::TNeumannHashTable<false, false> Table_;
+};
+
 } // namespace NKikimr::NMiniKQL::NJoinTable

ydb/library/yql/dq/comp_nodes/dq_hash_join_table.h

@@ -1,6 +1,7 @@
 #pragma once
 #include "dq_hash_join_table.h"
 #include <vector>
+#include <ydb/library/yql/dq/comp_nodes/hash_join_utils/block_layout_converter.h>
 #include <yql/essentials/minikql/computation/mkql_block_reader.h>
 #include <yql/essentials/minikql/computation/mkql_computation_node.h>
 #include <yql/essentials/minikql/computation/mkql_computation_node_holders.h>
@@ -12,6 +13,46 @@ struct TColumnsMetadata {
     std::vector<TType*> ColumnTypes;
 };
 
+enum class ESide { Probe, Build };
+
+template <typename T> struct TSides {
+    T Build;
+    T Probe;
+
+    T& SelectSide(ESide side) {
+        return side == ESide::Build ? Build : Probe;
+    }
+
+    const T& SelectSide(ESide side) const {
+        return side == ESide::Build ? Build : Probe;
+    }
+};
+
+
+/*
+  usage:
+  instead of copy pasting code and changing "build" to "probe", use TSides and ForEachSide to call same code twice.
+  example:
+
+void f(int buildSize, int probeSize, bool buildRequired, bool probeRequired){
+    use(buildSize + (buildRequired ? 0 : transform(buildSize)));
+    use(probeSize + (probeRequired ? 0 : transform(probeSize)));
+}
+
+vs
+
+void f(TSides<int> sizes, TSides<bool> required) {
+    ForEachSide([&](ESide side){
+        use(sizes.SelectSide(side) + (required.SelectSide(side) ? 0 : transform(sizes.SelectSide(side))));
+    });
+}
+  
+*/
+void ForEachSide(std::invocable<ESide> auto fn) {
+    fn(ESide::Build);
+    fn(ESide::Probe);
+}
+
 struct TJoinMetadata {
     TColumnsMetadata Build;
     TColumnsMetadata Probe;
@@ -21,7 +62,7 @@ struct TJoinMetadata {
 TKeyTypes KeyTypesFromColumns(const std::vector<TType*>& types, const std::vector<ui32>& keyIndexes);
 
 template <EJoinKind Kind> struct TRenamedOutput {
-    TRenamedOutput(TDqRenames renames, const std::vector<TType*>& leftColumnTypes,
+    TRenamedOutput(TDqUserRenames renames, const std::vector<TType*>& leftColumnTypes,
                    const std::vector<TType*>& rightColumnTypes)
         : OutputBuffer()
         , NullTuples(std::max(leftColumnTypes.size(), rightColumnTypes.size()), NYql::NUdf::TUnboxedValuePod{})
@@ -46,13 +87,6 @@ template <EJoinKind Kind> struct TRenamedOutput {
                     MKQL_ENSURE(tuple != nullptr, "null output row in semi/only join?");
                     for (int index = 0; index < std::ssize(Renames); ++index) {
                         auto thisRename = Renames[index];
-                        if constexpr (Kind == EJoinKind::LeftOnly || Kind == EJoinKind::LeftSemi) {
-                            MKQL_ENSURE(thisRename.Side == JoinSide::kLeft,
-                                        "rename has right column for LeftOnly or LeftSemi join??");
-                        } else {
-                            MKQL_ENSURE(thisRename.Side == JoinSide::kRight,
-                                        "rename has left column for RightOnly or RightSemi join??");
-                        }
                         OutputBuffer.push_back(tuple[thisRename.Index]);
                     }
                 };
@@ -67,7 +101,7 @@ template <EJoinKind Kind> struct TRenamedOutput {
                     }
                     for (int index = 0; index < std::ssize(Renames); ++index) {
                         auto thisRename = Renames[index];
-                        if (thisRename.Side == JoinSide::kLeft) {
+                        if (thisRename.Side == EJoinSide::kLeft) {
                             OutputBuffer.push_back(probe[thisRename.Index]);
                         } else {
                             OutputBuffer.push_back(build[thisRename.Index]);
@@ -80,13 +114,14 @@ template <EJoinKind Kind> struct TRenamedOutput {
 
   private:
     const std::vector<NYql::NUdf::TUnboxedValue> NullTuples;
-    const TDqRenames Renames;
+    const TDqUserRenames Renames;
 };
 
 // Some joins produce concatenation of 2 tuples, some produce one tuple(effectively)
-template <typename Fun>
-concept JoinMatchFun =
-    std::invocable<Fun, NJoinTable::TTuple> || std::invocable<Fun, NJoinTable::TTuple, NJoinTable::TTuple>;
+template <typename Fun, typename Tuple>
+concept JoinMatchFun = std::invocable<Fun, NJoinTable::TTuple> || std::invocable<Fun, TSides<Tuple>>;
+
+IBlockLayoutConverter::TPackResult Flatten(std::vector<IBlockLayoutConverter::TPackResult> tuples);
 
 template <typename Source, EJoinKind Kind> class TJoin : public TComputationValue<TJoin<Source, Kind>> {
     using TBase = TComputationValue<TJoin>;
@@ -120,7 +155,7 @@ template <typename Source, EJoinKind Kind> class TJoin : public TComputationValu
         return Build_.UserDataSize();
     }
 
-    EFetchResult MatchRows(TComputationContext& ctx, JoinMatchFun auto consumeOneOrTwoTuples) {
+    EFetchResult MatchRows(TComputationContext& ctx, auto consumeOneOrTwoTuples) {
         while (!Build_.Finished()) {
             auto res = Build_.ForEachRow(ctx, [&](auto tuple) { Table_.Add({tuple, tuple + Build_.UserDataSize()}); });
             switch (res) {
@@ -163,27 +198,32 @@ template <typename Source, EJoinKind Kind> class TJoin : public TComputationValu
             });
             switch (result) {
             case NYql::NUdf::EFetchStatus::Finish: {
+                int consumedTotal = 0;
                 if (Table_.UnusedTrackingOn()) {
                     if constexpr (Kind == EJoinKind::RightSemi) {
                         for (auto& v : Table_.MapView()) {
                             if (v.second.Used) {
                                 for (NJoinTable::TTuple used : v.second.Tuples) {
+
+                                    ++consumedTotal;
                                     consumeOneOrTwoTuples(used);
                                 }
                             }
                         }
                     }
                     Table_.ForEachUnused([&](NJoinTable::TTuple unused) {
                         if constexpr (Kind == EJoinKind::RightOnly) {
+                            ++consumedTotal;
                             consumeOneOrTwoTuples(unused);
                         }
                         if constexpr (Kind == EJoinKind::Exclusion || Kind == EJoinKind::Right ||
                                       Kind == EJoinKind::Full) {
+                            ++consumedTotal;
                             consumeOneOrTwoTuples(nullptr, unused);
                         }
                     });
                 }
-                return EFetchResult::One;
+                return consumedTotal == 0 ? EFetchResult::Finish : EFetchResult::One;
             }
             case NYql::NUdf::EFetchStatus::Yield: {
                 return EFetchResult::Yield;
@@ -208,4 +248,99 @@ template <typename Source, EJoinKind Kind> class TJoin : public TComputationValu
     NJoinTable::TStdJoinTable Table_;
 };
 
+template <typename Source> class TJoinPackedTuples {
+  public:
+    using TTable = NJoinTable::TNeumannJoinTable;
+
+    TJoinPackedTuples(TSides<Source> sources, NUdf::TLoggerPtr logger, TString componentName,
+                      TSides<const NPackedTuple::TTupleLayout*> layouts)
+        : Logger_(logger)
+        , LogComponent_(logger->RegisterComponent(componentName))
+        , Sources_(std::move(sources))
+        , Layouts_(layouts)
+        , Table_(Layouts_.Build)
+    {}
+
+    IBlockLayoutConverter::TPackResult Flatten(std::vector<IBlockLayoutConverter::TPackResult> tuples) {
+        IBlockLayoutConverter::TPackResult flattened;
+        flattened.NTuples = std::accumulate(tuples.begin(), tuples.end(), i64{0},
+                                            [](i64 summ, const auto& packRes) { return summ += packRes.NTuples; });
+
+        i64 totalTuplesSize = std::accumulate(tuples.begin(), tuples.end(), i64{0}, [](i64 summ, const auto& packRes) {
+            return summ += std::ssize(packRes.PackedTuples);
+        });
+        flattened.PackedTuples.reserve(totalTuplesSize);
+
+        i64 totaOverflowlSize =
+            std::accumulate(tuples.begin(), tuples.end(), i64{0},
+                            [](i64 summ, const auto& packRes) { return summ += std::ssize(packRes.Overflow); });
+        flattened.Overflow.reserve(totaOverflowlSize);
+
+        int tupleSize = Layouts_.Build->TotalRowSize;
+        for (const IBlockLayoutConverter::TPackResult& tupleBatch : tuples) {
+            Layouts_.Build->Concat(flattened.PackedTuples, flattened.Overflow,
+                                   std::ssize(flattened.PackedTuples) / tupleSize, tupleBatch.PackedTuples.data(),
+                                   tupleBatch.Overflow.data(), tupleBatch.PackedTuples.size() / tupleSize,
+                                   tupleBatch.Overflow.size());
+        }
+        return flattened;
+    }
+
+    EFetchResult MatchRows([[maybe_unused]] TComputationContext& ctx,
+                           JoinMatchFun<TTable::Tuple> auto consumeOneOrTwoTuples) {
+        while (!Sources_.Build.Finished()) {
+            FetchResult<IBlockLayoutConverter::TPackResult> var = Sources_.Build.FetchRow();
+            switch (AsStatus(var)) {
+            case NYql::NUdf::EFetchStatus::Finish: {
+                Table_.BuildWith(Flatten(BuildChunks_));
+                break;
+            }
+            case NYql::NUdf::EFetchStatus::Yield: {
+                return EFetchResult::Yield;
+            }
+            case NYql::NUdf::EFetchStatus::Ok: {
+                auto& packResult = std::get<One<IBlockLayoutConverter::TPackResult>>(var);
+                BuildChunks_.push_back(std::move(packResult.Data));
+                break;
+            }
+            default:
+                MKQL_ENSURE(false, "unreachable");
+            }
+        }
+        if (Table_.Empty()) {
+            return EFetchResult::Finish; // is it ok?
+        }
+
+        if (!Sources_.Probe.Finished()) {
+            const FetchResult<IBlockLayoutConverter::TPackResult> var = Sources_.Probe.FetchRow();
+            const NKikimr::NMiniKQL::EFetchResult resEnum = AsResult(var);
+
+            if (resEnum == EFetchResult::One) {
+                const IBlockLayoutConverter::TPackResult& thisPackResult =
+                    std::get<One<IBlockLayoutConverter::TPackResult>>(var).Data;
+                for (int index = 0; index < thisPackResult.NTuples; ++index) {
+                    const ui8* thisRow = &thisPackResult.PackedTuples[index * Layouts_.Probe->TotalRowSize];
+                    TTable::Tuple probeRow{thisRow, thisPackResult.Overflow.data()};
+                    Table_.Lookup(probeRow, [&](TTable::Tuple matchedBuildRow) {
+                        consumeOneOrTwoTuples(TSides<TTable::Tuple>{.Build = matchedBuildRow, .Probe = probeRow});
+                    });
+                }
+            }
+
+            return resEnum;
+        }
+
+        return EFetchResult::Finish;
+    }
+
+  private:
+    const NUdf::TLoggerPtr Logger_;
+    const NUdf::TLogComponentId LogComponent_;
+    TSides<Source> Sources_;
+    TSides<const NPackedTuple::TTupleLayout*> Layouts_;
+    TTable Table_;
+    IBlockLayoutConverter::TPackResult BuildData_;
+    std::vector<IBlockLayoutConverter::TPackResult> BuildChunks_;
+};
+
 } // namespace NKikimr::NMiniKQL

ydb/library/yql/dq/comp_nodes/dq_join_common.h

@@ -64,7 +64,7 @@ template <EJoinKind Kind> class TScalarHashJoinState : public TComputationValue<
                          IComputationWideFlowNode* rightFlow, const std::vector<ui32>& leftKeyColumns,
                          const std::vector<ui32>& rightKeyColumns, const std::vector<TType*>& leftColumnTypes,
                          const std::vector<TType*>& rightColumnTypes, NUdf::TLoggerPtr logger, TString componentName,
-                         TDqRenames renames)
+                         TDqUserRenames renames)
         : NKikimr::NMiniKQL::TComputationValue<TScalarHashJoinState>(memInfo)
         , Join_(memInfo, TScalarRowSource{leftFlow, leftColumnTypes}, TScalarRowSource{rightFlow, rightColumnTypes},
                 TJoinMetadata{TColumnsMetadata{rightKeyColumns, rightColumnTypes},
@@ -113,7 +113,7 @@ class TScalarHashJoinWrapper : public TStatefulWideFlowComputationNode<TScalarHa
                            TVector<TType*>&& leftColumnTypes,
                            TVector<ui32>&& leftKeyColumns,
                            TVector<TType*>&& rightColumnTypes,
-                           TVector<ui32>&& rightKeyColumns, TDqRenames renames)
+                           TVector<ui32>&& rightKeyColumns, TDqUserRenames renames)
         : TBaseComputation(mutables, nullptr, EValueRepresentation::Boxed)
         , LeftFlow_(leftFlow)
         , RightFlow_(rightFlow)
@@ -155,7 +155,7 @@ class TScalarHashJoinWrapper : public TStatefulWideFlowComputationNode<TScalarHa
     const TVector<ui32> LeftKeyColumns_;
     const TVector<TType*> RightColumnTypes_;
     const TVector<ui32> RightKeyColumns_;
-    const TDqRenames Renames_;
+    const TDqUserRenames Renames_;
 };
 
 } // namespace
@@ -216,7 +216,7 @@ IComputationWideFlowNode* WrapDqScalarHashJoin(TCallable& callable, const TCompu
     MKQL_ENSURE(rightFlow, "Expected WideFlow as a right input");
     MKQL_ENSURE(joinKind == EJoinKind::Inner, "Only inner is supported, see gh#26780 for details.");
 
-    TDqRenames renames =
+    TDqUserRenames renames =
         FromGraceFormat(TGraceJoinRenames::FromRuntimeNodes(callable.GetInput(5), callable.GetInput(6)));
     ValidateRenames(renames, joinKind, std::ssize(leftFlowItems), std::ssize(rightFlowItems));
     return new TScalarHashJoinWrapper<EJoinKind::Inner>(ctx.Mutables, leftFlow, rightFlow, std::move(joinItems),