Test: support push down filter in ut

dbms/src/Debug/MockStorage.cpp

@@ -11,8 +11,13 @@
 // 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/ExpressionBlockInputStream.h>
+#include <DataStreams/FilterBlockInputStream.h>
 #include <DataStreams/IBlockOutputStream.h>
 #include <Debug/MockStorage.h>
+#include <Flash/Coprocessor/DAGExpressionAnalyzer.h>
+#include <Flash/Coprocessor/DAGQueryInfo.h>
+#include <Flash/Coprocessor/InterpreterUtils.h>
 #include <Flash/Coprocessor/TiDBTableScan.h>
 #include <Interpreters/Context.h>
 #include <Parsers/ASTIdentifier.h>
@@ -119,10 +124,11 @@ void MockStorage::addTableDataForDeltaMerge(Context & context, const String & na
     }
 }
 
-BlockInputStreamPtr MockStorage::getStreamFromDeltaMerge(Context & context, Int64 id)
+BlockInputStreamPtr MockStorage::getStreamFromDeltaMerge(Context & context, Int64 table_id, const PushDownFilter * push_down_filter)
 {
-    auto storage = storage_delta_merge_map[id];
-    auto column_infos = table_schema_for_delta_merge[id];
+    assert(tableExistsForDeltaMerge(table_id));
+    auto storage = storage_delta_merge_map[table_id];
+    auto column_infos = table_schema_for_delta_merge[table_id];
     assert(storage);
     assert(!column_infos.empty());
     Names column_names;
@@ -134,10 +140,30 @@ BlockInputStreamPtr MockStorage::getStreamFromDeltaMerge(Context & context, Int6
     SelectQueryInfo query_info;
     query_info.query = std::make_shared<ASTSelectQuery>();
     query_info.mvcc_query_info = std::make_unique<MvccQueryInfo>(context.getSettingsRef().resolve_locks, std::numeric_limits<UInt64>::max(), scan_context);
-    BlockInputStreams ins = storage->read(column_names, query_info, context, stage, 8192, 1); // TODO: Support config max_block_size and num_streams
-
-    BlockInputStreamPtr in = ins[0];
-    return in;
+    if (push_down_filter && push_down_filter->hasValue())
+    {
+        auto analyzer = std::make_unique<DAGExpressionAnalyzer>(names_and_types_map_for_delta_merge[table_id], context);
+        query_info.dag_query = std::make_unique<DAGQueryInfo>(
+            push_down_filter->conditions,
+            analyzer->getPreparedSets(),
+            analyzer->getCurrentInputColumns(),
+            context.getTimezoneInfo());
+        auto [before_where, filter_column_name, project_after_where] = ::DB::buildPushDownFilter(*push_down_filter, *analyzer);
+        BlockInputStreams ins = storage->read(column_names, query_info, context, stage, 8192, 1); // TODO: Support config max_block_size and num_streams
+        // TODO: set num_streams, then ins.size() != 1
+        BlockInputStreamPtr in = ins[0];
+        in = std::make_shared<FilterBlockInputStream>(in, before_where, filter_column_name, "test");
+        in->setExtraInfo("push down filter");
+        in = std::make_shared<ExpressionBlockInputStream>(in, project_after_where, "test");
+        in->setExtraInfo("projection after push down filter");
+        return in;
+    }
+    else
+    {
+        BlockInputStreams ins = storage->read(column_names, query_info, context, stage, 8192, 1);
+        BlockInputStreamPtr in = ins[0];
+        return in;
+    }
 }
 
 void MockStorage::addTableInfoForDeltaMerge(const String & name, const MockColumnInfoVec & columns)

dbms/src/Debug/MockStorage.h

@@ -14,6 +14,7 @@
 #pragma once
 #include <Core/ColumnsWithTypeAndName.h>
 #include <DataStreams/IBlockInputStream.h>
+#include <Flash/Coprocessor/PushDownFilter.h>
 #include <Flash/Coprocessor/TiDBTableScan.h>
 #include <Storages/Transaction/TiDB.h>
 #include <common/types.h>
@@ -74,7 +75,7 @@ class MockStorage
 
     NamesAndTypes getNameAndTypesForDeltaMerge(Int64 table_id);
 
-    BlockInputStreamPtr getStreamFromDeltaMerge(Context & context, Int64 table_id);
+    BlockInputStreamPtr getStreamFromDeltaMerge(Context & context, Int64 table_id, const PushDownFilter * push_down_filter = nullptr);
 
     bool tableExistsForDeltaMerge(Int64 table_id);
 

dbms/src/Flash/Planner/PhysicalPlan.cpp

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include <Common/TiFlashMetrics.h>
+#include <Debug/MockStorage.h>
 #include <Flash/Coprocessor/DAGContext.h>
 #include <Flash/Coprocessor/FineGrainedShuffle.h>
 #include <Flash/Planner/ExecutorIdGenerator.h>
@@ -40,13 +41,21 @@ namespace DB
 {
 namespace
 {
-bool pushDownSelection(const PhysicalPlanNodePtr & plan, const String & executor_id, const tipb::Selection & selection)
+bool pushDownSelection(Context & context, const PhysicalPlanNodePtr & plan, const String & executor_id, const tipb::Selection & selection)
 {
     if (plan->tp() == PlanType::TableScan)
     {
         auto physical_table_scan = std::static_pointer_cast<PhysicalTableScan>(plan);
         return physical_table_scan->pushDownFilter(executor_id, selection);
     }
+    if (unlikely(plan->tp() == PlanType::MockTableScan && context.isExecutorTest()))
+    {
+        auto physical_mock_table_scan = std::static_pointer_cast<PhysicalMockTableScan>(plan);
+        if (context.mockStorage()->useDeltaMerge() && context.mockStorage()->tableExistsForDeltaMerge(physical_mock_table_scan->getLogicalTableID()))
+        {
+            return physical_mock_table_scan->pushDownFilter(context, executor_id, selection);
+        }
+    }
     return false;
 }
 
@@ -110,7 +119,7 @@ void PhysicalPlan::build(const String & executor_id, const tipb::Executor * exec
     {
         GET_METRIC(tiflash_coprocessor_executor_count, type_sel).Increment();
         auto child = popBack();
-        if (pushDownSelection(child, executor_id, executor->selection()))
+        if (pushDownSelection(context, child, executor_id, executor->selection()))
             pushBack(child);
         else
             pushBack(PhysicalFilter::build(context, executor_id, log, executor->selection(), child));

dbms/src/Flash/Planner/plans/PhysicalMockTableScan.cpp

@@ -82,10 +82,12 @@ PhysicalMockTableScan::PhysicalMockTableScan(
     const NamesAndTypes & schema_,
     const String & req_id,
     const Block & sample_block_,
-    const BlockInputStreams & mock_streams_)
+    const BlockInputStreams & mock_streams_,
+    Int64 table_id_)
     : PhysicalLeaf(executor_id_, PlanType::MockTableScan, schema_, req_id)
     , sample_block(sample_block_)
     , mock_streams(mock_streams_)
+    , table_id(table_id_)
 {}
 
 PhysicalPlanNodePtr PhysicalMockTableScan::build(
@@ -95,15 +97,15 @@ PhysicalPlanNodePtr PhysicalMockTableScan::build(
     const TiDBTableScan & table_scan)
 {
     assert(context.isTest());
-
     auto [schema, mock_streams] = mockSchemaAndStreams(context, executor_id, log, table_scan);
 
     auto physical_mock_table_scan = std::make_shared<PhysicalMockTableScan>(
         executor_id,
         schema,
         log->identifier(),
         Block(schema),
-        mock_streams);
+        mock_streams,
+        table_scan.getLogicalTableID());
     return physical_mock_table_scan;
 }
 
@@ -122,4 +124,38 @@ const Block & PhysicalMockTableScan::getSampleBlock() const
 {
     return sample_block;
 }
+
+void PhysicalMockTableScan::updateStreams(Context & context)
+{
+    mock_streams.clear();
+    assert(context.mockStorage()->tableExistsForDeltaMerge(table_id));
+    mock_streams.emplace_back(context.mockStorage()->getStreamFromDeltaMerge(context, table_id, &push_down_filter));
+}
+
+bool PhysicalMockTableScan::pushDownFilter(Context & context, const String & filter_executor_id, const tipb::Selection & selection)
+{
+    if (unlikely(hasPushDownFilter()))
+    {
+        return false;
+    }
+    push_down_filter = PushDownFilter::pushDownFilterFrom(filter_executor_id, selection);
+    updateStreams(context);
+    return true;
+}
+
+bool PhysicalMockTableScan::hasPushDownFilter() const
+{
+    return push_down_filter.hasValue();
+}
+
+const String & PhysicalMockTableScan::getPushDownFilterId() const
+{
+    assert(hasPushDownFilter());
+    return push_down_filter.executor_id;
+}
+
+Int64 PhysicalMockTableScan::getLogicalTableID() const
+{
+    return table_id;
+}
 } // namespace DB

dbms/src/Flash/Planner/plans/PhysicalMockTableScan.h

@@ -15,6 +15,7 @@
 #pragma once
 
 #include <DataStreams/IBlockInputStream.h>
+#include <Flash/Coprocessor/PushDownFilter.h>
 #include <Flash/Coprocessor/TiDBTableScan.h>
 #include <Flash/Planner/plans/PhysicalLeaf.h>
 #include <tipb/executor.pb.h>
@@ -40,17 +41,35 @@ class PhysicalMockTableScan : public PhysicalLeaf
         const NamesAndTypes & schema_,
         const String & req_id,
         const Block & sample_block_,
-        const BlockInputStreams & mock_streams_);
+        const BlockInputStreams & mock_streams_,
+        Int64 table_id_);
 
     void finalize(const Names & parent_require) override;
 
     const Block & getSampleBlock() const override;
 
+    void initStreams(Context & context);
+
+    // for delta-merge test
+    bool pushDownFilter(Context & context, const String & filter_executor_id, const tipb::Selection & selection);
+
+    bool hasPushDownFilter() const;
+
+    const String & getPushDownFilterId() const;
+
+    Int64 getLogicalTableID() const;
+
+    void updateStreams(Context & context);
+
 private:
     void transformImpl(DAGPipeline & pipeline, Context & /*context*/, size_t /*max_streams*/) override;
 
+private:
+    PushDownFilter push_down_filter;
     Block sample_block;
 
     BlockInputStreams mock_streams;
+
+    const Int64 table_id;
 };
 } // namespace DB

