Adjust BlockMapJoinCore program builder with key column drops

ydb-platform · Sep 10, 2024 · f48070e · f48070e
1 parent 9e26730
commit f48070e
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/ydb/library/yql/minikql/mkql_program_builder.cpp b/ydb/library/yql/minikql/mkql_program_builder.cpp
@@ -5855,7 +5855,8 @@ TRuntimeNode TProgramBuilder::ScalarApply(const TArrayRef<const TRuntimeNode>& a
 }
 
 TRuntimeNode TProgramBuilder::BlockMapJoinCore(TRuntimeNode flow, TRuntimeNode dict,
-    EJoinKind joinKind, const TArrayRef<const ui32>& leftKeyColumns
+    EJoinKind joinKind, const TArrayRef<const ui32>& leftKeyColumns,
+    const TArrayRef<const ui32>& leftKeyDrops
 ) {
     if constexpr (RuntimeVersion < 51U) {
         THROW yexception() << "Runtime version (" << RuntimeVersion << ") too old for " << __func__;
@@ -5864,6 +5865,11 @@ TRuntimeNode TProgramBuilder::BlockMapJoinCore(TRuntimeNode flow, TRuntimeNode d
                 joinKind == EJoinKind::LeftSemi || joinKind == EJoinKind::LeftOnly,
                 "Unsupported join kind");
     MKQL_ENSURE(!leftKeyColumns.empty(), "At least one key column must be specified");
+    const TSet<ui32> leftKeySet(leftKeyColumns.cbegin(), leftKeyColumns.cend());
+    for (const auto& drop : leftKeyDrops) {
+        MKQL_ENSURE(leftKeySet.contains(drop),
+                    "Only key columns has to be specified in drop column set");
+    }
 
     TRuntimeNode::TList leftKeyColumnsNodes;
     leftKeyColumnsNodes.reserve(leftKeyColumns.size());
@@ -5872,6 +5878,13 @@ TRuntimeNode TProgramBuilder::BlockMapJoinCore(TRuntimeNode flow, TRuntimeNode d
             return NewDataLiteral(idx);
         });
 
+    TRuntimeNode::TList leftKeyDropsNodes;
+    leftKeyDropsNodes.reserve(leftKeyDrops.size());
+    std::transform(leftKeyDrops.cbegin(), leftKeyDrops.cend(),
+        std::back_inserter(leftKeyDropsNodes), [this](const ui32 idx) {
+            return NewDataLiteral(idx);
+        });
+
     auto returnJoinItems = ValidateBlockFlowType(flow.GetStaticType(), false);
     const auto payloadType = AS_TYPE(TDictType, dict.GetStaticType())->GetPayloadType();
     const auto payloadItemType = payloadType->IsList()
@@ -5907,6 +5920,7 @@ TRuntimeNode TProgramBuilder::BlockMapJoinCore(TRuntimeNode flow, TRuntimeNode d
     callableBuilder.Add(dict);
     callableBuilder.Add(NewDataLiteral((ui32)joinKind));
     callableBuilder.Add(NewTuple(leftKeyColumnsNodes));
+    callableBuilder.Add(NewTuple(leftKeyDropsNodes));
 
     return TRuntimeNode(callableBuilder.Build(), false);
 }

diff --git a/ydb/library/yql/minikql/mkql_program_builder.h b/ydb/library/yql/minikql/mkql_program_builder.h
@@ -256,7 +256,8 @@ class TProgramBuilder : public TTypeBuilder {
     TRuntimeNode BlockPgResolvedCall(const std::string_view& name, ui32 id,
         const TArrayRef<const TRuntimeNode>& args, TType* returnType);
     TRuntimeNode BlockMapJoinCore(TRuntimeNode flow, TRuntimeNode dict,
-        EJoinKind joinKind, const TArrayRef<const ui32>& leftKeyColumns);
+        EJoinKind joinKind, const TArrayRef<const ui32>& leftKeyColumns,
+        const TArrayRef<const ui32>& leftKeyDrops = {});
 
     //-- logical functions
     TRuntimeNode BlockNot(TRuntimeNode data);