add HashJoinProbeBlockInputStream

dbms/src/DataStreams/HashJoinProbeBlockInputStream.cpp

@@ -0,0 +1,75 @@
+// Copyright 2022 PingCAP, Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <DataStreams/HashJoinProbeBlockInputStream.h>
+#include <Flash/Mpp/getMPPTaskLog.h>
+#include <Interpreters/ExpressionActions.h>
+
+namespace DB
+{
+HashJoinProbeBlockInputStream::HashJoinProbeBlockInputStream(
+    const BlockInputStreamPtr & input,
+    const ExpressionActionsPtr & expression_,
+    const LogWithPrefixPtr & log_)
+    : log(getMPPTaskLog(log_, name))
+{
+    children.push_back(input);
+
+    bool has_join_probe_action = false;
+    if (expression_)
+    {
+        for (const auto & action : expression_->getActions())
+        {
+            if (action.type == ExpressionAction::Type::JOIN)
+            {
+                has_join_probe_action = true;
+                break;
+            }
+        }
+    }
+    if (!has_join_probe_action)
+    {
+        throw Exception("join probe expression should have join probe action", ErrorCodes::LOGICAL_ERROR);
+    }
+    expression = expression_;
+}
+
+Block HashJoinProbeBlockInputStream::getTotals()
+{
+    if (IProfilingBlockInputStream * child = dynamic_cast<IProfilingBlockInputStream *>(&*children.back()))
+    {
+        totals = child->getTotals();
+        expression->executeOnTotals(totals);
+    }
+
+    return totals;
+}
+
+Block HashJoinProbeBlockInputStream::getHeader() const
+{
+    Block res = children.back()->getHeader();
+    expression->execute(res);
+    return res;
+}
+
+Block HashJoinProbeBlockInputStream::readImpl()
+{
+    Block res = children.back()->read();
+    if (!res)
+        return res;
+    expression->execute(res);
+    return res;
+}
+
+} // namespace DB

dbms/src/DataStreams/HashJoinProbeBlockInputStream.h

@@ -0,0 +1,55 @@
+// Copyright 2022 PingCAP, Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <DataStreams/IProfilingBlockInputStream.h>
+
+namespace DB
+{
+class ExpressionActions;
+
+/** Executes a certain expression over the block.
+  * Basically the same as ExpressionBlockInputStream,
+  * but requires that there must be a join probe action in the Expression.
+  *
+  * The join probe action is different from the general expression
+  * and needs to be executed after join hash map building.
+  * We should separate it from the ExpressionBlockInputStream.
+  */
+class HashJoinProbeBlockInputStream : public IProfilingBlockInputStream
+{
+private:
+    using ExpressionActionsPtr = std::shared_ptr<ExpressionActions>;
+    static constexpr auto name = "HashJoinProbe";
+
+public:
+    HashJoinProbeBlockInputStream(
+        const BlockInputStreamPtr & input,
+        const ExpressionActionsPtr & expression_,
+        const LogWithPrefixPtr & log_);
+
+    String getName() const override { return name; }
+    Block getTotals() override;
+    Block getHeader() const override;
+
+protected:
+    Block readImpl() override;
+
+private:
+    const LogWithPrefixPtr log;
+    ExpressionActionsPtr expression;
+};
+
+} // namespace DB

dbms/src/Flash/Coprocessor/DAGQueryBlockInterpreter.cpp

@@ -20,14 +20,14 @@
 #include <DataStreams/ExpressionBlockInputStream.h>
 #include <DataStreams/FilterBlockInputStream.h>
 #include <DataStreams/HashJoinBuildBlockInputStream.h>
+#include <DataStreams/HashJoinProbeBlockInputStream.h>
 #include <DataStreams/LimitBlockInputStream.h>
 #include <DataStreams/MergeSortingBlockInputStream.h>
 #include <DataStreams/NullBlockInputStream.h>
 #include <DataStreams/ParallelAggregatingBlockInputStream.h>
 #include <DataStreams/PartialSortingBlockInputStream.h>
 #include <DataStreams/SquashingBlockInputStream.h>
 #include <DataStreams/TiRemoteBlockInputStream.h>
-#include <DataStreams/UnionBlockInputStream.h>
 #include <DataTypes/DataTypeNullable.h>
 #include <DataTypes/DataTypesNumber.h>
 #include <DataTypes/getLeastSupertype.h>
@@ -724,7 +724,7 @@ void DAGQueryBlockInterpreter::handleJoin(const tipb::Join & join, DAGPipeline &
         }
     }
     for (auto & stream : pipeline.streams)
-        stream = std::make_shared<ExpressionBlockInputStream>(stream, chain.getLastActions(), taskLogger());
+        stream = std::make_shared<HashJoinProbeBlockInputStream>(stream, chain.getLastActions(), taskLogger());
 
     /// add a project to remove all the useless column
     NamesWithAliases project_cols;