dbms/src/Flash/Planner/tests/gtest_physical_plan.cpp

@@ -55,11 +55,11 @@ class PhysicalPlanTestRunner : public DB::tests::ExecutorTest
                                      toNullableVec<Int64>("s4", {1, 1, {}})});
 
         context.addExchangeReceiver("exchange_r_table",
-                                    {{"s1", TiDB::TP::TypeString}, {"join_c", TiDB::TP::TypeString}},
+                                    {{"s", TiDB::TP::TypeString}, {"join_c", TiDB::TP::TypeString}},
                                     {toNullableVec<String>("s", {"banana", "banana"}),
                                      toNullableVec<String>("join_c", {"apple", "banana"})});
         context.addExchangeReceiver("exchange_l_table",
-                                    {{"s1", TiDB::TP::TypeString}, {"join_c", TiDB::TP::TypeString}},
+                                    {{"s", TiDB::TP::TypeString}, {"join_c", TiDB::TP::TypeString}},
                                     {toNullableVec<String>("s", {"banana", "banana"}),
                                      toNullableVec<String>("join_c", {"apple", "banana"})});
 

dbms/src/Flash/tests/gtest_collation.cpp

@@ -56,7 +56,7 @@ class ExecutorCollation : public DB::tests::ExecutorTest
         /// For topn
         context.addMockTable({db_name, topn_table},
                              {{topn_col, TiDB::TP::TypeString}},
-                             {toNullableVec<String>("_col", ColumnWithString{"col0-0", "col0-1", "col0-2", {}, "col0-4", {}, "col0-6", "col0-7"})});
+                             {toNullableVec<String>(topn_col, ColumnWithString{"col0-0", "col0-1", "col0-2", {}, "col0-4", {}, "col0-6", "col0-7"})});
 
         /// For projection
         context.addMockTable({db_name, proj_table},

dbms/src/Flash/tests/gtest_filter_executor.cpp

@@ -29,6 +29,7 @@ class FilterExecutorTestRunner : public DB::tests::ExecutorTest
                              {{"s1", TiDB::TP::TypeString}, {"s2", TiDB::TP::TypeString}},
                              {toNullableVec<String>("s1", {"banana", {}, "banana"}),
                               toNullableVec<String>("s2", {"apple", {}, "banana"})});
+
         context.addExchangeReceiver("exchange1",
                                     {{"s1", TiDB::TP::TypeString}, {"s2", TiDB::TP::TypeString}},
                                     {toNullableVec<String>("s1", {"banana", {}, "banana"}),
@@ -251,6 +252,90 @@ try
 }
 CATCH
 
+TEST_F(FilterExecutorTestRunner, PushDownFilter)
+try
+{
+    context.mockStorage()->setUseDeltaMerge(true);
+    context.addMockDeltaMerge({"test_db", "test_table1"},
+                              {{"i1", TiDB::TP::TypeLongLong}, {"s2", TiDB::TP::TypeString}},
+                              {toVec<Int64>("i1", {1, 2, 3}),
+                               toNullableVec<String>("s2", {"apple", {}, "banana"})});
+
+    // Do not support push down filter test for DAGQueryBlockInterpreter
+    enablePlanner(true);
+
+    auto request = context
+                       .scan("test_db", "test_table1")
+                       .filter(lt(col("i1"), lit(Field(static_cast<Int64>(2)))))
+                       .build(context);
+
+    {
+        String expected = R"(
+Expression: <final projection>
+ Expression: <projection after push down filter>
+  Filter: <push down filter>
+   DeltaMergeSegmentThread)";
+        executeInterpreterWithDeltaMerge(expected, request, 10);
+    }
+    executeAndAssertColumnsEqual(
+        request,
+        {toNullableVec<Int64>({1}),
+         toNullableVec<String>({"apple"})});
+
+
+    request = context
+                  .scan("test_db", "test_table1")
+                  .filter(lt(col("i1"), lit(Field(static_cast<Int64>(3)))))
+                  .build(context);
+
+    executeAndAssertColumnsEqual(
+        request,
+        {toNullableVec<Int64>({1, 2}),
+         toNullableVec<String>({"apple", {}})});
+
+    for (size_t i = 4; i < 10; ++i)
+    {
+        request = context
+                      .scan("test_db", "test_table1")
+                      .filter(lt(col("i1"), lit(Field(static_cast<Int64>(i)))))
+                      .build(context);
+
+        executeAndAssertColumnsEqual(
+            request,
+            {toNullableVec<Int64>({1, 2, 3}),
+             toNullableVec<String>({"apple", {}, "banana"})});
+    }
+
+    for (size_t i = 0; i < 10; ++i)
+    {
+        request = context
+                      .scan("test_db", "test_table1")
+                      .filter(gt(col("i1"), lit(Field(static_cast<Int64>(-i)))))
+                      .build(context);
+
+        executeAndAssertColumnsEqual(
+            request,
+            {toNullableVec<Int64>({1, 2, 3}),
+             toNullableVec<String>({"apple", {}, "banana"})});
+    }
+
+    for (size_t i = 0; i < 10; ++i)
+    {
+        request = context
+                      .scan("test_db", "test_table1")
+                      .filter(gt(col("i1"), lit(Field(static_cast<Int64>(-i)))))
+                      .project({col("i1")})
+                      .build(context);
+
+        executeAndAssertColumnsEqual(
+            request,
+            {toNullableVec<Int64>({1, 2, 3})});
+    }
+
+    context.mockStorage()->setUseDeltaMerge(false);
+}
+CATCH
+
 /// TODO: more functions.
 
 } // namespace tests

dbms/src/Flash/tests/gtest_join_executor.cpp

@@ -48,12 +48,12 @@ class JoinExecutorTestRunner : public DB::tests::ExecutorTest
                               toVec<String>("join_c", {"apple", "banana"})});
 
         context.addExchangeReceiver("exchange_r_table",
-                                    {{"s1", TiDB::TP::TypeString}, {"join_c", TiDB::TP::TypeString}},
+                                    {{"s", TiDB::TP::TypeString}, {"join_c", TiDB::TP::TypeString}},
                                     {toNullableVec<String>("s", {"banana", "banana"}),
                                      toNullableVec<String>("join_c", {"apple", "banana"})});
 
         context.addExchangeReceiver("exchange_l_table",
-                                    {{"s1", TiDB::TP::TypeString}, {"join_c", TiDB::TP::TypeString}},
+                                    {{"s", TiDB::TP::TypeString}, {"join_c", TiDB::TP::TypeString}},
                                     {toNullableVec<String>("s", {"banana", "banana"}),
                                      toNullableVec<String>("join_c", {"apple", "banana"})});
     }
@@ -354,11 +354,16 @@ try
     executeAndAssertColumnsEqual(cast_request(), {createNullableColumn<Decimal256>(std::make_tuple(65, 0), {"0.12"}, {0}), createNullableColumn<Decimal256>(std::make_tuple(65, 0), {"0.12"}, {0})});
 
     /// datetime(1970-01-01 00:00:01) == timestamp(1970-01-01 00:00:01)
-    context.addMockTable("cast", "t1", {{"a", TiDB::TP::TypeDatetime}}, {createDateTimeColumn({{{1970, 1, 1, 0, 0, 1, 0}}}, 6)});
+    context.addMockTable("cast", "t1", {{"datetime", TiDB::TP::TypeDatetime}}, {createDateTimeColumn({{{1970, 1, 1, 0, 0, 1, 0}}}, 6)});
 
-    context.addMockTable("cast", "t2", {{"a", TiDB::TP::TypeTimestamp}}, {createDateTimeColumn({{{1970, 1, 1, 0, 0, 1, 0}}}, 6)});
+    context.addMockTable("cast", "t2", {{"datetime", TiDB::TP::TypeTimestamp}}, {createDateTimeColumn({{{1970, 1, 1, 0, 0, 1, 0}}}, 6)});
 
-    executeAndAssertColumnsEqual(cast_request(), {createDateTimeColumn({{{1970, 1, 1, 0, 0, 1, 0}}}, 0), createDateTimeColumn({{{1970, 1, 1, 0, 0, 1, 0}}}, 0)});
+    auto cast_request_1 = [&]() {
+        return context.scan("cast", "t1")
+            .join(context.scan("cast", "t2"), tipb::JoinType::TypeInnerJoin, {col("datetime")})
+            .build(context);
+    };
+    executeAndAssertColumnsEqual(cast_request_1(), {createDateTimeColumn({{{1970, 1, 1, 0, 0, 1, 0}}}, 0), createDateTimeColumn({{{1970, 1, 1, 0, 0, 1, 0}}}, 0)});
 }
 CATCH