diff --git a/contrib/tipb b/contrib/tipb index fe7badb76b6..474dc14add1 160000 --- a/contrib/tipb +++ b/contrib/tipb @@ -1 +1 @@ -Subproject commit fe7badb76b660abaedd8706accd809809c475280 +Subproject commit 474dc14add1ef2d1c77bc6f0a42f93546583323f diff --git a/dbms/src/Debug/MockStorage.cpp b/dbms/src/Debug/MockStorage.cpp index 9a2c1d04db3..57a7599bd40 100644 --- a/dbms/src/Debug/MockStorage.cpp +++ b/dbms/src/Debug/MockStorage.cpp @@ -159,7 +159,7 @@ bool shouldEnableMultiStageLateMaterializationForMockDeltaMerge( const TiDB::ColumnInfos & scan_column_infos) { const auto & settings = context.getSettingsRef(); - if (!settings.dt_enable_multi_stage_late_materialization) + if (settings.dt_enable_multi_stage_late_materialization == 0) return false; if (filter_conditions == nullptr || !filter_conditions->hasValue()) return false; @@ -376,7 +376,8 @@ BlockInputStreamPtr MockStorage::getStreamFromDeltaMerge( std::vector runtime_filter_ids, int rf_max_wait_time_ms, const google::protobuf::RepeatedPtrField * pushed_down_filters, - const TiDB::ColumnInfos * table_scan_column_infos) + const TiDB::ColumnInfos * table_scan_column_infos, + const DM::MultiStageLateMaterializationTopNDescriptionPtr & multi_stage_late_materialization_topn) { static const google::protobuf::RepeatedPtrField empty_pushed_down_filters{}; static const auto empty_ann_query_info = tipb::ANNQueryInfo{}; @@ -400,6 +401,7 @@ BlockInputStreamPtr MockStorage::getStreamFromDeltaMerge( runtime_filter_ids, rf_max_wait_time_ms, context.getTimezoneInfo()); + query_info.multi_stage_late_materialization_topn = multi_stage_late_materialization_topn; BlockInputStreams ins = storage->read( column_names, query_info, @@ -476,7 +478,8 @@ void MockStorage::buildExecFromDeltaMerge( int rf_max_wait_time_ms, const google::protobuf::RepeatedPtrField * pushed_down_filters, const String & table_scan_executor_id, - const TiDB::ColumnInfos * table_scan_column_infos) + const TiDB::ColumnInfos * table_scan_column_infos, + const DM::MultiStageLateMaterializationTopNDescriptionPtr & multi_stage_late_materialization_topn) { static const google::protobuf::RepeatedPtrField empty_pushed_down_filters{}; static const auto empty_ann_query_info = tipb::ANNQueryInfo{}; @@ -501,9 +504,11 @@ void MockStorage::buildExecFromDeltaMerge( if (enable_multi_stage_late_materialization) { multi_stage_late_materialization_runtime_stats - = std::make_shared(); + = std::make_shared( + fmt::format("mock table_scan_executor_id={}", table_scan_executor_id)); if (auto * dag_context = context.getDAGContext(); dag_context != nullptr && !table_scan_executor_id.empty()) { + dag_context->scan_context_map[table_scan_executor_id] = query_info.mvcc_query_info->scan_context; dag_context->setExecutorRowsOverride( table_scan_executor_id, std::shared_ptr>( @@ -513,8 +518,10 @@ void MockStorage::buildExecFromDeltaMerge( filter_conditions->executor_id, std::shared_ptr>( multi_stage_late_materialization_runtime_stats, - &multi_stage_late_materialization_runtime_stats->stage1_output_rows)); + &multi_stage_late_materialization_runtime_stats->final_rest_input_rows)); } + query_info.mvcc_query_info->scan_context->setMultiStageLateMaterializationRuntimeStats( + multi_stage_late_materialization_runtime_stats); } query_info.dag_query = std::make_unique( filter_conditions->conditions, @@ -526,6 +533,7 @@ void MockStorage::buildExecFromDeltaMerge( context.getTimezoneInfo()); query_info.enable_multi_stage_late_materialization = enable_multi_stage_late_materialization; query_info.multi_stage_late_materialization_runtime_stats = multi_stage_late_materialization_runtime_stats; + query_info.multi_stage_late_materialization_topn = multi_stage_late_materialization_topn; storage->read( exec_context_, group_builder, @@ -607,6 +615,8 @@ void MockStorage::addTableInfoForDeltaMerge(const String & name, const MockColum TiDB::ColumnInfo ret; ret.name = column.name; ret.tp = column.type; + ret.collate = column.collate; + ret.elems = column.elems; if (!column.nullable) ret.setNotNullFlag(); @@ -890,6 +900,7 @@ TiDB::ColumnInfos mockColumnInfosToTiDBColumnInfos(const MockColumnInfoVec & moc column_info.name = mock_column_info.name; column_info.tp = mock_column_info.type; column_info.collate = mock_column_info.collate; + column_info.elems = mock_column_info.elems; column_info.id = col_id++; // TODO: find a way to assign decimal field's flen. if (column_info.tp == TiDB::TP::TypeNewDecimal) diff --git a/dbms/src/Debug/MockStorage.h b/dbms/src/Debug/MockStorage.h index 1259b9ad1dc..84ea3a6919d 100644 --- a/dbms/src/Debug/MockStorage.h +++ b/dbms/src/Debug/MockStorage.h @@ -20,12 +20,15 @@ #include #include #include +#include #include #include #include #include #include +#include +#include namespace DB { @@ -36,10 +39,26 @@ struct SelectQueryInfo; struct MockColumnInfo { + MockColumnInfo() = default; + + MockColumnInfo( + String name_, + TiDB::TP type_, + bool nullable_ = true, + Poco::Dynamic::Var collate_ = {}, + std::vector> elems_ = {}) + : name(std::move(name_)) + , type(type_) + , nullable(nullable_) + , collate(std::move(collate_)) + , elems(std::move(elems_)) + {} + String name; TiDB::TP type; bool nullable = true; Poco::Dynamic::Var collate{}; // default empty means no collation. + std::vector> elems; }; using MockColumnInfoVec = std::vector; using TableInfo = TiDB::TableInfo; @@ -107,7 +126,8 @@ class MockStorage std::vector runtime_filter_ids = std::vector(), int rf_max_wait_time_ms = 0, const google::protobuf::RepeatedPtrField * pushed_down_filters = nullptr, - const TiDB::ColumnInfos * scan_column_infos = nullptr); + const TiDB::ColumnInfos * scan_column_infos = nullptr, + const DM::MultiStageLateMaterializationTopNDescriptionPtr & multi_stage_late_materialization_topn = nullptr); void buildExecFromDeltaMerge( PipelineExecutorContext & exec_context_, @@ -121,7 +141,8 @@ class MockStorage int rf_max_wait_time_ms = 0, const google::protobuf::RepeatedPtrField * pushed_down_filters = nullptr, const String & table_scan_executor_id = "", - const TiDB::ColumnInfos * scan_column_infos = nullptr); + const TiDB::ColumnInfos * scan_column_infos = nullptr, + const DM::MultiStageLateMaterializationTopNDescriptionPtr & multi_stage_late_materialization_topn = nullptr); bool tableExistsForDeltaMerge(Int64 table_id); diff --git a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp index 0a5307cb0c3..154450989a2 100644 --- a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp +++ b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp @@ -382,11 +382,13 @@ DAGStorageInterpreter::DAGStorageInterpreter( Context & context_, const TiDBTableScan & table_scan_, const FilterConditions & filter_conditions_, - size_t max_streams_) + size_t max_streams_, + const DM::MultiStageLateMaterializationTopNDescriptionPtr & multi_stage_late_materialization_topn_) : context(context_) , table_scan(table_scan_) , filter_conditions(filter_conditions_) , max_streams(max_streams_) + , multi_stage_late_materialization_topn(multi_stage_late_materialization_topn_) , log(Logger::get(context.getDAGContext()->log ? context.getDAGContext()->log->identifier() : "")) , logical_table_id(table_scan.getLogicalTableID()) , tmt(context.getTMTContext()) @@ -1044,7 +1046,8 @@ std::unordered_map DAGStorageInterpreter::generateSele if (enable_multi_stage_late_materialization) { multi_stage_late_materialization_runtime_stats - = std::make_shared(); + = std::make_shared( + fmt::format("{} table_scan_executor_id={}", log->identifier(), table_scan.getTableScanExecutorID())); dagContext().setExecutorRowsOverride( table_scan.getTableScanExecutorID(), std::shared_ptr>( @@ -1054,7 +1057,13 @@ std::unordered_map DAGStorageInterpreter::generateSele filter_conditions.executor_id, std::shared_ptr>( multi_stage_late_materialization_runtime_stats, - &multi_stage_late_materialization_runtime_stats->stage1_output_rows)); + &multi_stage_late_materialization_runtime_stats->final_rest_input_rows)); + if (auto scan_context_it = dagContext().scan_context_map.find(table_scan.getTableScanExecutorID()); + scan_context_it != dagContext().scan_context_map.end() && scan_context_it->second != nullptr) + { + scan_context_it->second->setMultiStageLateMaterializationRuntimeStats( + multi_stage_late_materialization_runtime_stats); + } } auto create_query_info = [&](Int64 table_id) -> SelectQueryInfo { @@ -1074,6 +1083,7 @@ std::unordered_map DAGStorageInterpreter::generateSele query_info.is_fast_scan = table_scan.isFastScan(); query_info.enable_multi_stage_late_materialization = enable_multi_stage_late_materialization; query_info.multi_stage_late_materialization_runtime_stats = multi_stage_late_materialization_runtime_stats; + query_info.multi_stage_late_materialization_topn = multi_stage_late_materialization_topn; return query_info; }; RUNTIME_CHECK_MSG(mvcc_query_info->scan_context != nullptr, "Unexpected null scan_context"); @@ -1790,7 +1800,7 @@ std::pair> DAGStorageInterpreter::getColumnsForTableSc bool DAGStorageInterpreter::shouldEnableMultiStageLateMaterialization() const { - if (!context.getSettingsRef().dt_enable_multi_stage_late_materialization) + if (context.getSettingsRef().dt_enable_multi_stage_late_materialization == 0) return false; auto disable = [&](const String & reason) { diff --git a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.h b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.h index 40dcdd72631..dc0bf1e3e0d 100644 --- a/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.h +++ b/dbms/src/Flash/Coprocessor/DAGStorageInterpreter.h @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -47,7 +48,8 @@ class DAGStorageInterpreter Context & context_, const TiDBTableScan & table_scan, const FilterConditions & filter_conditions_, - size_t max_streams_); + size_t max_streams_, + const DM::MultiStageLateMaterializationTopNDescriptionPtr & multi_stage_late_materialization_topn_ = nullptr); ~DAGStorageInterpreter(); @@ -146,6 +148,7 @@ class DAGStorageInterpreter const TiDBTableScan & table_scan; const FilterConditions & filter_conditions; const size_t max_streams; + const DM::MultiStageLateMaterializationTopNDescriptionPtr multi_stage_late_materialization_topn; LoggerPtr log; /// derived from other members, doesn't change during DAGStorageInterpreter's lifetime diff --git a/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.cpp b/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.cpp index 6597e82a52f..72b013838da 100644 --- a/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.cpp +++ b/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.cpp @@ -73,7 +73,8 @@ std::pair mockSchemaAndStreams( table_scan.getRuntimeFilterIDs(), 10000, &table_scan.getPushedDownFilters(), - use_table_scan_columns ? &table_scan.getColumns() : nullptr)); + use_table_scan_columns ? &table_scan.getColumns() : nullptr, + nullptr)); } else { @@ -184,7 +185,8 @@ void PhysicalMockTableScan::buildPipelineExecGroupImpl( rf_max_wait_time_ms, &pushed_down_filters, execId(), - use_table_scan_columns_for_delta_merge ? &table_scan_columns : nullptr); + use_table_scan_columns_for_delta_merge ? &table_scan_columns : nullptr, + multi_stage_late_materialization_topn); for (size_t i = 0; i < group_builder.concurrency(); ++i) { if (auto * source_op = dynamic_cast(group_builder.getCurBuilder(i).source_op.get())) @@ -240,7 +242,8 @@ bool PhysicalMockTableScan::setFilterConditions( runtime_filter_ids, 10000, &pushed_down_filters, - use_table_scan_columns_for_delta_merge ? &table_scan_columns : nullptr)); + use_table_scan_columns_for_delta_merge ? &table_scan_columns : nullptr, + multi_stage_late_materialization_topn)); return true; } @@ -261,6 +264,12 @@ const String & PhysicalMockTableScan::getFilterConditionsId() const return filter_conditions.executor_id; } +void PhysicalMockTableScan::setMultiStageLateMaterializationTopN( + const DM::MultiStageLateMaterializationTopNDescriptionPtr & topn) +{ + multi_stage_late_materialization_topn = topn; +} + void PhysicalMockTableScan::buildRuntimeFilterInLocalStream(Context & context) { for (const auto & local_stream : mock_streams) diff --git a/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.h b/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.h index 2b2c2d29548..246297727c1 100644 --- a/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.h +++ b/dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.h @@ -19,6 +19,7 @@ #include #include #include +#include #include namespace DB @@ -63,6 +64,12 @@ class PhysicalMockTableScan : public PhysicalLeaf const String & getFilterConditionsId() const; + const TiDB::ColumnInfos & getTableScanColumns() const { return table_scan_columns; } + + const google::protobuf::RepeatedPtrField & getPushedDownFilters() const { return pushed_down_filters; } + + void setMultiStageLateMaterializationTopN(const DM::MultiStageLateMaterializationTopNDescriptionPtr & topn); + private: void buildBlockInputStreamImpl(DAGPipeline & pipeline, Context & /*context*/, size_t /*max_streams*/) override; @@ -92,6 +99,8 @@ class PhysicalMockTableScan : public PhysicalLeaf TiDB::ColumnInfos table_scan_columns; + DM::MultiStageLateMaterializationTopNDescriptionPtr multi_stage_late_materialization_topn; + bool use_table_scan_columns_for_delta_merge; const int rf_max_wait_time_ms = 10000; diff --git a/dbms/src/Flash/Planner/Plans/PhysicalTableScan.cpp b/dbms/src/Flash/Planner/Plans/PhysicalTableScan.cpp index 6165bbe448e..2076961ef08 100644 --- a/dbms/src/Flash/Planner/Plans/PhysicalTableScan.cpp +++ b/dbms/src/Flash/Planner/Plans/PhysicalTableScan.cpp @@ -114,7 +114,12 @@ void PhysicalTableScan::buildBlockInputStreamImpl(DAGPipeline & pipeline, Contex } else { - DAGStorageInterpreter storage_interpreter(context, tidb_table_scan, filter_conditions, max_streams); + DAGStorageInterpreter storage_interpreter( + context, + tidb_table_scan, + filter_conditions, + max_streams, + multi_stage_late_materialization_topn); storage_interpreter.execute(pipeline); } buildProjection(pipeline); @@ -137,7 +142,12 @@ void PhysicalTableScan::buildPipeline( } else { - DAGStorageInterpreter storage_interpreter(context, tidb_table_scan, filter_conditions, context.getMaxStreams()); + DAGStorageInterpreter storage_interpreter( + context, + tidb_table_scan, + filter_conditions, + context.getMaxStreams(), + multi_stage_late_materialization_topn); storage_interpreter.execute(exec_context, pipeline_exec_builder); } buildProjection(exec_context, pipeline_exec_builder); @@ -214,4 +224,10 @@ const String & PhysicalTableScan::getFilterConditionsId() const RUNTIME_CHECK(hasFilterConditions()); return filter_conditions.executor_id; } + +void PhysicalTableScan::setMultiStageLateMaterializationTopN( + const DM::MultiStageLateMaterializationTopNDescriptionPtr & topn) +{ + multi_stage_late_materialization_topn = topn; +} } // namespace DB diff --git a/dbms/src/Flash/Planner/Plans/PhysicalTableScan.h b/dbms/src/Flash/Planner/Plans/PhysicalTableScan.h index 284b75d0c73..bd654def8a8 100644 --- a/dbms/src/Flash/Planner/Plans/PhysicalTableScan.h +++ b/dbms/src/Flash/Planner/Plans/PhysicalTableScan.h @@ -18,6 +18,7 @@ #include #include #include +#include #include namespace DB @@ -48,6 +49,10 @@ class PhysicalTableScan : public PhysicalLeaf const String & getFilterConditionsId() const; + const TiDBTableScan & getTiDBTableScan() const { return tidb_table_scan; } + + void setMultiStageLateMaterializationTopN(const DM::MultiStageLateMaterializationTopNDescriptionPtr & topn); + void buildPipeline(PipelineBuilder & builder, Context & context, PipelineExecutorContext & exec_context) override; private: @@ -67,6 +72,8 @@ class PhysicalTableScan : public PhysicalLeaf TiDBTableScan tidb_table_scan; + DM::MultiStageLateMaterializationTopNDescriptionPtr multi_stage_late_materialization_topn; + Block sample_block; PipelineExecGroupBuilder pipeline_exec_builder; diff --git a/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp b/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp index e6b3e7022f8..3249b3d5e6d 100644 --- a/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp +++ b/dbms/src/Flash/Planner/Plans/PhysicalTopN.cpp @@ -13,18 +13,117 @@ // limitations under the License. #include +#include #include #include #include +#include #include #include #include #include +#include +#include #include #include +#include +#include namespace DB { +namespace +{ +DM::MultiStageLateMaterializationTopNDescriptionPtr tryBuildMultiStageLateMaterializationTopN( + const tipb::TopN & top_n, + const TiDB::ColumnInfos & table_scan_columns, + const google::protobuf::RepeatedPtrField & pushed_down_filters, + bool has_filter_conditions, + const LoggerPtr & log) +{ + auto disable = [&](const String & reason) -> DM::MultiStageLateMaterializationTopNDescriptionPtr { + LOG_DEBUG(log, "Disable TopN-enhanced multi-stage late materialization, reason={}", reason); + return nullptr; + }; + + if (!has_filter_conditions) + return disable("no residual filter conditions"); + if (pushed_down_filters.empty()) + return disable("no stage0 pushed down filters"); + if (top_n.limit() == 0) + return disable("topk is zero"); + if (top_n.limit() > DM::multi_stage_late_materialization_topn_max_topk) + return disable(fmt::format("topk is too large: {}", top_n.limit())); + if (top_n.order_by_size() > static_cast(DM::multi_stage_late_materialization_topn_max_order_by_columns)) + return disable(fmt::format("too many order by columns: {}", top_n.order_by_size())); + if (top_n.partition_by_size() != 0) + return disable( + fmt::format("partition TopN is unsupported, partition_by_columns={}", top_n.partition_by_size())); + + auto desc = std::make_shared(); + desc->topk = top_n.limit(); + desc->order_by_columns.reserve(top_n.order_by_size()); + + for (const auto & by_item : top_n.order_by()) + { + if (!isColumnExpr(by_item.expr())) + return disable("order by expression is not direct ColumnRef"); + + // The ColumnRef index is relative to TopN's input schema. TopN-enhanced MSLM is only attached + // when the child is a direct TableScan, whose output schema is generated from table_scan_columns + // with the same order, so the index can be used to locate the original table scan column here. + const auto column_index = decodeDAGInt64(by_item.expr().val()); + if (column_index < 0 || column_index >= static_cast(table_scan_columns.size())) + return disable(fmt::format( + "order by ColumnRef index out of range, index={}, table_scan_column_size={}", + column_index, + table_scan_columns.size())); + + const auto & column = table_scan_columns[column_index]; + if (column.hasGeneratedColumnFlag()) + return disable(fmt::format("order by generated column is unsupported, column_id={}", column.id)); + if (column.id == ExtraTableIDColumnID || column.id == ExtraCommitTSColumnID) + return disable(fmt::format("order by virtual column is unsupported, column_id={}", column.id)); + + desc->order_by_columns.push_back( + DM::MultiStageLateMaterializationTopNOrderByColumn{column.id, by_item.desc() ? -1 : 1}); + } + + return desc; +} + +void tryAttachMultiStageLateMaterializationTopN( + const tipb::TopN & top_n, + const PhysicalPlanNodePtr & child, + const LoggerPtr & log) +{ + if (auto table_scan = std::dynamic_pointer_cast(child)) + { + auto topn = tryBuildMultiStageLateMaterializationTopN( + top_n, + table_scan->getTiDBTableScan().getColumns(), + table_scan->getTiDBTableScan().getPushedDownFilters(), + table_scan->hasFilterConditions(), + log); + table_scan->setMultiStageLateMaterializationTopN(topn); + return; + } + + if (auto table_scan = std::dynamic_pointer_cast(child)) + { + auto topn = tryBuildMultiStageLateMaterializationTopN( + top_n, + table_scan->getTableScanColumns(), + table_scan->getPushedDownFilters(), + table_scan->hasFilterConditions(), + log); + table_scan->setMultiStageLateMaterializationTopN(topn); + return; + } + + LOG_DEBUG(log, "Disable TopN-enhanced multi-stage late materialization, reason=child is not direct table scan"); +} +} // namespace + PhysicalPlanNodePtr PhysicalTopN::build( const Context & context, const String & executor_id, @@ -45,6 +144,10 @@ PhysicalPlanNodePtr PhysicalTopN::build( auto order_columns = analyzer.buildOrderColumns(before_sort_actions, top_n.order_by()); SortDescription order_descr = getSortDescription(order_columns, top_n.order_by()); + if (context.getSettingsRef().dt_enable_multi_stage_late_materialization >= 2) + tryAttachMultiStageLateMaterializationTopN(top_n, child, log); + else + LOG_DEBUG(log, "Disable TopN-enhanced multi-stage late materialization, reason=setting disabled"); auto physical_top_n = std::make_shared( executor_id, diff --git a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp index 8e6019b7774..c5ce0a57e04 100644 --- a/dbms/src/Flash/tests/gtest_executors_with_dm.cpp +++ b/dbms/src/Flash/tests/gtest_executors_with_dm.cpp @@ -14,11 +14,13 @@ #include #include +#include #include #include #include #include #include +#include #include #include #include @@ -145,10 +147,42 @@ class ExecutorsWithDMTestRunner : public DB::tests::ExecutorTest namespace { +constexpr UInt64 MSLMModeDisabled = 0; +constexpr UInt64 MSLMModeSelection = 1; +constexpr UInt64 MSLMModeTopN = 2; + +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationSettingKeepsUInt64Mode) +try +{ + Settings settings; + ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeTopN); + + settings.dt_enable_multi_stage_late_materialization = MSLMModeDisabled; + ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeDisabled); + + settings.dt_enable_multi_stage_late_materialization = MSLMModeSelection; + ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeSelection); + + settings.dt_enable_multi_stage_late_materialization = MSLMModeTopN; + ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), MSLMModeTopN); + + settings.dt_enable_multi_stage_late_materialization = 999; + ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), 999); + + settings.set("dt_enable_multi_stage_late_materialization", Field(static_cast(3))); + ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), 3); + + settings.set("dt_enable_multi_stage_late_materialization", String("999")); + ASSERT_EQ(settings.dt_enable_multi_stage_late_materialization.get(), 999); +} +CATCH + tipb::TableScan * findMutableTableScan(tipb::Executor * executor) { switch (executor->tp()) { + case tipb::ExecType::TypeTopN: + return findMutableTableScan(executor->mutable_topn()->mutable_child()); case tipb::ExecType::TypeTableScan: return executor->mutable_tbl_scan(); case tipb::ExecType::TypeSelection: @@ -164,6 +198,8 @@ tipb::Selection * findMutableSelection(tipb::Executor * executor) { switch (executor->tp()) { + case tipb::ExecType::TypeTopN: + return findMutableSelection(executor->mutable_topn()->mutable_child()); case tipb::ExecType::TypeSelection: return executor->mutable_selection(); case tipb::ExecType::TypeProjection: @@ -294,6 +330,63 @@ void assertColumnsEqual(const ColumnsWithTypeAndName & expected, const ColumnsWi << getColumnsContent(actual); } +DM::MultiStageLateMaterializationRuntimeStatsPtr getMultiStageRuntimeStats(DAGContext & dag_context) +{ + DM::MultiStageLateMaterializationRuntimeStatsPtr stats; + for (const auto & scan_context_entry : dag_context.scan_context_map) + { + const auto & scan_context = scan_context_entry.second; + if (scan_context == nullptr || scan_context->multi_stage_late_materialization_runtime_stats == nullptr) + continue; + RUNTIME_CHECK_MSG(stats == nullptr, "Found multiple table scan contexts with multi-stage runtime stats"); + stats = scan_context->multi_stage_late_materialization_runtime_stats; + } + RUNTIME_CHECK_MSG(stats != nullptr, "Cannot find table scan context with multi-stage runtime stats"); + return stats; +} + +void assertTopNEnhancedMultiStageRowsOverride( + DAGContext & dag_context, + UInt64 expected_table_scan_rows, + UInt64 min_selection_rows) +{ + auto table_scan_rows = dag_context.getExecutorRowsOverride("table_scan_0"); + auto selection_rows = dag_context.getExecutorRowsOverride("selection_1"); + ASSERT_NE(table_scan_rows, nullptr); + ASSERT_NE(selection_rows, nullptr); + ASSERT_EQ(table_scan_rows->load(), expected_table_scan_rows); + ASSERT_GE(selection_rows->load(), min_selection_rows); + ASSERT_LT(selection_rows->load(), expected_table_scan_rows); + + auto stats = getMultiStageRuntimeStats(dag_context); + ASSERT_NE(stats, nullptr); + ASSERT_TRUE(stats->topn_enabled.load()); + ASSERT_GT(stats->running_topn_input_rows.load(), 0); + ASSERT_EQ(selection_rows->load(), stats->final_rest_input_rows.load()); +} + +void assertTopNEnhancedMultiStageMatchesDisabled( + ExecutorsWithDMTestRunner & runner, + MockDAGRequestContext & context, + const std::shared_ptr & request, + const String & log_identifier, + UInt64 expected_table_scan_rows, + UInt64 min_selection_rows) +{ + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; + DAGContext disabled_dag_context(*request, log_identifier, /*concurrency=*/1); + auto expected = runner.executeStreams(&disabled_dag_context); + ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr); + ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr); + + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeTopN; + DAGContext enabled_dag_context(*request, log_identifier, /*concurrency=*/1); + auto actual = runner.executeStreams(&enabled_dag_context); + + assertColumnsEqual(expected, actual); + assertTopNEnhancedMultiStageRowsOverride(enabled_dag_context, expected_table_scan_rows, min_selection_rows); +} + void rewriteTableScanColumnAsCommitTS(tipb::DAGRequest * request, Int32 column_index) { auto * table_scan = findMutableTableScan(request->mutable_root_executor()); @@ -364,6 +457,46 @@ std::shared_ptr buildDAGRequestWithPushedDownFilterAndProjecti return request; } +std::shared_ptr buildDAGRequestWithPushedDownFilterAndTopN( + MockDAGRequestContext & context, + const String & table_name, + const ASTPtr & pushed_down_filter, + const ASTPtr & residual_filter, + MockOrderByItemVec order_by_items, + UInt64 limit) +{ + auto pushed_down_filter_request = context.scan("test_db", table_name).filter(pushed_down_filter).build(context); + RUNTIME_CHECK(pushed_down_filter_request->root_executor().tp() == tipb::ExecType::TypeSelection); + RUNTIME_CHECK(pushed_down_filter_request->root_executor().selection().conditions_size() == 1); + const auto & pushed_down_filter_expr = pushed_down_filter_request->root_executor().selection().conditions(0); + + auto request = context.scan("test_db", table_name) + .filter(residual_filter) + .topN(std::move(order_by_items), static_cast(limit)) + .build(context); + *findMutableTableScan(request->mutable_root_executor())->add_pushed_down_filter_conditions() + = pushed_down_filter_expr; + return request; +} + +std::shared_ptr buildDAGRequestWithPushedDownFilterAndTopN( + MockDAGRequestContext & context, + const String & table_name, + const ASTPtr & pushed_down_filter, + const ASTPtr & residual_filter, + const String & order_by_column, + bool is_desc, + UInt64 limit) +{ + return buildDAGRequestWithPushedDownFilterAndTopN( + context, + table_name, + pushed_down_filter, + residual_filter, + MockOrderByItemVec{{order_by_column, is_desc}}, + limit); +} + ASTPtr buildAllGreaterThanMinusOneFilter(size_t begin_col_id, size_t end_col_id) { RUNTIME_CHECK(begin_col_id < end_col_id); @@ -404,6 +537,74 @@ void addMockMultiStageLateMaterializationTable( context.addMockDeltaMerge({"test_db", table_name}, columns, data, /*concurrency_hint=*/4); } +void addMockMultiStageLateMaterializationTopNOrderByTable( + MockDAGRequestContext & context, + const String & table_name, + const MockColumnInfo & order_by_info, + ColumnWithTypeAndName order_by_column) +{ + constexpr size_t rows = 32; + MockColumnInfoVec columns; + ColumnsWithTypeAndName data; + + columns.push_back({"c0", TiDB::TP::TypeLongLong}); + columns.push_back(order_by_info); + columns.push_back({"c1", TiDB::TP::TypeLongLong}); + + std::vector c0_values; + std::vector c1_values; + c0_values.reserve(rows); + c1_values.reserve(rows); + for (size_t row = 0; row < rows; ++row) + { + c0_values.push_back(static_cast(row)); + c1_values.push_back(static_cast(row % 8)); + } + + data.emplace_back(toVec("c0", c0_values)); + data.emplace_back(std::move(order_by_column)); + data.emplace_back(toVec("c1", c1_values)); + + for (size_t col_id = 2; col_id <= 14; ++col_id) + { + const auto name = fmt::format("c{}", col_id); + columns.push_back({name, TiDB::TP::TypeLongLong}); + + std::vector values; + values.reserve(rows); + for (size_t row = 0; row < rows; ++row) + values.push_back(static_cast(col_id * 1000 + row)); + data.emplace_back(toVec(name, values)); + } + + context.addMockDeltaMerge({"test_db", table_name}, columns, data, /*concurrency_hint=*/4); +} + +void addMockMultiStageLateMaterializationTopNOrderByTable( + MockDAGRequestContext & context, + const String & table_name, + TiDB::TP order_by_type, + ColumnWithTypeAndName order_by_column) +{ + addMockMultiStageLateMaterializationTopNOrderByTable( + context, + table_name, + MockColumnInfo{"ord", order_by_type}, + std::move(order_by_column)); +} + +ColumnWithTypeAndName toEnum16Vec( + const String & name, + const std::vector & values, + const DataTypeEnum16::Values & enum_values) +{ + auto type = std::make_shared(enum_values); + auto column = type->createColumn(); + for (const auto value : values) + column->insert(Field(static_cast(value))); + return {std::move(column), type, name}; +} + void rewriteComparisonRightLiteral( tipb::Expr * condition, tipb::ScalarFuncSig sig, @@ -453,6 +654,45 @@ void assertMultiStageRowsOverride(DAGContext & dag_context, UInt64 expected_sele ASSERT_NE(selection_rows, nullptr); ASSERT_EQ(selection_rows->load(), expected_selection_rows); ASSERT_GE(table_scan_rows->load(), selection_rows->load()); + + auto stats = getMultiStageRuntimeStats(dag_context); + ASSERT_NE(stats, nullptr); + ASSERT_EQ(selection_rows->load(), stats->final_rest_input_rows.load()); +} + +void assertTopNEnhancedMultiStageFallbackStats(DAGContext & dag_context, UInt64 expected_selection_rows) +{ + auto stats = getMultiStageRuntimeStats(dag_context); + ASSERT_NE(stats, nullptr); + ASSERT_FALSE(stats->topn_enabled.load()); + ASSERT_EQ(stats->running_topn_input_rows.load(), 0); + ASSERT_EQ(stats->running_topn_selected_rows.load(), 0); + ASSERT_EQ(stats->running_topn_bypass_rows.load(), 0); + ASSERT_EQ(stats->running_topn_filtered_rows.load(), 0); + ASSERT_EQ(stats->stage1_output_rows.load(), expected_selection_rows); + ASSERT_EQ(stats->final_rest_input_rows.load(), expected_selection_rows); +} + +void assertTopNEnhancedMultiStageFallsBackToSelection( + ExecutorsWithDMTestRunner & runner, + MockDAGRequestContext & context, + const std::shared_ptr & request, + const String & log_identifier, + UInt64 expected_selection_rows) +{ + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; + DAGContext disabled_dag_context(*request, log_identifier, /*concurrency=*/1); + auto expected = runner.executeStreams(&disabled_dag_context); + ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr); + ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr); + + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeTopN; + DAGContext enabled_dag_context(*request, log_identifier, /*concurrency=*/1); + auto actual = runner.executeStreams(&enabled_dag_context); + + assertColumnsEqual(expected, actual); + assertMultiStageRowsOverride(enabled_dag_context, expected_selection_rows); + assertTopNEnhancedMultiStageFallbackStats(enabled_dag_context, expected_selection_rows); } } // namespace @@ -535,13 +775,13 @@ try lt(col("c0"), lit(Field(static_cast(16)))), eq(makeASTFunction("bitand", col("c1"), lit(Field(static_cast(3)))), lit(Field(static_cast(0))))); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto expected = executeStreams(&disabled_dag_context); ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr); ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection; DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto actual = executeStreams(&enabled_dag_context); ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false)) @@ -565,13 +805,13 @@ try lt(col("c0"), lit(Field(static_cast(16)))), eq(makeASTFunction("bitand", col("c1"), lit(Field(static_cast(8)))), lit(Field(static_cast(0))))); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto expected = executeStreams(&disabled_dag_context); ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr); ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection; DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto actual = executeStreams(&enabled_dag_context); ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false)) @@ -583,6 +823,651 @@ try } CATCH +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNAsc) +try +{ + enablePipeline(true); + context.context->setSetting("max_block_size", Field(static_cast(64))); + + auto request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "c1", + /*is_desc=*/false, + /*limit=*/2); + + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; + DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1); + auto expected = executeStreams(&disabled_dag_context); + ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr); + ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr); + + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection; + DAGContext selection_only_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1); + auto selection_only_actual = executeStreams(&selection_only_dag_context); + ASSERT_TRUE(columnsEqual(expected, selection_only_actual, /*_restrict=*/false)) + << "\n expect_block: \n" + << getColumnsContent(expected) << "\n actual_block: \n" + << getColumnsContent(selection_only_actual); + assertMultiStageRowsOverride(selection_only_dag_context, 16); + + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeTopN; + DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1); + auto actual = executeStreams(&enabled_dag_context); + ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false)) + << "\n expect_block: \n" + << getColumnsContent(expected) << "\n actual_block: \n" + << getColumnsContent(actual); + + auto table_scan_rows = enabled_dag_context.getExecutorRowsOverride("table_scan_0"); + auto selection_rows = enabled_dag_context.getExecutorRowsOverride("selection_1"); + ASSERT_NE(table_scan_rows, nullptr); + ASSERT_NE(selection_rows, nullptr); + ASSERT_EQ(table_scan_rows->load(), 16); + ASSERT_GE(selection_rows->load(), 2); + ASSERT_LT(selection_rows->load(), 16); +} +CATCH + +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDesc) +try +{ + enablePipeline(true); + context.context->setSetting("max_block_size", Field(static_cast(64))); + + auto request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "c1", + /*is_desc=*/true, + /*limit=*/2); + + assertTopNEnhancedMultiStageMatchesDisabled( + *this, + context, + request, + dag_context_ptr->log->identifier(), + /*expected_table_scan_rows=*/16, + /*min_selection_rows=*/2); +} +CATCH + +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNMultiColumnOrderBy) +try +{ + enablePipeline(true); + context.context->setSetting("max_block_size", Field(static_cast(64))); + + auto request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + MockOrderByItemVec{{"c1", false}, {"c2", true}}, + /*limit=*/3); + + assertTopNEnhancedMultiStageMatchesDisabled( + *this, + context, + request, + dag_context_ptr->log->identifier(), + /*expected_table_scan_rows=*/16, + /*min_selection_rows=*/3); +} +CATCH + +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDateTime) +try +{ + enablePipeline(true); + context.context->setSetting("max_block_size", Field(static_cast(64))); + + MockColumnInfoVec columns; + columns.push_back({"c0", TiDB::TP::TypeLongLong}); + columns.push_back({"ts_col", TiDB::TP::TypeDatetime}); + columns.push_back({"c1", TiDB::TP::TypeLongLong}); + for (size_t col_id = 2; col_id <= 14; ++col_id) + columns.push_back({fmt::format("c{}", col_id), TiDB::TP::TypeLongLong}); + + constexpr size_t rows = 32; + std::vector c0_values; + ColumnWithNullableMyDateTime ts_values; + std::vector c1_values; + c0_values.reserve(rows); + ts_values.reserve(rows); + c1_values.reserve(rows); + for (size_t row = 0; row < rows; ++row) + { + c0_values.push_back(row); + ts_values.push_back(MyDateTime(2020, 1, 1, static_cast(row), 0, 0, 0).toPackedUInt()); + c1_values.push_back(row % 8); + } + + ColumnsWithTypeAndName data; + data.emplace_back(toVec("c0", c0_values)); + data.emplace_back(toNullableVec("ts_col", ts_values)); + data.emplace_back(toVec("c1", c1_values)); + for (size_t col_id = 2; col_id <= 14; ++col_id) + { + std::vector values; + values.reserve(rows); + for (size_t row = 0; row < rows; ++row) + values.push_back(static_cast(col_id * 1000 + row)); + data.emplace_back(toVec(fmt::format("c{}", col_id), values)); + } + context.addMockDeltaMerge( + {"test_db", "multi_stage_lm_datetime_topn"}, + columns, + data, + /*concurrency_hint=*/4); + + auto request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm_datetime_topn", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "ts_col", + /*is_desc=*/false, + /*limit=*/2); + + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; + DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1); + auto expected = executeStreams(&disabled_dag_context); + + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeTopN; + DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1); + auto actual = executeStreams(&enabled_dag_context); + ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false)) + << "\n expect_block: \n" + << getColumnsContent(expected) << "\n actual_block: \n" + << getColumnsContent(actual); + + auto table_scan_rows = enabled_dag_context.getExecutorRowsOverride("table_scan_0"); + auto selection_rows = enabled_dag_context.getExecutorRowsOverride("selection_1"); + ASSERT_NE(table_scan_rows, nullptr); + ASSERT_NE(selection_rows, nullptr); + ASSERT_EQ(table_scan_rows->load(), 16); + ASSERT_GE(selection_rows->load(), 2); + ASSERT_LT(selection_rows->load(), 16); +} +CATCH + +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNNullableOrderByNullOrdering) +try +{ + enablePipeline(true); + context.context->setSetting("max_block_size", Field(static_cast(64))); + + addMockMultiStageLateMaterializationTopNOrderByTable( + context, + "multi_stage_lm_nullable_topn", + TiDB::TP::TypeLongLong, + toNullableVec("ord", {{}, 100, {}, 90, 80, 70, 60, 50, 40, 30, 20, 10, 9, 8, 7, 6, + 5, 4, 3, 2, 1, 0, -1, -2, -3, -4, -5, -6, -7, -8, -9, -10})); + + for (const auto is_desc : {false, true}) + { + auto request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm_nullable_topn", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "ord", + is_desc, + /*limit=*/2); + + assertTopNEnhancedMultiStageMatchesDisabled( + *this, + context, + request, + dag_context_ptr->log->identifier(), + /*expected_table_scan_rows=*/16, + /*min_selection_rows=*/2); + } +} +CATCH + +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNFloatDesc) +try +{ + enablePipeline(true); + context.context->setSetting("max_block_size", Field(static_cast(64))); + + std::vector order_by_values; + order_by_values.reserve(32); + for (size_t row = 0; row < 32; ++row) + order_by_values.push_back(static_cast(row) + 0.5); + addMockMultiStageLateMaterializationTopNOrderByTable( + context, + "multi_stage_lm_float_topn", + TiDB::TP::TypeDouble, + toVec("ord", order_by_values)); + + auto request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm_float_topn", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "ord", + /*is_desc=*/true, + /*limit=*/3); + + assertTopNEnhancedMultiStageMatchesDisabled( + *this, + context, + request, + dag_context_ptr->log->identifier(), + /*expected_table_scan_rows=*/16, + /*min_selection_rows=*/3); +} +CATCH + +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDecimalDesc) +try +{ + enablePipeline(true); + context.context->setSetting("max_block_size", Field(static_cast(64))); + + std::vector> order_by_values; + order_by_values.reserve(32); + for (size_t row = 0; row < 32; ++row) + order_by_values.emplace_back(static_cast(row * 10), 0); + addMockMultiStageLateMaterializationTopNOrderByTable( + context, + "multi_stage_lm_decimal_topn", + TiDB::TP::TypeNewDecimal, + toVec("ord", order_by_values)); + + auto request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm_decimal_topn", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "ord", + /*is_desc=*/true, + /*limit=*/3); + + assertTopNEnhancedMultiStageMatchesDisabled( + *this, + context, + request, + dag_context_ptr->log->identifier(), + /*expected_table_scan_rows=*/16, + /*min_selection_rows=*/3); +} +CATCH + +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDateDesc) +try +{ + enablePipeline(true); + context.context->setSetting("max_block_size", Field(static_cast(64))); + + std::vector::FieldType> order_by_values; + order_by_values.reserve(32); + for (size_t row = 0; row < 32; ++row) + order_by_values.push_back(MyDate(2020, 1, static_cast(row % 28 + 1)).toPackedUInt()); + addMockMultiStageLateMaterializationTopNOrderByTable( + context, + "multi_stage_lm_date_topn", + TiDB::TP::TypeDate, + toVec("ord", order_by_values)); + + auto request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm_date_topn", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "ord", + /*is_desc=*/true, + /*limit=*/3); + + assertTopNEnhancedMultiStageMatchesDisabled( + *this, + context, + request, + dag_context_ptr->log->identifier(), + /*expected_table_scan_rows=*/16, + /*min_selection_rows=*/3); +} +CATCH + +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDateTimeDesc) +try +{ + enablePipeline(true); + context.context->setSetting("max_block_size", Field(static_cast(64))); + + std::vector::FieldType>> order_by_values; + order_by_values.reserve(32); + for (size_t row = 0; row < 32; ++row) + order_by_values.push_back(MyDateTime(2020, 1, 1, static_cast(row % 24), 0, 0, 0).toPackedUInt()); + addMockMultiStageLateMaterializationTopNOrderByTable( + context, + "multi_stage_lm_datetime_desc_topn", + TiDB::TP::TypeDatetime, + toNullableVec("ord", order_by_values)); + + auto request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm_datetime_desc_topn", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "ord", + /*is_desc=*/true, + /*limit=*/3); + + assertTopNEnhancedMultiStageMatchesDisabled( + *this, + context, + request, + dag_context_ptr->log->identifier(), + /*expected_table_scan_rows=*/16, + /*min_selection_rows=*/3); +} +CATCH + +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNTimestampDesc) +try +{ + enablePipeline(true); + context.context->setSetting("max_block_size", Field(static_cast(64))); + + std::vector::FieldType>> order_by_values; + order_by_values.reserve(32); + for (size_t row = 0; row < 32; ++row) + order_by_values.push_back(MyDateTime(2020, 1, 1, static_cast(row % 24), 0, 0, 0).toPackedUInt()); + addMockMultiStageLateMaterializationTopNOrderByTable( + context, + "multi_stage_lm_timestamp_topn", + TiDB::TP::TypeTimestamp, + toNullableVec("ord", order_by_values)); + + auto request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm_timestamp_topn", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "ord", + /*is_desc=*/true, + /*limit=*/3); + request->set_time_zone_name("Asia/Shanghai"); + + assertTopNEnhancedMultiStageFallsBackToSelection( + *this, + context, + request, + dag_context_ptr->log->identifier(), + /*expected_selection_rows=*/16); +} +CATCH + +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDisabledForPartitionBy) +try +{ + enablePipeline(true); + context.context->setSetting("max_block_size", Field(static_cast(64))); + + auto request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "c1", + /*is_desc=*/false, + /*limit=*/2); + + auto * topn = request->mutable_root_executor()->mutable_topn(); + ASSERT_EQ(topn->order_by_size(), 1); + *topn->add_partition_by() = topn->order_by(0); + + assertTopNEnhancedMultiStageFallsBackToSelection( + *this, + context, + request, + dag_context_ptr->log->identifier(), + /*expected_selection_rows=*/16); +} +CATCH + +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNVersionColumn) +try +{ + enablePipeline(true); + context.context->setSetting("max_block_size", Field(static_cast(64))); + + MockColumnInfoVec columns; + columns.push_back({"pk", TiDB::TP::TypeLongLong}); + columns.push_back({"commit_ts", TiDB::TP::TypeLongLong}); + for (size_t col_id = 1; col_id <= 14; ++col_id) + columns.push_back({fmt::format("c{}", col_id), TiDB::TP::TypeLongLong}); + + constexpr size_t rows = 32; + std::vector pk_values; + std::vector commit_ts_dummy_values; + pk_values.reserve(rows); + commit_ts_dummy_values.reserve(rows); + for (size_t row = 0; row < rows; ++row) + { + pk_values.push_back(row); + commit_ts_dummy_values.push_back(-1); + } + + ColumnsWithTypeAndName data; + data.emplace_back(toVec("pk", pk_values)); + data.emplace_back(toVec("commit_ts", commit_ts_dummy_values)); + for (size_t col_id = 1; col_id <= 14; ++col_id) + { + std::vector values; + values.reserve(rows); + for (size_t row = 0; row < rows; ++row) + { + if (col_id == 1) + values.push_back(row % 8); + else + values.push_back(static_cast(col_id * 1000 + row)); + } + data.emplace_back(toVec(fmt::format("c{}", col_id), values)); + } + context.addMockDeltaMerge( + {"test_db", "multi_stage_lm_version_topn"}, + columns, + data, + /*concurrency_hint=*/4); + + auto request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm_version_topn", + lt(col("pk"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "commit_ts", + /*is_desc=*/true, + /*limit=*/3); + rewriteTableScanColumnAsCommitTS(request.get(), /*column_index=*/1); + + const auto selected_rows = std::vector{15, 14, 13}; + auto expected_commit_ts_col = toNullableVec( + MutableSupport::version_column_name, + buildNullableInt64Values(selected_rows, [](Int64 row) { return row + 1; })); + auto expected = buildExpectedPkSpecialPayloadColumnsForDAG(selected_rows, expected_commit_ts_col); + + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; + DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1); + assertColumnsEqual(expected, executeStreams(&disabled_dag_context)); + ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr); + ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr); + + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeTopN; + DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1); + assertColumnsEqual(expected, executeStreams(&enabled_dag_context)); + assertMultiStageRowsOverride(enabled_dag_context, 16); +} +CATCH + +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDisabledForStringOrderBy) +try +{ + enablePipeline(true); + context.context->setSetting("max_block_size", Field(static_cast(64))); + + MockColumnInfoVec columns; + columns.push_back({"c0", TiDB::TP::TypeLongLong}); + columns.push_back({"c1", TiDB::TP::TypeLongLong}); + columns.push_back({"s", TiDB::TP::TypeString}); + for (size_t col_id = 2; col_id <= 14; ++col_id) + columns.push_back({fmt::format("c{}", col_id), TiDB::TP::TypeLongLong}); + + constexpr size_t rows = 32; + std::vector c0_values; + std::vector c1_values; + std::vector s_values; + c0_values.reserve(rows); + c1_values.reserve(rows); + s_values.reserve(rows); + for (size_t row = 0; row < rows; ++row) + { + c0_values.push_back(row); + c1_values.push_back(row % 8); + s_values.push_back(fmt::format("s{:02}", 31 - row)); + } + + ColumnsWithTypeAndName data; + data.emplace_back(toVec("c0", c0_values)); + data.emplace_back(toVec("c1", c1_values)); + data.emplace_back(toVec("s", s_values)); + for (size_t col_id = 2; col_id <= 14; ++col_id) + { + std::vector values; + values.reserve(rows); + for (size_t row = 0; row < rows; ++row) + values.push_back(static_cast(col_id * 1000 + row)); + data.emplace_back(toVec(fmt::format("c{}", col_id), values)); + } + context.addMockDeltaMerge( + {"test_db", "multi_stage_lm_string_topn"}, + columns, + data, + /*concurrency_hint=*/4); + + auto request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm_string_topn", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "s", + /*is_desc=*/false, + /*limit=*/2); + + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; + DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1); + auto expected = executeStreams(&disabled_dag_context); + + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeTopN; + DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/1); + auto actual = executeStreams(&enabled_dag_context); + ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false)) + << "\n expect_block: \n" + << getColumnsContent(expected) << "\n actual_block: \n" + << getColumnsContent(actual); + + assertMultiStageRowsOverride(enabled_dag_context, 16); + assertTopNEnhancedMultiStageFallbackStats(enabled_dag_context, 16); +} +CATCH + +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNDisabledForBitSetEnumOrderBy) +try +{ + enablePipeline(true); + context.context->setSetting("max_block_size", Field(static_cast(64))); + + constexpr size_t rows = 32; + std::vector uint64_order_by_values; + uint64_order_by_values.reserve(rows); + for (size_t row = 0; row < rows; ++row) + uint64_order_by_values.push_back(static_cast(rows - row)); + + addMockMultiStageLateMaterializationTopNOrderByTable( + context, + "multi_stage_lm_bit_topn", + TiDB::TP::TypeBit, + toVec("ord", uint64_order_by_values)); + auto bit_request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm_bit_topn", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "ord", + /*is_desc=*/false, + /*limit=*/2); + assertTopNEnhancedMultiStageFallsBackToSelection( + *this, + context, + bit_request, + dag_context_ptr->log->identifier(), + /*expected_selection_rows=*/16); + + MockColumnInfo set_info{"ord", TiDB::TP::TypeSet}; + for (Int16 value = 1; value <= 6; ++value) + set_info.elems.emplace_back(fmt::format("set_{}", value), value); + addMockMultiStageLateMaterializationTopNOrderByTable( + context, + "multi_stage_lm_set_topn", + set_info, + toVec("ord", uint64_order_by_values)); + auto set_request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm_set_topn", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "ord", + /*is_desc=*/false, + /*limit=*/2); + assertTopNEnhancedMultiStageFallsBackToSelection( + *this, + context, + set_request, + dag_context_ptr->log->identifier(), + /*expected_selection_rows=*/16); + + MockColumnInfo enum_info{"ord", TiDB::TP::TypeEnum}; + DataTypeEnum16::Values enum_values; + enum_values.reserve(rows); + std::vector enum_order_by_values; + enum_order_by_values.reserve(rows); + for (size_t row = 0; row < rows; ++row) + { + const auto value = static_cast(row + 1); + enum_info.elems.emplace_back(fmt::format("enum_{}", value), value); + enum_values.emplace_back(fmt::format("enum_{}", value), value); + enum_order_by_values.push_back(static_cast(rows - row)); + } + addMockMultiStageLateMaterializationTopNOrderByTable( + context, + "multi_stage_lm_enum_topn", + enum_info, + toEnum16Vec("ord", enum_order_by_values, enum_values)); + auto enum_request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm_enum_topn", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "ord", + /*is_desc=*/false, + /*limit=*/2); + assertTopNEnhancedMultiStageFallsBackToSelection( + *this, + context, + enum_request, + dag_context_ptr->log->identifier(), + /*expected_selection_rows=*/16); +} +CATCH + TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationEnabledAtRestColumnThreshold) try { @@ -595,13 +1480,13 @@ try lt(col("c0"), lit(Field(static_cast(16)))), buildAllGreaterThanMinusOneFilter(/*begin_col_id=*/1, /*end_col_id=*/6)); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto expected = executeStreams(&disabled_dag_context); ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr); ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection; DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto actual = executeStreams(&enabled_dag_context); ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false)) @@ -627,13 +1512,13 @@ try lt(col("c0"), lit(Field(static_cast(16)))), lt(col("c1"), lit(Field(static_cast(4))))); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto expected = executeStreams(&disabled_dag_context); ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr); ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection; DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto actual = executeStreams(&enabled_dag_context); ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false)) @@ -646,6 +1531,35 @@ try } CATCH +TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationTopNFallsBackWhenRestColumnsBelowThreshold) +try +{ + enablePipeline(true); + context.context->setSetting("max_block_size", Field(static_cast(64))); + + addMockMultiStageLateMaterializationTable( + context, + "multi_stage_lm_topn_rest_below_threshold", + /*column_count=*/11); + + auto request = buildDAGRequestWithPushedDownFilterAndTopN( + context, + "multi_stage_lm_topn_rest_below_threshold", + lt(col("c0"), lit(Field(static_cast(16)))), + gt(col("c1"), lit(Field(static_cast(-1)))), + "c2", + /*is_desc=*/false, + /*limit=*/2); + + assertTopNEnhancedMultiStageFallsBackToSelection( + *this, + context, + request, + dag_context_ptr->log->identifier(), + /*expected_selection_rows=*/16); +} +CATCH + TEST_F(ExecutorsWithDMTestRunner, MultiStageLateMaterializationDisabledWhenRestColumnsRatioBelowThreshold) try { @@ -663,13 +1577,13 @@ try lt(col("c0"), lit(Field(static_cast(16)))), buildAllGreaterThanMinusOneFilter(/*begin_col_id=*/1, /*end_col_id=*/9)); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto expected = executeStreams(&disabled_dag_context); ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("table_scan_0"), nullptr); ASSERT_EQ(disabled_dag_context.getExecutorRowsOverride("selection_1"), nullptr); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection; DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto actual = executeStreams(&enabled_dag_context); ASSERT_TRUE(columnsEqual(expected, actual, /*_restrict=*/false)) @@ -704,11 +1618,11 @@ try for (const auto & condition : pushed_down_filter_request->root_executor().selection().conditions()) *table_scan->add_pushed_down_filter_conditions() = condition; - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto expected = executeStreams(&disabled_dag_context); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection; DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto actual = executeStreams(&enabled_dag_context); @@ -778,7 +1692,7 @@ try *findMutableTableScan(request->mutable_root_executor())->add_pushed_down_filter_conditions() = pushed_down_filter_expr; - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); const auto selected_rows = selectRows(rows, [](size_t) { return true; }); auto expected = buildExpectedPkSpecialPayloadColumnsForDAG( @@ -787,7 +1701,7 @@ try "project_2_"); assertColumnsEqual(expected, executeStreams(&disabled_dag_context)); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection; DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto actual = executeStreams(&enabled_dag_context); @@ -858,7 +1772,7 @@ try })); auto expected = buildExpectedPkSpecialPayloadColumnsForDAG(selected_rows, expected_ts_col, "table_scan_0_"); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection; DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto actual = executeStreams(&enabled_dag_context); @@ -928,7 +1842,7 @@ try tipb::ScalarFuncSig::LTDuration, context.getCollation()); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); const auto selected_rows = selectRows(rows, [](size_t row) { return row % 8 > 2 && row % 8 < 6; }); auto expected = buildExpectedPkSpecialPayloadColumnsForDAG( @@ -937,7 +1851,7 @@ try "project_2_"); assertColumnsEqual(expected, executeStreams(&disabled_dag_context)); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection; DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto actual = executeStreams(&enabled_dag_context); @@ -1010,11 +1924,11 @@ try context.getCollation()); request->set_time_zone_name("Asia/Shanghai"); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto expected = executeStreams(&disabled_dag_context); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection; DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto actual = executeStreams(&enabled_dag_context); @@ -1039,13 +1953,13 @@ try gt(col("c0"), lit(Field(static_cast(16)))), lt(col("c0"), lit(Field(static_cast(24))))); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); const auto selected_rows = selectRows(32, [](size_t row) { return row > 16 && row < 24; }); auto expected = buildExpectedPayloadColumns(selected_rows); assertColumnsEqual(expected, executeStreams(&disabled_dag_context)); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection; DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto actual = executeStreams(&enabled_dag_context); @@ -1113,11 +2027,11 @@ try buildNullableInt64Values(selected_rows, [](Int64 row) { return row + 1; })); auto expected = buildExpectedPkSpecialPayloadColumnsForDAG(selected_rows, expected_commit_ts_col); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = false; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeDisabled; DAGContext disabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); assertColumnsEqual(expected, executeStreams(&disabled_dag_context)); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection; DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto actual = executeStreams(&enabled_dag_context); assertColumnsEqual(expected, actual); @@ -1144,7 +2058,7 @@ try gt(col("c1"), lit(Field(static_cast(-1))))); rewriteTableScanColumnAsExtraTableID(request.get(), /*column_index=*/1); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection; DAGContext enabled_dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); auto actual = executeStreams(&enabled_dag_context); @@ -1168,7 +2082,7 @@ try eq(makeASTFunction("bitand", col("c1"), lit(Field(static_cast(3)))), lit(Field(static_cast(0))))); request->set_collect_execution_summaries(true); - context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = true; + context.context->getSettingsRef().dt_enable_multi_stage_late_materialization = MSLMModeSelection; DAGContext dag_context(*request, dag_context_ptr->log->identifier(), /*concurrency=*/4); executeStreams(&dag_context); diff --git a/dbms/src/Interpreters/Settings.h b/dbms/src/Interpreters/Settings.h index 703dabdf632..583ef02d162 100644 --- a/dbms/src/Interpreters/Settings.h +++ b/dbms/src/Interpreters/Settings.h @@ -226,7 +226,7 @@ struct Settings M(SettingUInt64, dt_max_sharing_column_bytes_for_all, 2048 * Constant::MB, "Memory limitation for data sharing of all requests, include those sharing blocks in block queue. 0 means disable data sharing") \ M(SettingUInt64, dt_max_sharing_column_count, 5, "Deprecated") \ M(SettingBool, dt_enable_bitmap_filter, true, "Use bitmap filter to read data or not") \ - M(SettingBool, dt_enable_multi_stage_late_materialization, true, "Enable multi-stage late materialization for residual filters in DeltaTree Engine.") \ + M(SettingUInt64, dt_enable_multi_stage_late_materialization, 2, "Multi-stage late materialization mode in DeltaTree Engine. 0: disabled, 1: selection only, >=2: TopN enhanced.") \ M(SettingDouble, dt_read_thread_count_scale, 2.0, "Number of read thread = number of logical cpu cores * dt_read_thread_count_scale. Only has meaning at server startup.") \ M(SettingDouble, io_thread_count_scale, 5.0, "Number of thread of IOThreadPool = number of logical cpu cores * io_thread_count_scale. Only has meaning at server startup.") \ M(SettingUInt64, init_thread_count_scale, 100, "Number of thread = number of logical cpu cores * init_thread_count_scale. It just works for thread pool for initStores and loadMetadata. Only has meaning at server startup.") \ diff --git a/dbms/src/Storages/DeltaMerge/File/DMFilePackFilter.cpp b/dbms/src/Storages/DeltaMerge/File/DMFilePackFilter.cpp index 1ffa12e0cb6..45e41cebb1b 100644 --- a/dbms/src/Storages/DeltaMerge/File/DMFilePackFilter.cpp +++ b/dbms/src/Storages/DeltaMerge/File/DMFilePackFilter.cpp @@ -87,7 +87,7 @@ DMFilePackFilterResultPtr DMFilePackFilter::load(ReadTag read_tag) SCOPE_EXIT({ scan_context->total_rs_pack_filter_check_time_ns += watch.elapsed(); }); const size_t pack_count = dmfile->getPacks(); DMFilePackFilterResult result(index_cache, read_limiter, pack_count); - result.param.record_trim_metrics = read_tag == ReadTag::Query; + result.param.record_trim_metrics = isQueryRoughSetStatsOwner(read_tag); auto read_all_packs = (rowkey_ranges.size() == 1 && rowkey_ranges[0].all()) || rowkey_ranges.empty(); if (!read_all_packs) { @@ -185,11 +185,14 @@ DMFilePackFilterResultPtr DMFilePackFilter::load(ReadTag read_tag) // In table scanning, DMFilePackFilter of a DMFile may be created several times: // 1. When building MVCC bitmap (ReadTag::MVCC). // 2. When building LM filter stream (ReadTag::LMFilter). - // 3. When building MSLM stage1 residual filter stream (ReadTag::MSLMStage1Filter). - // 4. When building stream of other columns (ReadTag::Query). - // Only need to count the filter result once. + // 3. When building MSLM pushed filter stream (ReadTag::MSLMPushedFilter). + // 4. When building MSLM candidate stream (ReadTag::MSLMCandidate). + // 5. When building MSLM final rest stream (ReadTag::MSLMFinalRest). + // 6. When building stream of other columns (ReadTag::Query). + // Query owns these statistics for normal/LM reads. MSLMPushedFilter owns them for MSLM because it always + // traverses the logical MSLM input. // TODO: We can create DMFilePackFilter at the beginning and pass it to the stages described above. - if (read_tag == ReadTag::Query) + if (isQueryRoughSetStatsOwner(read_tag)) { scan_context->rs_pack_filter_none += none_count; scan_context->rs_pack_filter_some += some_count; diff --git a/dbms/src/Storages/DeltaMerge/File/DMFileReader.cpp b/dbms/src/Storages/DeltaMerge/File/DMFileReader.cpp index 3a3126a346b..f25515164ad 100644 --- a/dbms/src/Storages/DeltaMerge/File/DMFileReader.cpp +++ b/dbms/src/Storages/DeltaMerge/File/DMFileReader.cpp @@ -138,7 +138,8 @@ size_t DMFileReader::skipNextBlock() return 0; addSkippedRows(read_rows); - scan_context->late_materialization_skip_rows += read_rows; + if (read_tag == ReadTag::Query) + scan_context->late_materialization_skip_rows += read_rows; return read_rows; } @@ -162,7 +163,7 @@ std::pair DMFileReader::getReadRows() } next_row_offset += read_rows; - if (read_tag == ReadTag::Query && last_pack_res.allMatch()) + if (isQueryRoughSetStatsOwner(read_tag) && last_pack_res.allMatch()) scan_context->rs_dmfile_read_with_all += next_pack_id - start_pack_id; return {read_rows, last_pack_res}; } @@ -285,7 +286,7 @@ bool DMFileReader::isCacheableColumn(const ColumnDefine & cd) Block DMFileReader::read() { Stopwatch watch; - SCOPE_EXIT(scan_context->total_dmfile_read_time_ns += watch.elapsed();); + SCOPE_EXIT(scan_context->addDMFileReadTime(watch.elapsed(), read_tag);); /// 1. Skip filtered out packs. if (size_t skip_rows; !getSkippedRows(skip_rows)) @@ -694,41 +695,12 @@ bool DMFileReader::getCachedPacks( void DMFileReader::addScannedRows(UInt64 rows) { - switch (read_tag) - { - case ReadTag::Query: - scan_context->dmfile_data_scanned_rows += rows; - break; - case ReadTag::MVCC: - scan_context->dmfile_mvcc_scanned_rows += rows; - break; - case ReadTag::LMFilter: - scan_context->dmfile_lm_filter_scanned_rows += rows; - break; - case ReadTag::MSLMStage1Filter: - scan_context->dmfile_mslm_stage1_filter_scanned_rows += rows; - break; - default: - break; - } + scan_context->addDMFileScannedRows(rows, read_tag); } void DMFileReader::addSkippedRows(UInt64 rows) { - switch (read_tag) - { - case ReadTag::Query: - scan_context->dmfile_data_skipped_rows += rows; - break; - case ReadTag::MVCC: - scan_context->dmfile_mvcc_skipped_rows += rows; - break; - case ReadTag::LMFilter: - scan_context->dmfile_lm_filter_skipped_rows += rows; - break; - default: - break; - } + scan_context->addDMFileSkippedRows(rows, read_tag); } void DMFileReader::initAllMatchBlockInfo() diff --git a/dbms/src/Storages/DeltaMerge/Filter/PushDownFilter.h b/dbms/src/Storages/DeltaMerge/Filter/PushDownFilter.h index 538531afe0a..a813a47a08f 100644 --- a/dbms/src/Storages/DeltaMerge/Filter/PushDownFilter.h +++ b/dbms/src/Storages/DeltaMerge/Filter/PushDownFilter.h @@ -22,6 +22,7 @@ #include #include #include +#include namespace DB { @@ -45,7 +46,8 @@ class PushDownFilter const ColumnDefinesPtr & filter_columns_, const String filter_column_name_, const ExpressionActionsPtr & extra_cast_, - const ColumnDefinesPtr & columns_after_cast_) + const ColumnDefinesPtr & columns_after_cast_, + const MultiStageLateMaterializationTopNDescriptionPtr & topn_ = nullptr) : rs_operator(rs_operator_) , before_where(before_where_) , project_after_where(project_after_where_) @@ -53,10 +55,12 @@ class PushDownFilter , filter_columns(filter_columns_) , extra_cast(extra_cast_) , columns_after_cast(columns_after_cast_) + , topn(topn_) {} explicit PushDownFilter(const RSOperatorPtr & rs_operator_) : rs_operator(rs_operator_) + , topn(nullptr) {} Poco::JSON::Object::Ptr toJSONObject() const; @@ -94,6 +98,8 @@ class PushDownFilter const ExpressionActionsPtr extra_cast; // If the extra_cast is not null, the types of the columns may be changed const ColumnDefinesPtr columns_after_cast; + // Optional storage-side TopN metadata used only by multi-stage late materialization. + const MultiStageLateMaterializationTopNDescriptionPtr topn; }; } // namespace DB::DM diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp index cbb6002f0fd..16874b53cf1 100644 --- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp +++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.cpp @@ -17,10 +17,15 @@ #include #include +#include + namespace DB::DM { namespace { +constexpr UInt64 runtime_stats_flush_logical_rows = static_cast(DEFAULT_MERGE_BLOCK_SIZE) * 64; +constexpr UInt64 runtime_stats_flush_blocks = 64; + void filterBlock(Block & block, const IColumn::Filter & filter, size_t passed_count) { if (!block) @@ -57,11 +62,22 @@ MultiStageLateMaterializationBlockInputStream::MultiStageLateMaterializationBloc buildResidualFilterHeader(residual_filter_), residual_filter_->before_where, residual_filter_->filter_column_name) + , running_topn( + residual_filter_->topn != nullptr + ? std::make_unique(*residual_filter_->topn, *residual_filter_->filter_columns) + : nullptr) , log(Logger::get(NAME, req_id_)) { RUNTIME_CHECK(residual_filter != nullptr); RUNTIME_CHECK(residual_filter->before_where != nullptr); RUNTIME_CHECK(residual_filter->filter_columns != nullptr); + if (running_topn != nullptr && runtime_stats) + runtime_stats->recordRunningTopNEnabled(); +} + +MultiStageLateMaterializationBlockInputStream::~MultiStageLateMaterializationBlockInputStream() +{ + flushRuntimeStatsNoThrow(); } Block MultiStageLateMaterializationBlockInputStream::buildResidualFilterHeader( @@ -278,6 +294,53 @@ Block MultiStageLateMaterializationBlockInputStream::buildDirectModeBlock( return full_block; } +void MultiStageLateMaterializationBlockInputStream::updateTopNAdaptiveState( + UInt64 residual_passed_rows, + UInt64 topn_candidate_rows) +{ + if (running_topn == nullptr || topn_adaptive_disabled || residual_passed_rows == 0) + return; + + const auto warmup_rows + = std::max(running_topn->topK() * 4, multi_stage_late_materialization_topn_adaptive_rows_threshold); + if (!topn_adaptive_warmed_up) + { + topn_adaptive_warmup_observed_rows += residual_passed_rows; + if (topn_adaptive_warmup_observed_rows < warmup_rows) + return; + + topn_adaptive_warmed_up = true; + return; + } + + topn_adaptive_input_rows += residual_passed_rows; + topn_adaptive_candidate_rows += topn_candidate_rows; + + const auto min_input_rows_to_check + = std::max(running_topn->topK() * 4, multi_stage_late_materialization_topn_adaptive_rows_threshold); + if (topn_adaptive_input_rows < min_input_rows_to_check) + return; + + if (topn_adaptive_candidate_rows * multi_stage_late_materialization_topn_adaptive_disable_ratio_denominator + < topn_adaptive_input_rows * multi_stage_late_materialization_topn_adaptive_disable_ratio_numerator) + { + return; + } + + topn_adaptive_disabled = true; + LOG_INFO( + log, + "Disable running local TopN for multi-stage late materialization adaptively, warmup_rows={} input_rows={} " + "candidate_rows={} min_input_rows_to_check={} heap_size={} disable_ratio={}/{}", + topn_adaptive_warmup_observed_rows, + topn_adaptive_input_rows, + topn_adaptive_candidate_rows, + min_input_rows_to_check, + running_topn->heapSize(), + multi_stage_late_materialization_topn_adaptive_disable_ratio_numerator, + multi_stage_late_materialization_topn_adaptive_disable_ratio_denominator); +} + Block MultiStageLateMaterializationBlockInputStream::read() { while (true) @@ -287,19 +350,18 @@ Block MultiStageLateMaterializationBlockInputStream::read() stage0_block = stage0_filter_stream->read(stage0_filter, true); if (!stage0_block) { - logSummary(); + finishRuntimeStats(); return {}; } + ++runtime_stats_pending_blocks; auto effective_stage0_filter = buildStage0EffectiveFilter(stage0_block, stage0_filter); - if (runtime_stats) - runtime_stats->stage0_output_rows.fetch_add( - effective_stage0_filter.passed_count, - std::memory_order_relaxed); + local_runtime_stats.recordPushedFilter(stage0_block.rows(), effective_stage0_filter.passed_count); if (effective_stage0_filter.passed_count == 0) { skipNextBlockOrRead(stage1_filter_stream, "stage1_filter"); skipNextBlockOrRead(final_rest_stream, "final_rest"); + flushRuntimeStatsIfNeeded(); continue; } @@ -323,19 +385,64 @@ Block MultiStageLateMaterializationBlockInputStream::read() Block filter_eval_block; FilterPtr residual_filter_ptr = nullptr; const auto residual_passed_rows = executeResidualFilter(stage1_block, filter_eval_block, residual_filter_ptr); - if (runtime_stats) - runtime_stats->stage1_output_rows.fetch_add(residual_passed_rows, std::memory_order_relaxed); + local_runtime_stats.recordResidualFilter(stage1_block.rows(), residual_passed_rows); + + if (!shouldUseRunningTopN()) + { + if (running_topn != nullptr) + local_runtime_stats.recordRunningTopNBypass(residual_passed_rows); + else + local_runtime_stats.recordFinalRestInputRows(residual_passed_rows); + } if (residual_passed_rows == 0) { ++late_mode_blocks; skipNextBlockOrRead(final_rest_stream, "final_rest"); + flushRuntimeStatsIfNeeded(); continue; } + if (shouldUseRunningTopN()) + { + auto topn_result = running_topn->update(stage1_block, residual_filter_ptr, residual_passed_rows); + local_runtime_stats.recordRunningTopN(residual_passed_rows, topn_result.passed_count); + updateTopNAdaptiveState(residual_passed_rows, topn_result.passed_count); + + if (topn_result.passed_count == 0) + { + ++late_mode_blocks; + skipNextBlockOrRead(final_rest_stream, "final_rest"); + flushRuntimeStatsIfNeeded(); + continue; + } + + if (topn_result.passed_count == stage1_block.rows()) + { + ++direct_mode_blocks; + flushRuntimeStatsIfNeeded(); + return buildDirectModeBlock( + stage1_block, + stage0_filter_ptr, + effective_stage0_filter.passed_count, + nullptr, + topn_result.passed_count); + } + + ++late_mode_blocks; + flushRuntimeStatsIfNeeded(); + return buildLateModeBlock( + stage1_block, + stage0_filter_ptr, + stage0_block.rows(), + topn_result.filter, + topn_result.passed_count); + } + if (residual_passed_rows == stage1_block.rows()) { ++direct_mode_blocks; + flushRuntimeStatsIfNeeded(); return buildDirectModeBlock( stage1_block, stage0_filter_ptr, @@ -346,6 +453,7 @@ Block MultiStageLateMaterializationBlockInputStream::read() RUNTIME_CHECK(residual_filter_ptr != nullptr); ++late_mode_blocks; + flushRuntimeStatsIfNeeded(); return buildLateModeBlock( stage1_block, stage0_filter_ptr, @@ -355,17 +463,52 @@ Block MultiStageLateMaterializationBlockInputStream::read() } } -void MultiStageLateMaterializationBlockInputStream::logSummary() +void MultiStageLateMaterializationBlockInputStream::flushRuntimeStats() { - if (summary_logged) + if (local_runtime_stats.empty()) + { + runtime_stats_pending_blocks = 0; return; - summary_logged = true; + } + if (runtime_stats) + runtime_stats->merge(local_runtime_stats); + local_runtime_stats.reset(); + runtime_stats_pending_blocks = 0; +} - LOG_INFO( - log, - "Multi-stage late materialization finished, late_mode_blocks={} direct_mode_blocks={}", - late_mode_blocks, - direct_mode_blocks); +void MultiStageLateMaterializationBlockInputStream::flushRuntimeStatsNoThrow() +{ + try + { + finishRuntimeStats(); + } + catch (...) + {} +} + +void MultiStageLateMaterializationBlockInputStream::flushRuntimeStatsIfNeeded() +{ + if (local_runtime_stats.pushed_filter_input_rows >= runtime_stats_flush_logical_rows + || runtime_stats_pending_blocks >= runtime_stats_flush_blocks) + flushRuntimeStats(); +} + +void MultiStageLateMaterializationBlockInputStream::finishRuntimeStats() +{ + if (runtime_stats_finished) + return; + runtime_stats_finished = true; + + flushRuntimeStats(); + if (runtime_stats) + runtime_stats->finishStream( + late_mode_blocks, + direct_mode_blocks, + running_topn != nullptr ? running_topn->heapSize() : 0, + topn_adaptive_disabled, + topn_adaptive_warmup_observed_rows, + topn_adaptive_input_rows, + topn_adaptive_candidate_rows); } } // namespace DB::DM diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h index 90ac71bc11d..68f0a132c0f 100644 --- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h +++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationBlockInputStream.h @@ -40,6 +40,8 @@ class MultiStageLateMaterializationBlockInputStream : public IBlockInputStream const String & req_id_, const MultiStageLateMaterializationRuntimeStatsPtr & runtime_stats_ = nullptr); + ~MultiStageLateMaterializationBlockInputStream() override; + String getName() const override { return NAME; } Block getHeader() const override { return header; } @@ -90,7 +92,17 @@ class MultiStageLateMaterializationBlockInputStream : public IBlockInputStream const IColumn::Filter * residual_filter, size_t residual_passed_rows); - void logSummary(); + bool shouldUseRunningTopN() const { return running_topn != nullptr && !topn_adaptive_disabled; } + + void updateTopNAdaptiveState(UInt64 residual_passed_rows, UInt64 topn_candidate_rows); + + void flushRuntimeStats(); + + void flushRuntimeStatsNoThrow(); + + void flushRuntimeStatsIfNeeded(); + + void finishRuntimeStats(); private: Block header; @@ -100,11 +112,19 @@ class MultiStageLateMaterializationBlockInputStream : public IBlockInputStream PushDownFilterPtr residual_filter; BitmapFilterPtr bitmap_filter; MultiStageLateMaterializationRuntimeStatsPtr runtime_stats; + MultiStageLateMaterializationRuntimeStatsDelta local_runtime_stats; + UInt64 runtime_stats_pending_blocks = 0; FilterTransformAction residual_filter_action; + std::unique_ptr running_topn; size_t late_mode_blocks = 0; size_t direct_mode_blocks = 0; - bool summary_logged = false; + UInt64 topn_adaptive_warmup_observed_rows = 0; + UInt64 topn_adaptive_input_rows = 0; + UInt64 topn_adaptive_candidate_rows = 0; + bool topn_adaptive_warmed_up = false; + bool topn_adaptive_disabled = false; + bool runtime_stats_finished = false; const LoggerPtr log; }; diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h index 9bc7d221adc..4403e03c2f1 100644 --- a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h +++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationRuntimeStats.h @@ -14,6 +14,7 @@ #pragma once +#include #include #include @@ -22,10 +23,280 @@ namespace DB::DM { +struct MultiStageLateMaterializationRuntimeStatsDelta +{ + void reset() { *this = {}; } + + bool empty() const + { + return pushed_filter_input_rows == 0 && residual_filter_input_rows == 0 && final_rest_input_rows == 0 + && running_topn_input_rows == 0; + } + + void recordPushedFilter(UInt64 input_rows, UInt64 selected_rows) + { + pushed_filter_input_rows += input_rows; + pushed_filter_selected_rows += selected_rows; + pushed_filter_filtered_rows += input_rows - selected_rows; + + // Keep the old counter name for existing actRows/tests. + stage0_output_rows += selected_rows; + } + + void recordResidualFilter(UInt64 input_rows, UInt64 selected_rows) + { + residual_filter_input_rows += input_rows; + residual_filter_selected_rows += selected_rows; + residual_filter_filtered_rows += input_rows - selected_rows; + + // Keep the old counter name for existing actRows/tests. + stage1_output_rows += selected_rows; + } + + void recordFinalRestInputRows(UInt64 rows) + { + final_rest_input_rows += rows; + + // Keep the old counter name for existing actRows/tests. + topn_candidate_rows += rows; + } + + void recordRunningTopN(UInt64 input_rows, UInt64 selected_rows) + { + running_topn_input_rows += input_rows; + running_topn_selected_rows += selected_rows; + running_topn_filtered_rows += input_rows - selected_rows; + recordFinalRestInputRows(selected_rows); + } + + void recordRunningTopNBypass(UInt64 rows) + { + running_topn_input_rows += rows; + running_topn_bypass_rows += rows; + recordFinalRestInputRows(rows); + } + + UInt64 pushed_filter_input_rows = 0; + UInt64 pushed_filter_selected_rows = 0; + UInt64 pushed_filter_filtered_rows = 0; + UInt64 residual_filter_input_rows = 0; + UInt64 residual_filter_selected_rows = 0; + UInt64 residual_filter_filtered_rows = 0; + UInt64 final_rest_input_rows = 0; + UInt64 running_topn_input_rows = 0; + UInt64 running_topn_selected_rows = 0; + UInt64 running_topn_bypass_rows = 0; + UInt64 running_topn_filtered_rows = 0; + + // Legacy names kept because actRows overrides and existing unit tests still refer to them. + UInt64 stage0_output_rows = 0; + UInt64 stage1_output_rows = 0; + UInt64 topn_candidate_rows = 0; +}; + struct MultiStageLateMaterializationRuntimeStats { + MultiStageLateMaterializationRuntimeStats() = default; + + explicit MultiStageLateMaterializationRuntimeStats(const String & log_id_) + : log(Logger::get("MultiStageLateMaterialization", log_id_)) + {} + + ~MultiStageLateMaterializationRuntimeStats() + { + try + { + logSummary(); + } + catch (...) + {} + } + + void finishStream( + UInt64 stream_late_mode_blocks, + UInt64 stream_direct_mode_blocks, + UInt64 stream_topn_heap_size, + bool stream_topn_adaptive_disabled, + UInt64 stream_topn_adaptive_warmup_rows, + UInt64 stream_topn_adaptive_input_rows, + UInt64 stream_topn_adaptive_candidate_rows) + { + finished_streams.fetch_add(1, std::memory_order_relaxed); + late_mode_blocks.fetch_add(stream_late_mode_blocks, std::memory_order_relaxed); + direct_mode_blocks.fetch_add(stream_direct_mode_blocks, std::memory_order_relaxed); + topn_heap_size_sum.fetch_add(stream_topn_heap_size, std::memory_order_relaxed); + topn_adaptive_warmup_rows.fetch_add(stream_topn_adaptive_warmup_rows, std::memory_order_relaxed); + topn_adaptive_post_warmup_input_rows.fetch_add(stream_topn_adaptive_input_rows, std::memory_order_relaxed); + topn_adaptive_post_warmup_candidate_rows.fetch_add( + stream_topn_adaptive_candidate_rows, + std::memory_order_relaxed); + if (stream_topn_adaptive_disabled) + topn_adaptive_disabled_streams.fetch_add(1, std::memory_order_relaxed); + } + + void recordPushedFilter(UInt64 input_rows, UInt64 selected_rows) + { + pushed_filter_input_rows.fetch_add(input_rows, std::memory_order_relaxed); + pushed_filter_selected_rows.fetch_add(selected_rows, std::memory_order_relaxed); + pushed_filter_filtered_rows.fetch_add(input_rows - selected_rows, std::memory_order_relaxed); + + // Keep the old counter name for existing actRows/tests. + stage0_output_rows.fetch_add(selected_rows, std::memory_order_relaxed); + } + + void recordResidualFilter(UInt64 input_rows, UInt64 selected_rows) + { + residual_filter_input_rows.fetch_add(input_rows, std::memory_order_relaxed); + residual_filter_selected_rows.fetch_add(selected_rows, std::memory_order_relaxed); + residual_filter_filtered_rows.fetch_add(input_rows - selected_rows, std::memory_order_relaxed); + + // Keep the old counter name for existing actRows/tests. + stage1_output_rows.fetch_add(selected_rows, std::memory_order_relaxed); + } + + void recordFinalRestInputRows(UInt64 rows) + { + final_rest_input_rows.fetch_add(rows, std::memory_order_relaxed); + + // Keep the old counter name for existing actRows/tests. + topn_candidate_rows.fetch_add(rows, std::memory_order_relaxed); + } + + void recordRunningTopNEnabled() { topn_enabled.store(true, std::memory_order_relaxed); } + + void recordRunningTopN(UInt64 input_rows, UInt64 selected_rows) + { + recordRunningTopNEnabled(); + running_topn_input_rows.fetch_add(input_rows, std::memory_order_relaxed); + running_topn_selected_rows.fetch_add(selected_rows, std::memory_order_relaxed); + running_topn_filtered_rows.fetch_add(input_rows - selected_rows, std::memory_order_relaxed); + recordFinalRestInputRows(selected_rows); + } + + void recordRunningTopNBypass(UInt64 rows) + { + recordRunningTopNEnabled(); + running_topn_input_rows.fetch_add(rows, std::memory_order_relaxed); + running_topn_bypass_rows.fetch_add(rows, std::memory_order_relaxed); + recordFinalRestInputRows(rows); + } + + void merge(const MultiStageLateMaterializationRuntimeStatsDelta & delta) + { + if (delta.empty()) + return; + + pushed_filter_input_rows.fetch_add(delta.pushed_filter_input_rows, std::memory_order_relaxed); + pushed_filter_selected_rows.fetch_add(delta.pushed_filter_selected_rows, std::memory_order_relaxed); + pushed_filter_filtered_rows.fetch_add(delta.pushed_filter_filtered_rows, std::memory_order_relaxed); + residual_filter_input_rows.fetch_add(delta.residual_filter_input_rows, std::memory_order_relaxed); + residual_filter_selected_rows.fetch_add(delta.residual_filter_selected_rows, std::memory_order_relaxed); + residual_filter_filtered_rows.fetch_add(delta.residual_filter_filtered_rows, std::memory_order_relaxed); + final_rest_input_rows.fetch_add(delta.final_rest_input_rows, std::memory_order_relaxed); + running_topn_input_rows.fetch_add(delta.running_topn_input_rows, std::memory_order_relaxed); + running_topn_selected_rows.fetch_add(delta.running_topn_selected_rows, std::memory_order_relaxed); + running_topn_bypass_rows.fetch_add(delta.running_topn_bypass_rows, std::memory_order_relaxed); + running_topn_filtered_rows.fetch_add(delta.running_topn_filtered_rows, std::memory_order_relaxed); + + stage0_output_rows.fetch_add(delta.stage0_output_rows, std::memory_order_relaxed); + stage1_output_rows.fetch_add(delta.stage1_output_rows, std::memory_order_relaxed); + topn_candidate_rows.fetch_add(delta.topn_candidate_rows, std::memory_order_relaxed); + } + + void merge(const MultiStageLateMaterializationRuntimeStats & other) + { + pushed_filter_input_rows.fetch_add(other.pushed_filter_input_rows.load(std::memory_order_relaxed)); + pushed_filter_selected_rows.fetch_add(other.pushed_filter_selected_rows.load(std::memory_order_relaxed)); + pushed_filter_filtered_rows.fetch_add(other.pushed_filter_filtered_rows.load(std::memory_order_relaxed)); + residual_filter_input_rows.fetch_add(other.residual_filter_input_rows.load(std::memory_order_relaxed)); + residual_filter_selected_rows.fetch_add(other.residual_filter_selected_rows.load(std::memory_order_relaxed)); + residual_filter_filtered_rows.fetch_add(other.residual_filter_filtered_rows.load(std::memory_order_relaxed)); + final_rest_input_rows.fetch_add(other.final_rest_input_rows.load(std::memory_order_relaxed)); + running_topn_input_rows.fetch_add(other.running_topn_input_rows.load(std::memory_order_relaxed)); + running_topn_selected_rows.fetch_add(other.running_topn_selected_rows.load(std::memory_order_relaxed)); + running_topn_bypass_rows.fetch_add(other.running_topn_bypass_rows.load(std::memory_order_relaxed)); + running_topn_filtered_rows.fetch_add(other.running_topn_filtered_rows.load(std::memory_order_relaxed)); + + stage0_output_rows.fetch_add(other.stage0_output_rows.load(std::memory_order_relaxed)); + stage1_output_rows.fetch_add(other.stage1_output_rows.load(std::memory_order_relaxed)); + topn_candidate_rows.fetch_add(other.topn_candidate_rows.load(std::memory_order_relaxed)); + finished_streams.fetch_add(other.finished_streams.load(std::memory_order_relaxed)); + late_mode_blocks.fetch_add(other.late_mode_blocks.load(std::memory_order_relaxed)); + direct_mode_blocks.fetch_add(other.direct_mode_blocks.load(std::memory_order_relaxed)); + topn_heap_size_sum.fetch_add(other.topn_heap_size_sum.load(std::memory_order_relaxed)); + topn_adaptive_warmup_rows.fetch_add(other.topn_adaptive_warmup_rows.load(std::memory_order_relaxed)); + topn_adaptive_post_warmup_input_rows.fetch_add( + other.topn_adaptive_post_warmup_input_rows.load(std::memory_order_relaxed)); + topn_adaptive_post_warmup_candidate_rows.fetch_add( + other.topn_adaptive_post_warmup_candidate_rows.load(std::memory_order_relaxed)); + topn_adaptive_disabled_streams.fetch_add(other.topn_adaptive_disabled_streams.load(std::memory_order_relaxed)); + topn_enabled.store( + topn_enabled.load(std::memory_order_relaxed) || other.topn_enabled.load(std::memory_order_relaxed), + std::memory_order_relaxed); + } + + void logSummary() const + { + if (log == nullptr) + return; + + LOG_INFO( + log, + "Multi-stage late materialization finished, streams={} late_mode_blocks={} direct_mode_blocks={} " + "pushed_filter_input_rows={} pushed_filter_selected_rows={} pushed_filter_filtered_rows={} " + "residual_filter_input_rows={} residual_filter_selected_rows={} residual_filter_filtered_rows={} " + "final_rest_input_rows={} topn_enabled={} running_topn_input_rows={} running_topn_selected_rows={} " + "running_topn_bypass_rows={} running_topn_filtered_rows={} running_topn_heap_size_sum={} " + "topn_adaptive_warmup_rows={} topn_adaptive_post_warmup_input_rows={} " + "topn_adaptive_post_warmup_candidate_rows={} topn_adaptive_disabled_streams={}", + finished_streams.load(std::memory_order_relaxed), + late_mode_blocks.load(std::memory_order_relaxed), + direct_mode_blocks.load(std::memory_order_relaxed), + pushed_filter_input_rows.load(std::memory_order_relaxed), + pushed_filter_selected_rows.load(std::memory_order_relaxed), + pushed_filter_filtered_rows.load(std::memory_order_relaxed), + residual_filter_input_rows.load(std::memory_order_relaxed), + residual_filter_selected_rows.load(std::memory_order_relaxed), + residual_filter_filtered_rows.load(std::memory_order_relaxed), + final_rest_input_rows.load(std::memory_order_relaxed), + topn_enabled.load(std::memory_order_relaxed), + running_topn_input_rows.load(std::memory_order_relaxed), + running_topn_selected_rows.load(std::memory_order_relaxed), + running_topn_bypass_rows.load(std::memory_order_relaxed), + running_topn_filtered_rows.load(std::memory_order_relaxed), + topn_heap_size_sum.load(std::memory_order_relaxed), + topn_adaptive_warmup_rows.load(std::memory_order_relaxed), + topn_adaptive_post_warmup_input_rows.load(std::memory_order_relaxed), + topn_adaptive_post_warmup_candidate_rows.load(std::memory_order_relaxed), + topn_adaptive_disabled_streams.load(std::memory_order_relaxed)); + } + + std::atomic pushed_filter_input_rows{0}; + std::atomic pushed_filter_selected_rows{0}; + std::atomic pushed_filter_filtered_rows{0}; + std::atomic residual_filter_input_rows{0}; + std::atomic residual_filter_selected_rows{0}; + std::atomic residual_filter_filtered_rows{0}; + std::atomic final_rest_input_rows{0}; + std::atomic running_topn_input_rows{0}; + std::atomic running_topn_selected_rows{0}; + std::atomic running_topn_bypass_rows{0}; + std::atomic running_topn_filtered_rows{0}; + + // Legacy names kept because actRows overrides and existing unit tests still refer to them. std::atomic stage0_output_rows{0}; std::atomic stage1_output_rows{0}; + std::atomic topn_candidate_rows{0}; + std::atomic finished_streams{0}; + std::atomic late_mode_blocks{0}; + std::atomic direct_mode_blocks{0}; + std::atomic topn_heap_size_sum{0}; + std::atomic topn_adaptive_warmup_rows{0}; + std::atomic topn_adaptive_post_warmup_input_rows{0}; + std::atomic topn_adaptive_post_warmup_candidate_rows{0}; + std::atomic topn_adaptive_disabled_streams{0}; + std::atomic topn_enabled{false}; + LoggerPtr log; }; using MultiStageLateMaterializationRuntimeStatsPtr = std::shared_ptr; diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp new file mode 100644 index 00000000000..9b90cebf250 --- /dev/null +++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.cpp @@ -0,0 +1,441 @@ +// Copyright 2026 PingCAP, Inc. +// +// 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace DB::DM +{ +namespace +{ +struct SortKeyTypeInfo +{ + SortKeyKind kind; + SortKeyColumnView (*build_column_view)(const IColumn & column); + void (*extract_field)(const SortKeyColumnView & column, size_t row, SortKeyField & field); + int (*compare_column_with_owned)(const SortKeyColumnView & column, size_t row, const SortKeyField & field); +}; + +std::pair unwrapNullableColumn(const IColumn & column) +{ + if (const auto * nullable_column = typeid_cast(&column)) + return {&nullable_column->getNestedColumn(), &nullable_column->getNullMapData()}; + + return {&column, nullptr}; +} + +template +const ColumnType & castColumn(const IColumn & column) +{ + const auto * typed_column = typeid_cast(&column); + RUNTIME_CHECK_MSG( + typed_column != nullptr, + "Unexpected order by column type in TopN-enhanced multi-stage late materialization, column_type={}", + column.getName()); + return *typed_column; +} + +template +SortKeyColumnView buildSortKeyColumnViewImpl(const IColumn & column) +{ + const auto [nested_column, null_map] = unwrapNullableColumn(column); + return { + .data = &castColumn(*nested_column).getData(), + .null_map = null_map, + }; +} + +template +const typename ColumnType::Container & getColumnData(const SortKeyColumnView & column) +{ + return *static_cast(column.data); +} + +bool isNullAt(const IColumn::Filter * null_map, size_t row) +{ + return null_map != nullptr && (*null_map)[row] != 0; +} + +template +void extractFieldImpl(const SortKeyColumnView & column, size_t row, SortKeyField & field) +{ + field.kind = kind; + + field.is_null = isNullAt(column.null_map, row); + if (field.is_null) + return; + + const auto & data = getColumnData(column); + field.value = static_cast(data[row]); +} + +template +void extractDecimalFieldImpl(const SortKeyColumnView & column, size_t row, SortKeyField & field) +{ + field.kind = kind; + + field.is_null = isNullAt(column.null_map, row); + if (field.is_null) + return; + + const auto & data = getColumnData>(column); + field.value = data[row]; +} + +template +int compareValue(const T & lhs, const T & rhs) +{ + return lhs > rhs ? 1 : (lhs < rhs ? -1 : 0); +} + +int compareNull(bool lhs_is_null, bool rhs_is_null) +{ + if (lhs_is_null && rhs_is_null) + return 0; + return lhs_is_null ? -1 : 1; +} + +template +int compareColumnWithOwnedFieldImpl(const SortKeyColumnView & column, size_t row, const SortKeyField & rhs) +{ + const bool lhs_is_null = isNullAt(column.null_map, row); + if (lhs_is_null || rhs.is_null) + return compareNull(lhs_is_null, rhs.is_null); + + const auto & data = getColumnData(column); + return compareValue(static_cast(data[row]), std::get(rhs.value)); +} + +template +int compareFloatColumnWithOwnedFieldImpl(const SortKeyColumnView & column, size_t row, const SortKeyField & rhs) +{ + const bool lhs_is_null = isNullAt(column.null_map, row); + if (lhs_is_null || rhs.is_null) + return compareNull(lhs_is_null, rhs.is_null); + + const auto & data = getColumnData(column); + return CompareHelper::compare(data[row], std::get(rhs.value), -1); +} + +template +int compareDecimalColumnWithOwnedFieldImpl(const SortKeyColumnView & column, size_t row, const SortKeyField & rhs) +{ + const bool lhs_is_null = isNullAt(column.null_map, row); + if (lhs_is_null || rhs.is_null) + return compareNull(lhs_is_null, rhs.is_null); + + const auto & data = getColumnData>(column); + return compareValue(data[row], std::get(rhs.value)); +} + +template +SortKeyTypeInfo makeSortKeyTypeInfo() +{ + return { + .kind = kind, + .build_column_view = buildSortKeyColumnViewImpl, + .extract_field = extractFieldImpl, + .compare_column_with_owned = compareColumnWithOwnedFieldImpl, + }; +} + +template +SortKeyTypeInfo makeFloatSortKeyTypeInfo() +{ + return { + .kind = kind, + .build_column_view = buildSortKeyColumnViewImpl, + .extract_field = extractFieldImpl, + .compare_column_with_owned = compareFloatColumnWithOwnedFieldImpl, + }; +} + +template +SortKeyTypeInfo makeDecimalSortKeyTypeInfo() +{ + return { + .kind = kind, + .build_column_view = buildSortKeyColumnViewImpl>, + .extract_field = extractDecimalFieldImpl, + .compare_column_with_owned = compareDecimalColumnWithOwnedFieldImpl, + }; +} + +int compareSortKeyFields(const SortKeyField & lhs, const SortKeyField & rhs) +{ + RUNTIME_CHECK(lhs.kind == rhs.kind); + + if (lhs.is_null || rhs.is_null) + { + if (lhs.is_null && rhs.is_null) + return 0; + return lhs.is_null ? -1 : 1; + } + + switch (lhs.kind) + { + case SortKeyKind::Int64: + return compareValue(std::get(lhs.value), std::get(rhs.value)); + case SortKeyKind::UInt64: + case SortKeyKind::Date: + case SortKeyKind::DateTime: + return compareValue(std::get(lhs.value), std::get(rhs.value)); + case SortKeyKind::Float32: + return CompareHelper::compare(std::get(lhs.value), std::get(rhs.value), -1); + case SortKeyKind::Float64: + return CompareHelper::compare(std::get(lhs.value), std::get(rhs.value), -1); + case SortKeyKind::Decimal32: + return compareValue(std::get(lhs.value), std::get(rhs.value)); + case SortKeyKind::Decimal64: + return compareValue(std::get(lhs.value), std::get(rhs.value)); + case SortKeyKind::Decimal128: + return compareValue(std::get(lhs.value), std::get(rhs.value)); + case SortKeyKind::Decimal256: + return compareValue(std::get(lhs.value), std::get(rhs.value)); + } + RUNTIME_CHECK_MSG(false, "Unexpected sort key kind"); + return 0; +} + +SortKeyTypeInfo getSortKeyTypeInfo(const DataTypePtr & type) +{ + const auto type_not_null = removeNullable(type); + if (checkDataType(type_not_null.get())) + return makeSortKeyTypeInfo>(); + if (checkDataType(type_not_null.get())) + return makeSortKeyTypeInfo>(); + if (checkDataType(type_not_null.get())) + return makeSortKeyTypeInfo>(); + if (checkDataType(type_not_null.get())) + return makeSortKeyTypeInfo>(); + if (checkDataType(type_not_null.get())) + return makeSortKeyTypeInfo>(); + if (checkDataType(type_not_null.get())) + return makeSortKeyTypeInfo>(); + if (checkDataType(type_not_null.get())) + return makeSortKeyTypeInfo>(); + if (checkDataType(type_not_null.get())) + return makeSortKeyTypeInfo>(); + if (checkDataType(type_not_null.get())) + return makeFloatSortKeyTypeInfo>(); + if (checkDataType(type_not_null.get())) + return makeFloatSortKeyTypeInfo>(); + if (checkDataType(type_not_null.get())) + return makeDecimalSortKeyTypeInfo(); + if (checkDataType(type_not_null.get())) + return makeDecimalSortKeyTypeInfo(); + if (checkDataType(type_not_null.get())) + return makeDecimalSortKeyTypeInfo(); + if (checkDataType(type_not_null.get())) + return makeDecimalSortKeyTypeInfo(); + if (checkDataType(type_not_null.get())) + return makeSortKeyTypeInfo>(); + if (checkDataType(type_not_null.get())) + return makeSortKeyTypeInfo>(); + if (checkDataType(type_not_null.get())) + return makeSortKeyTypeInfo>(); + if (checkDataType(type_not_null.get())) + return makeSortKeyTypeInfo>(); + + throw Exception("Unsupported order by type for TopN-enhanced multi-stage late materialization: " + type->getName()); +} +} // namespace + +bool RunningLocalTopN::HeapComparator::operator()(const HeapEntry & lhs, const HeapEntry & rhs) const +{ + RUNTIME_CHECK(owner != nullptr); + return owner->compareOwnedKeys(lhs.key, rhs.key) < 0; +} + +RunningLocalTopN::RunningLocalTopN( + const MultiStageLateMaterializationTopNDescription & description, + const ColumnDefines & stage1_columns) + : topk(description.topk) + , heap(HeapComparator{this}) +{ + RUNTIME_CHECK(topk > 0); + RUNTIME_CHECK(description.order_by_columns.size() <= multi_stage_late_materialization_topn_max_order_by_columns); + + sort_key_columns.reserve(description.order_by_columns.size()); + for (const auto & order_by_column : description.order_by_columns) + { + const auto column_it = std::find_if(stage1_columns.begin(), stage1_columns.end(), [&](const auto & column) { + return column.id == order_by_column.column_id; + }); + RUNTIME_CHECK_MSG( + column_it != stage1_columns.end(), + "Order by column is not found in stage1 columns, column_id={}", + order_by_column.column_id); + + auto type_info = getSortKeyTypeInfo(column_it->type); + sort_key_columns.push_back(SortKeyColumnDesc{ + .column_id = order_by_column.column_id, + .column_pos = static_cast(column_it - stage1_columns.begin()), + .kind = type_info.kind, + .direction = order_by_column.direction, + .build_column_view = type_info.build_column_view, + .extract_field = type_info.extract_field, + .compare_column_with_owned = type_info.compare_column_with_owned, + }); + } +} + +int RunningLocalTopN::compareOwnedKeys(const OwnedSortKey & lhs, const OwnedSortKey & rhs) const +{ + RUNTIME_CHECK(lhs.size == rhs.size); + RUNTIME_CHECK(lhs.size == sort_key_columns.size()); + for (size_t i = 0; i < lhs.size; ++i) + { + if (const auto cmp = compareSortKeyFields(lhs.fields[i], rhs.fields[i]); cmp != 0) + return cmp * sort_key_columns[i].direction; + } + return 0; +} + +int RunningLocalTopN::compareRowWithOwnedKey( + const std::vector & sort_columns, + size_t row, + const OwnedSortKey & rhs) const +{ + RUNTIME_CHECK(sort_columns.size() == sort_key_columns.size()); + RUNTIME_CHECK(rhs.size == sort_key_columns.size()); + + for (size_t i = 0; i < sort_key_columns.size(); ++i) + { + const auto cmp = sort_key_columns[i].compare_column_with_owned(sort_columns[i], row, rhs.fields[i]); + if (cmp != 0) + return cmp * sort_key_columns[i].direction; + } + return 0; +} + +OwnedSortKey RunningLocalTopN::materializeOwnedKey(const std::vector & sort_columns, size_t row) + const +{ + RUNTIME_CHECK(sort_columns.size() == sort_key_columns.size()); + + OwnedSortKey key; + key.size = sort_key_columns.size(); + for (size_t i = 0; i < sort_key_columns.size(); ++i) + sort_key_columns[i].extract_field(sort_columns[i], row, key.fields[i]); + return key; +} + +RunningLocalTopNUpdateResult RunningLocalTopN::update( + const Block & stage1_block, + const IColumn::Filter * residual_filter, + size_t residual_passed_rows) +{ + ++current_block_sequence; + + const auto rows = stage1_block.rows(); + RunningLocalTopNUpdateResult result; + + if (residual_passed_rows == 0) + return result; + + std::vector materialized_const_columns; + std::vector sort_columns; + sort_columns.reserve(sort_key_columns.size()); + for (const auto & desc : sort_key_columns) + { + auto column = stage1_block.getByPosition(desc.column_pos).column; + if (auto full_column = column->convertToFullColumnIfConst()) + { + materialized_const_columns.push_back(std::move(full_column)); + column = materialized_const_columns.back(); + } + sort_columns.push_back(desc.build_column_view(*column)); + } + + auto ensure_filter = [&]() -> IColumn::Filter & { + if (result.filter.empty()) + result.filter.resize_fill(rows, 0); + return result.filter; + }; + auto mark_candidate = [&](size_t row) { + auto & filter = ensure_filter(); + if (filter[row] == 0) + { + filter[row] = 1; + ++result.passed_count; + } + }; + auto unmark_current_block_candidate = [&](const HeapEntry & entry) { + if (entry.block_sequence != current_block_sequence) + return; + + RUNTIME_CHECK(!result.filter.empty()); + RUNTIME_CHECK(entry.row_index_in_stage1_block < result.filter.size()); + if (result.filter[entry.row_index_in_stage1_block] != 0) + { + result.filter[entry.row_index_in_stage1_block] = 0; + --result.passed_count; + } + }; + + for (size_t row = 0; row < rows; ++row) + { + if (residual_filter != nullptr && (*residual_filter)[row] == 0) + continue; + + if (heap.size() < topk) + { + mark_candidate(row); + heap.push(HeapEntry{ + .key = materializeOwnedKey(sort_columns, row), + .block_sequence = current_block_sequence, + .row_index_in_stage1_block = static_cast(row), + }); + continue; + } + + const auto cmp = compareRowWithOwnedKey(sort_columns, row, heap.top().key); + // Equal sort keys are not kept as extra candidates. This matches the existing local TopN behavior, + // where ties are broken locally by the entries already selected by the partial TopN. + if (cmp < 0) + { + unmark_current_block_candidate(heap.top()); + heap.pop(); + + mark_candidate(row); + heap.push(HeapEntry{ + .key = materializeOwnedKey(sort_columns, row), + .block_sequence = current_block_sequence, + .row_index_in_stage1_block = static_cast(row), + }); + } + } + + return result; +} + +} // namespace DB::DM diff --git a/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.h b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.h new file mode 100644 index 00000000000..7d4a2d97b4f --- /dev/null +++ b/dbms/src/Storages/DeltaMerge/MultiStageLateMaterializationTopN.h @@ -0,0 +1,156 @@ +// Copyright 2026 PingCAP, Inc. +// +// 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 +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace DB::DM +{ + +inline constexpr UInt64 multi_stage_late_materialization_topn_max_topk = 2048; +inline constexpr size_t multi_stage_late_materialization_topn_max_order_by_columns = 4; +inline constexpr UInt64 multi_stage_late_materialization_topn_adaptive_rows_threshold = 8192; +inline constexpr UInt64 multi_stage_late_materialization_topn_adaptive_disable_ratio_numerator = 7; +inline constexpr UInt64 multi_stage_late_materialization_topn_adaptive_disable_ratio_denominator = 10; + +struct MultiStageLateMaterializationTopNOrderByColumn +{ + ColumnID column_id; + int direction; +}; + +struct MultiStageLateMaterializationTopNDescription +{ + UInt64 topk = 0; + std::vector order_by_columns; +}; + +using MultiStageLateMaterializationTopNDescriptionPtr = std::shared_ptr; + +enum class SortKeyKind +{ + Int64, + UInt64, + Float32, + Float64, + Decimal32, + Decimal64, + Decimal128, + Decimal256, + Date, + DateTime, +}; + +using SortKeyValue = std::variant; + +struct SortKeyField +{ + bool is_null = false; + SortKeyKind kind = SortKeyKind::Int64; + SortKeyValue value = Int64{}; +}; + +struct OwnedSortKey +{ + std::array fields; + size_t size = 0; +}; + +struct RunningLocalTopNUpdateResult +{ + IColumn::Filter filter; + size_t passed_count = 0; +}; + +struct SortKeyColumnView +{ + const void * data = nullptr; + const IColumn::Filter * null_map = nullptr; +}; + +class RunningLocalTopN +{ +public: + RunningLocalTopN( + const MultiStageLateMaterializationTopNDescription & description, + const ColumnDefines & stage1_columns); + + RunningLocalTopNUpdateResult update( + const Block & stage1_block, + const IColumn::Filter * residual_filter, + size_t residual_passed_rows); + + size_t heapSize() const { return heap.size(); } + UInt64 topK() const { return topk; } + +private: + struct SortKeyColumnDesc + { + using BuildColumnViewFn = SortKeyColumnView (*)(const IColumn & column); + using ExtractFieldFn = void (*)(const SortKeyColumnView & column, size_t row, SortKeyField & field); + using CompareColumnWithOwnedFn + = int (*)(const SortKeyColumnView & column, size_t row, const SortKeyField & field); + + ColumnID column_id = EmptyColumnID; + size_t column_pos = 0; + SortKeyKind kind = SortKeyKind::Int64; + int direction = 1; + BuildColumnViewFn build_column_view = nullptr; + ExtractFieldFn extract_field = nullptr; + CompareColumnWithOwnedFn compare_column_with_owned = nullptr; + }; + + struct HeapEntry + { + OwnedSortKey key; + UInt64 block_sequence = 0; + UInt32 row_index_in_stage1_block = 0; + }; + + struct HeapComparator + { + const RunningLocalTopN * owner = nullptr; + + bool operator()(const HeapEntry & lhs, const HeapEntry & rhs) const; + }; + + using Heap = std::priority_queue, HeapComparator>; + + int compareOwnedKeys(const OwnedSortKey & lhs, const OwnedSortKey & rhs) const; + int compareRowWithOwnedKey( + const std::vector & sort_columns, + size_t row, + const OwnedSortKey & rhs) const; + OwnedSortKey materializeOwnedKey(const std::vector & sort_columns, size_t row) const; + +private: + UInt64 topk; + std::vector sort_key_columns; + UInt64 current_block_sequence = 0; + Heap heap; +}; + +} // namespace DB::DM diff --git a/dbms/src/Storages/DeltaMerge/ReadMode.h b/dbms/src/Storages/DeltaMerge/ReadMode.h index 607e0788a62..aa3ce574012 100644 --- a/dbms/src/Storages/DeltaMerge/ReadMode.h +++ b/dbms/src/Storages/DeltaMerge/ReadMode.h @@ -45,9 +45,18 @@ enum class ReadTag Query, // Read columns required by queries. MVCC, // Read columns to build MVCC bitmap. LMFilter, // Read columns required by late-materialization filter. - MSLMStage1Filter, // Read columns required by multi-stage late-materialization residual filter. + MSLMPushedFilter, // Read columns required by multi-stage late-materialization pushed filter. + MSLMCandidate, // Read columns required by multi-stage late-materialization residual/order-by candidate read. + MSLMFinalRest, // Read final rest columns for multi-stage late-materialization. }; +/// A DMFile may be opened once for every physical MSLM read stage, but query-level +/// rough-set statistics must be recorded by exactly one stage. +inline constexpr bool isQueryRoughSetStatsOwner(ReadTag read_tag) +{ + return read_tag == ReadTag::Query || read_tag == ReadTag::MSLMPushedFilter; +} + enum class ReadRUType { MVCC_ESTIMATE, diff --git a/dbms/src/Storages/DeltaMerge/ScanContext.cpp b/dbms/src/Storages/DeltaMerge/ScanContext.cpp index 5e7e4366fa3..97bd28eb752 100644 --- a/dbms/src/Storages/DeltaMerge/ScanContext.cpp +++ b/dbms/src/Storages/DeltaMerge/ScanContext.cpp @@ -25,6 +25,73 @@ namespace DB::DM { +namespace +{ +constexpr uint64_t NS_TO_MS = 1'000'000; + +bool isQueryReadTag(ReadTag read_tag) +{ + switch (read_tag) + { + case ReadTag::Query: + case ReadTag::LMFilter: + case ReadTag::MSLMPushedFilter: + case ReadTag::MSLMCandidate: + case ReadTag::MSLMFinalRest: + return true; + default: + return false; + } +} + +void serializeReadStage(const MSLMReadStageScanContext & read_stage, tipb::TiFlashMSLMReadStageContext * proto) +{ + proto->set_dmfile_scanned_rows(read_stage.dmfile_scanned_rows); + proto->set_dmfile_skipped_rows(read_stage.dmfile_skipped_rows); + proto->set_read_bytes(read_stage.read_bytes); + proto->set_read_time_ms(read_stage.read_time_ns / NS_TO_MS); +} + +void mergeReadStage(MSLMReadStageScanContext & read_stage, const tipb::TiFlashMSLMReadStageContext & proto) +{ + read_stage.dmfile_scanned_rows += proto.dmfile_scanned_rows(); + read_stage.dmfile_skipped_rows += proto.dmfile_skipped_rows(); + read_stage.read_bytes += proto.read_bytes(); + read_stage.read_time_ns += proto.read_time_ms() * NS_TO_MS; +} + +void serializeFilter( + uint64_t input_rows, + uint64_t selected_rows, + uint64_t filtered_rows, + tipb::TiFlashMSLMFilterContext * proto) +{ + proto->set_input_rows(input_rows); + proto->set_selected_rows(selected_rows); + proto->set_filtered_rows(filtered_rows); +} + +Poco::JSON::Object::Ptr readStageToJson(const MSLMReadStageScanContext & read_stage) +{ + static constexpr double NS_TO_MS_SCALE = 1'000'000.0; + Poco::JSON::Object::Ptr json = new Poco::JSON::Object(); + json->set("dmfile_scanned_rows", read_stage.dmfile_scanned_rows.load()); + json->set("dmfile_skipped_rows", read_stage.dmfile_skipped_rows.load()); + json->set("read_bytes", read_stage.read_bytes.load()); + json->set("read_time", fmt::format("{:.3f}ms", read_stage.read_time_ns.load() / NS_TO_MS_SCALE)); + return json; +} + +Poco::JSON::Object::Ptr filterToJson(uint64_t input_rows, uint64_t selected_rows, uint64_t filtered_rows) +{ + Poco::JSON::Object::Ptr json = new Poco::JSON::Object(); + json->set("input_rows", input_rows); + json->set("selected_rows", selected_rows); + json->set("filtered_rows", filtered_rows); + return json; +} +} // namespace + void ScanContext::setRegionNumOfCurrentInstance(uint64_t region_num) { region_num_of_instance[current_instance_id] = region_num; @@ -103,7 +170,6 @@ String ScanContext::toJson() const json->set("dmfile_mvcc_skipped_rows", dmfile_mvcc_skipped_rows.load()); json->set("dmfile_lm_filter_scanned_rows", dmfile_lm_filter_scanned_rows.load()); json->set("dmfile_lm_filter_skipped_rows", dmfile_lm_filter_skipped_rows.load()); - json->set("dmfile_mslm_stage1_filter_scanned_rows", dmfile_mslm_stage1_filter_scanned_rows.load()); json->set("dmfile_read_time", fmt::format("{:.3f}ms", total_dmfile_read_time_ns.load() / NS_TO_MS_SCALE)); json->set( @@ -191,6 +257,48 @@ String ScanContext::toJson() const json->set("pushdown", pushdown_executor->toJSONObject()); } + if (hasMultiStageLateMaterializationContext()) + { + const auto & stats = getMultiStageLateMaterializationRuntimeStats(); + Poco::JSON::Object::Ptr mslm = new Poco::JSON::Object(); + mslm->set("streams", stats.finished_streams.load()); + mslm->set("late_mode_blocks", stats.late_mode_blocks.load()); + mslm->set("direct_mode_blocks", stats.direct_mode_blocks.load()); + mslm->set("pushed_filter_read", readStageToJson(mslm_pushed_filter_read)); + mslm->set("candidate_read", readStageToJson(mslm_candidate_read)); + mslm->set("final_rest_read", readStageToJson(mslm_final_rest_read)); + mslm->set( + "pushed_filter", + filterToJson( + stats.pushed_filter_input_rows.load(), + stats.pushed_filter_selected_rows.load(), + stats.pushed_filter_filtered_rows.load())); + mslm->set( + "residual_filter", + filterToJson( + stats.residual_filter_input_rows.load(), + stats.residual_filter_selected_rows.load(), + stats.residual_filter_filtered_rows.load())); + if (stats.topn_enabled.load()) + { + Poco::JSON::Object::Ptr running_topn = new Poco::JSON::Object(); + running_topn->set("input_rows", stats.running_topn_input_rows.load()); + running_topn->set("selected_rows", stats.running_topn_selected_rows.load()); + running_topn->set("bypass_rows", stats.running_topn_bypass_rows.load()); + running_topn->set("filtered_rows", stats.running_topn_filtered_rows.load()); + running_topn->set("heap_size_sum", stats.topn_heap_size_sum.load()); + running_topn->set("adaptive_warmup_rows", stats.topn_adaptive_warmup_rows.load()); + running_topn->set("adaptive_post_warmup_input_rows", stats.topn_adaptive_post_warmup_input_rows.load()); + running_topn->set( + "adaptive_post_warmup_candidate_rows", + stats.topn_adaptive_post_warmup_candidate_rows.load()); + running_topn->set("adaptive_disabled_streams", stats.topn_adaptive_disabled_streams.load()); + mslm->set("running_topn", running_topn); + } + mslm->set("final_rest_input_rows", stats.final_rest_input_rows.load()); + json->set("multi_stage_late_materialization", mslm); + } + std::stringstream buf; json->stringify(buf); return buf.str(); @@ -275,7 +383,7 @@ std::optional ScanContext::newLACBytesCollector(ReadTag read_ { if (resource_group_name.empty()) return std::nullopt; - if (read_tag != ReadTag::Query && read_tag != ReadTag::LMFilter && read_tag != ReadTag::MSLMStage1Filter) + if (!isQueryReadTag(read_tag)) return std::nullopt; return LACBytesCollector(resource_group_name); } @@ -285,8 +393,7 @@ void ScanContext::addUserReadBytes( ReadTag read_tag, std::optional & lac_bytes_collector) { - if (read_tag != ReadTag::Query && read_tag != ReadTag::LMFilter && read_tag != ReadTag::MSLMStage1Filter - && read_tag != ReadTag::MVCC) + if (!isQueryReadTag(read_tag) && read_tag != ReadTag::MVCC) return; if (read_tag == ReadTag::MVCC) { @@ -297,6 +404,8 @@ void ScanContext::addUserReadBytes( else { query_read_bytes += bytes; + if (auto * read_stage = getMutableMSLMReadStage(read_tag); read_stage != nullptr) + read_stage->read_bytes += bytes; if (query_read_bytes_counter) query_read_bytes_counter->Increment(bytes); if (lac_bytes_collector) @@ -304,4 +413,230 @@ void ScanContext::addUserReadBytes( } } +bool ScanContext::hasMultiStageLateMaterializationContext() const +{ + return multi_stage_late_materialization_enabled.load(std::memory_order_relaxed) + || multi_stage_late_materialization_runtime_stats != nullptr; +} + +const MultiStageLateMaterializationRuntimeStats & ScanContext::getMultiStageLateMaterializationRuntimeStats() const +{ + if (multi_stage_late_materialization_runtime_stats) + return *multi_stage_late_materialization_runtime_stats; + return merged_multi_stage_late_materialization_runtime_stats; +} + +MultiStageLateMaterializationRuntimeStats & ScanContext::getMergedMultiStageLateMaterializationRuntimeStats() +{ + multi_stage_late_materialization_enabled.store(true, std::memory_order_relaxed); + return merged_multi_stage_late_materialization_runtime_stats; +} + +MSLMReadStageScanContext * ScanContext::getMutableMSLMReadStage(ReadTag read_tag) +{ + switch (read_tag) + { + case ReadTag::MSLMPushedFilter: + multi_stage_late_materialization_enabled.store(true, std::memory_order_relaxed); + return &mslm_pushed_filter_read; + case ReadTag::MSLMCandidate: + multi_stage_late_materialization_enabled.store(true, std::memory_order_relaxed); + return &mslm_candidate_read; + case ReadTag::MSLMFinalRest: + multi_stage_late_materialization_enabled.store(true, std::memory_order_relaxed); + return &mslm_final_rest_read; + default: + return nullptr; + } +} + +const MSLMReadStageScanContext * ScanContext::getMSLMReadStage(ReadTag read_tag) const +{ + switch (read_tag) + { + case ReadTag::MSLMPushedFilter: + return &mslm_pushed_filter_read; + case ReadTag::MSLMCandidate: + return &mslm_candidate_read; + case ReadTag::MSLMFinalRest: + return &mslm_final_rest_read; + default: + return nullptr; + } +} + +void ScanContext::addDMFileReadTime(uint64_t ns, ReadTag read_tag) +{ + total_dmfile_read_time_ns += ns; + if (auto * read_stage = getMutableMSLMReadStage(read_tag); read_stage != nullptr) + read_stage->read_time_ns += ns; +} + +void ScanContext::addDMFileScannedRows(uint64_t rows, ReadTag read_tag) +{ + switch (read_tag) + { + case ReadTag::Query: + dmfile_data_scanned_rows += rows; + break; + case ReadTag::MVCC: + dmfile_mvcc_scanned_rows += rows; + break; + case ReadTag::LMFilter: + dmfile_lm_filter_scanned_rows += rows; + break; + case ReadTag::MSLMPushedFilter: + case ReadTag::MSLMCandidate: + case ReadTag::MSLMFinalRest: + getMutableMSLMReadStage(read_tag)->dmfile_scanned_rows += rows; + break; + default: + break; + } +} + +void ScanContext::addDMFileSkippedRows(uint64_t rows, ReadTag read_tag) +{ + switch (read_tag) + { + case ReadTag::Query: + dmfile_data_skipped_rows += rows; + break; + case ReadTag::MVCC: + dmfile_mvcc_skipped_rows += rows; + break; + case ReadTag::LMFilter: + dmfile_lm_filter_skipped_rows += rows; + break; + case ReadTag::MSLMPushedFilter: + case ReadTag::MSLMCandidate: + case ReadTag::MSLMFinalRest: + getMutableMSLMReadStage(read_tag)->dmfile_skipped_rows += rows; + break; + default: + break; + } +} + +void ScanContext::setMultiStageLateMaterializationRuntimeStats( + const MultiStageLateMaterializationRuntimeStatsPtr & stats) +{ + multi_stage_late_materialization_runtime_stats = stats; + if (stats != nullptr) + multi_stage_late_materialization_enabled.store(true, std::memory_order_relaxed); +} + +void ScanContext::serializeMultiStageLateMaterialization(tipb::TiFlashScanContext & proto) const +{ + if (!hasMultiStageLateMaterializationContext()) + return; + + const auto & stats = getMultiStageLateMaterializationRuntimeStats(); + auto * mslm = proto.mutable_multi_stage_late_materialization(); + mslm->set_streams(stats.finished_streams.load(std::memory_order_relaxed)); + mslm->set_late_mode_blocks(stats.late_mode_blocks.load(std::memory_order_relaxed)); + mslm->set_direct_mode_blocks(stats.direct_mode_blocks.load(std::memory_order_relaxed)); + serializeReadStage(mslm_pushed_filter_read, mslm->mutable_pushed_filter_read()); + serializeReadStage(mslm_candidate_read, mslm->mutable_candidate_read()); + serializeReadStage(mslm_final_rest_read, mslm->mutable_final_rest_read()); + serializeFilter( + stats.pushed_filter_input_rows.load(std::memory_order_relaxed), + stats.pushed_filter_selected_rows.load(std::memory_order_relaxed), + stats.pushed_filter_filtered_rows.load(std::memory_order_relaxed), + mslm->mutable_pushed_filter()); + serializeFilter( + stats.residual_filter_input_rows.load(std::memory_order_relaxed), + stats.residual_filter_selected_rows.load(std::memory_order_relaxed), + stats.residual_filter_filtered_rows.load(std::memory_order_relaxed), + mslm->mutable_residual_filter()); + if (stats.topn_enabled.load(std::memory_order_relaxed)) + { + auto * topn = mslm->mutable_running_topn(); + topn->set_input_rows(stats.running_topn_input_rows.load(std::memory_order_relaxed)); + topn->set_selected_rows(stats.running_topn_selected_rows.load(std::memory_order_relaxed)); + topn->set_bypass_rows(stats.running_topn_bypass_rows.load(std::memory_order_relaxed)); + topn->set_filtered_rows(stats.running_topn_filtered_rows.load(std::memory_order_relaxed)); + topn->set_heap_size_sum(stats.topn_heap_size_sum.load(std::memory_order_relaxed)); + topn->set_adaptive_warmup_rows(stats.topn_adaptive_warmup_rows.load(std::memory_order_relaxed)); + topn->set_adaptive_post_warmup_input_rows( + stats.topn_adaptive_post_warmup_input_rows.load(std::memory_order_relaxed)); + topn->set_adaptive_post_warmup_candidate_rows( + stats.topn_adaptive_post_warmup_candidate_rows.load(std::memory_order_relaxed)); + topn->set_adaptive_disabled_streams(stats.topn_adaptive_disabled_streams.load(std::memory_order_relaxed)); + } + mslm->set_final_rest_input_rows(stats.final_rest_input_rows.load(std::memory_order_relaxed)); +} + +void ScanContext::deserializeMultiStageLateMaterialization(const tipb::TiFlashScanContext & proto) +{ + if (!proto.has_multi_stage_late_materialization()) + return; + mergeMultiStageLateMaterialization(proto); +} + +void ScanContext::mergeMultiStageLateMaterialization(const ScanContext & other) +{ + if (!other.hasMultiStageLateMaterializationContext()) + return; + + multi_stage_late_materialization_enabled.store(true, std::memory_order_relaxed); + mslm_pushed_filter_read.merge(other.mslm_pushed_filter_read); + mslm_candidate_read.merge(other.mslm_candidate_read); + mslm_final_rest_read.merge(other.mslm_final_rest_read); + merged_multi_stage_late_materialization_runtime_stats.merge(other.getMultiStageLateMaterializationRuntimeStats()); +} + +void ScanContext::mergeMultiStageLateMaterialization(const tipb::TiFlashScanContext & other) +{ + if (!other.has_multi_stage_late_materialization()) + return; + + const auto & mslm = other.multi_stage_late_materialization(); + multi_stage_late_materialization_enabled.store(true, std::memory_order_relaxed); + + if (mslm.has_pushed_filter_read()) + mergeReadStage(mslm_pushed_filter_read, mslm.pushed_filter_read()); + if (mslm.has_candidate_read()) + mergeReadStage(mslm_candidate_read, mslm.candidate_read()); + if (mslm.has_final_rest_read()) + mergeReadStage(mslm_final_rest_read, mslm.final_rest_read()); + + auto & stats = getMergedMultiStageLateMaterializationRuntimeStats(); + stats.finished_streams += mslm.streams(); + stats.late_mode_blocks += mslm.late_mode_blocks(); + stats.direct_mode_blocks += mslm.direct_mode_blocks(); + if (mslm.has_pushed_filter()) + { + const auto & pushed_filter = mslm.pushed_filter(); + stats.pushed_filter_input_rows += pushed_filter.input_rows(); + stats.pushed_filter_selected_rows += pushed_filter.selected_rows(); + stats.pushed_filter_filtered_rows += pushed_filter.filtered_rows(); + stats.stage0_output_rows += pushed_filter.selected_rows(); + } + if (mslm.has_residual_filter()) + { + const auto & residual_filter = mslm.residual_filter(); + stats.residual_filter_input_rows += residual_filter.input_rows(); + stats.residual_filter_selected_rows += residual_filter.selected_rows(); + stats.residual_filter_filtered_rows += residual_filter.filtered_rows(); + stats.stage1_output_rows += residual_filter.selected_rows(); + } + stats.final_rest_input_rows += mslm.final_rest_input_rows(); + stats.topn_candidate_rows += mslm.final_rest_input_rows(); + if (mslm.has_running_topn()) + { + const auto & running_topn = mslm.running_topn(); + stats.topn_enabled.store(true, std::memory_order_relaxed); + stats.running_topn_input_rows += running_topn.input_rows(); + stats.running_topn_selected_rows += running_topn.selected_rows(); + stats.running_topn_bypass_rows += running_topn.bypass_rows(); + stats.running_topn_filtered_rows += running_topn.filtered_rows(); + stats.topn_heap_size_sum += running_topn.heap_size_sum(); + stats.topn_adaptive_warmup_rows += running_topn.adaptive_warmup_rows(); + stats.topn_adaptive_post_warmup_input_rows += running_topn.adaptive_post_warmup_input_rows(); + stats.topn_adaptive_post_warmup_candidate_rows += running_topn.adaptive_post_warmup_candidate_rows(); + stats.topn_adaptive_disabled_streams += running_topn.adaptive_disabled_streams(); + } +} + } // namespace DB::DM diff --git a/dbms/src/Storages/DeltaMerge/ScanContext.h b/dbms/src/Storages/DeltaMerge/ScanContext.h index 14e1b5e4e5e..0eaa038e5d7 100644 --- a/dbms/src/Storages/DeltaMerge/ScanContext.h +++ b/dbms/src/Storages/DeltaMerge/ScanContext.h @@ -18,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -34,6 +35,23 @@ namespace DB::DM { class PushDownFilter; using PushDownFilterPtr = std::shared_ptr; + +struct MSLMReadStageScanContext +{ + std::atomic dmfile_scanned_rows{0}; + std::atomic dmfile_skipped_rows{0}; + std::atomic read_bytes{0}; + std::atomic read_time_ns{0}; + + void merge(const MSLMReadStageScanContext & other) + { + dmfile_scanned_rows += other.dmfile_scanned_rows; + dmfile_skipped_rows += other.dmfile_skipped_rows; + read_bytes += other.read_bytes; + read_time_ns += other.read_time_ns; + } +}; + /// ScanContext is used to record statistical information in table scan for current query. /// For each table scan(one executor id), there is only one ScanContext. /// ScanContext helps to collect the statistical information of the table scan to show in `EXPLAIN ANALYZE`. @@ -46,7 +64,12 @@ class ScanContext std::atomic dmfile_mvcc_skipped_rows{0}; std::atomic dmfile_lm_filter_scanned_rows{0}; std::atomic dmfile_lm_filter_skipped_rows{0}; - std::atomic dmfile_mslm_stage1_filter_scanned_rows{0}; + std::atomic multi_stage_late_materialization_enabled{false}; + MSLMReadStageScanContext mslm_pushed_filter_read; + MSLMReadStageScanContext mslm_candidate_read; + MSLMReadStageScanContext mslm_final_rest_read; + MultiStageLateMaterializationRuntimeStats merged_multi_stage_late_materialization_runtime_stats; + MultiStageLateMaterializationRuntimeStatsPtr multi_stage_late_materialization_runtime_stats; std::atomic total_dmfile_read_time_ns{0}; std::atomic total_rs_pack_filter_check_time_ns{0}; @@ -130,6 +153,7 @@ class ScanContext dmfile_mvcc_skipped_rows = tiflash_scan_context_pb.dmfile_mvcc_skipped_rows(); dmfile_lm_filter_scanned_rows = tiflash_scan_context_pb.dmfile_lm_filter_scanned_rows(); dmfile_lm_filter_skipped_rows = tiflash_scan_context_pb.dmfile_lm_filter_skipped_rows(); + deserializeMultiStageLateMaterialization(tiflash_scan_context_pb); total_rs_pack_filter_check_time_ns = tiflash_scan_context_pb.total_dmfile_rs_check_ms() * 1000000; // TODO: rs_pack_filter_none, rs_pack_filter_some, rs_pack_filter_all,rs_pack_filter_all_null // rs_dmfile_read_with_all @@ -186,6 +210,7 @@ class ScanContext tiflash_scan_context_pb.set_dmfile_mvcc_skipped_rows(dmfile_mvcc_skipped_rows); tiflash_scan_context_pb.set_dmfile_lm_filter_scanned_rows(dmfile_lm_filter_scanned_rows); tiflash_scan_context_pb.set_dmfile_lm_filter_skipped_rows(dmfile_lm_filter_skipped_rows); + serializeMultiStageLateMaterialization(tiflash_scan_context_pb); tiflash_scan_context_pb.set_total_dmfile_rs_check_ms(total_rs_pack_filter_check_time_ns / 1000000); // TODO: pack_filter_none, pack_filter_some, pack_filter_all tiflash_scan_context_pb.set_total_dmfile_read_ms(total_dmfile_read_time_ns / 1000000); @@ -241,7 +266,7 @@ class ScanContext dmfile_mvcc_skipped_rows += other.dmfile_mvcc_skipped_rows; dmfile_lm_filter_scanned_rows += other.dmfile_lm_filter_scanned_rows; dmfile_lm_filter_skipped_rows += other.dmfile_lm_filter_skipped_rows; - dmfile_mslm_stage1_filter_scanned_rows += other.dmfile_mslm_stage1_filter_scanned_rows; + mergeMultiStageLateMaterialization(other); total_rs_pack_filter_check_time_ns += other.total_rs_pack_filter_check_time_ns; rs_pack_filter_none += other.rs_pack_filter_none; rs_pack_filter_some += other.rs_pack_filter_some; @@ -303,6 +328,7 @@ class ScanContext dmfile_mvcc_skipped_rows += other.dmfile_mvcc_skipped_rows(); dmfile_lm_filter_scanned_rows += other.dmfile_lm_filter_scanned_rows(); dmfile_lm_filter_skipped_rows += other.dmfile_lm_filter_skipped_rows(); + mergeMultiStageLateMaterialization(other); total_rs_pack_filter_check_time_ns += other.total_dmfile_rs_check_ms() * 1000000; // TODO: rs_pack_filter_none, rs_pack_filter_some, rs_pack_filter_all, rs_pack_filter_all_null // rs_dmfile_read_with_all @@ -360,9 +386,22 @@ class ScanContext // LACBytesCollector is not thread-safe, to avoid locking, we create a new one for each stream. std::optional newLACBytesCollector(ReadTag read_tag); void addUserReadBytes(size_t bytes, ReadTag read_tag, std::optional & lac_bytes_collector); + void addDMFileReadTime(uint64_t ns, ReadTag read_tag); + void addDMFileScannedRows(uint64_t rows, ReadTag read_tag); + void addDMFileSkippedRows(uint64_t rows, ReadTag read_tag); + void setMultiStageLateMaterializationRuntimeStats(const MultiStageLateMaterializationRuntimeStatsPtr & stats); uint64_t userReadBytes() const { return query_read_bytes + mvcc_read_bytes; } private: + bool hasMultiStageLateMaterializationContext() const; + const MultiStageLateMaterializationRuntimeStats & getMultiStageLateMaterializationRuntimeStats() const; + MultiStageLateMaterializationRuntimeStats & getMergedMultiStageLateMaterializationRuntimeStats(); + MSLMReadStageScanContext * getMutableMSLMReadStage(ReadTag read_tag); + const MSLMReadStageScanContext * getMSLMReadStage(ReadTag read_tag) const; + void serializeMultiStageLateMaterialization(tipb::TiFlashScanContext & proto) const; + void deserializeMultiStageLateMaterialization(const tipb::TiFlashScanContext & proto); + void mergeMultiStageLateMaterialization(const ScanContext & other); + void mergeMultiStageLateMaterialization(const tipb::TiFlashScanContext & other); void serializeRegionNumOfInstance(tipb::TiFlashScanContext & proto) const; void deserializeRegionNumberOfInstance(const tipb::TiFlashScanContext & proto); void mergeRegionNumberOfInstance(const ScanContext & other); diff --git a/dbms/src/Storages/DeltaMerge/Segment.cpp b/dbms/src/Storages/DeltaMerge/Segment.cpp index 31abaeee776..9fbe02fb691 100644 --- a/dbms/src/Storages/DeltaMerge/Segment.cpp +++ b/dbms/src/Storages/DeltaMerge/Segment.cpp @@ -3627,6 +3627,8 @@ BlockInputStreamPtr Segment::getLateMaterializationStream( size_t expected_block_size) { const auto & filter_columns = filter->filter_columns; + const bool enable_multi_stage_late_materialization + = multi_stage_late_materialization_filter && multi_stage_late_materialization_filter->before_where; BlockInputStreamPtr filter_column_stream = getConcatSkippableBlockInputStream( bitmap_filter, segment_snap, @@ -3636,7 +3638,7 @@ BlockInputStreamPtr Segment::getLateMaterializationStream( filter->rs_operator, start_ts, expected_block_size, - ReadTag::LMFilter); + enable_multi_stage_late_materialization ? ReadTag::MSLMPushedFilter : ReadTag::LMFilter); if (unlikely(filter_columns->size() == columns_to_read.size())) { @@ -3682,7 +3684,7 @@ BlockInputStreamPtr Segment::getLateMaterializationStream( dm_context.tracing_id); filter_column_stream->setExtraInfo("push down filter"); - if (multi_stage_late_materialization_filter && multi_stage_late_materialization_filter->before_where) + if (enable_multi_stage_late_materialization) { const auto & stage1_filter_columns = multi_stage_late_materialization_filter->filter_columns; RUNTIME_CHECK(stage1_filter_columns != nullptr); @@ -3705,7 +3707,7 @@ BlockInputStreamPtr Segment::getLateMaterializationStream( filter->rs_operator, start_ts, expected_block_size, - ReadTag::MSLMStage1Filter); + ReadTag::MSLMCandidate); auto final_rest_columns_to_read = std::make_shared(columns_to_read); for (const auto & col : *stage1_filter_columns) @@ -3730,9 +3732,9 @@ BlockInputStreamPtr Segment::getLateMaterializationStream( filter->rs_operator, start_ts, expected_block_size, - ReadTag::Query); + ReadTag::MSLMFinalRest); - LOG_INFO( + LOG_DEBUG( segment_snap->log, "Use multi-stage late materialization, stage0_filter_columns={} stage1_filter_columns={} " "final_rest_columns={}", diff --git a/dbms/src/Storages/DeltaMerge/tests/gtest_dm_delta_merge_store.cpp b/dbms/src/Storages/DeltaMerge/tests/gtest_dm_delta_merge_store.cpp index 71e6d953c92..cbccfbd477d 100644 --- a/dbms/src/Storages/DeltaMerge/tests/gtest_dm_delta_merge_store.cpp +++ b/dbms/src/Storages/DeltaMerge/tests/gtest_dm_delta_merge_store.cpp @@ -268,9 +268,12 @@ try in->readSuffix(); ASSERT_EQ(output_rows, rows); - ASSERT_EQ(scan_context->dmfile_lm_filter_scanned_rows, rows); - ASSERT_EQ(scan_context->dmfile_mslm_stage1_filter_scanned_rows, rows); - ASSERT_NE(scan_context->toJson().find("dmfile_mslm_stage1_filter_scanned_rows"), String::npos); + ASSERT_EQ(scan_context->dmfile_lm_filter_scanned_rows, 0); + ASSERT_EQ(scan_context->mslm_pushed_filter_read.dmfile_scanned_rows, rows); + ASSERT_EQ(scan_context->mslm_candidate_read.dmfile_scanned_rows, rows); + ASSERT_EQ(scan_context->mslm_final_rest_read.dmfile_scanned_rows, rows); + ASSERT_NE(scan_context->toJson().find("multi_stage_late_materialization"), String::npos); + ASSERT_NE(scan_context->toJson().find("candidate_read"), String::npos); } CATCH diff --git a/dbms/src/Storages/DeltaMerge/tests/gtest_dm_file.cpp b/dbms/src/Storages/DeltaMerge/tests/gtest_dm_file.cpp index 3c4324ab3c4..d69fd348a2f 100644 --- a/dbms/src/Storages/DeltaMerge/tests/gtest_dm_file.cpp +++ b/dbms/src/Storages/DeltaMerge/tests/gtest_dm_file.cpp @@ -1705,6 +1705,55 @@ RSOperatorPtr toRSFilter(const ColumnDefine & cd, const HandleRange & range) } } // namespace +TEST_P(DMFileTest, MSLMRoughSetStatsRecordedOnce) +try +{ + auto cols = DMTestEnv::getDefaultColumns(); + ColumnDefine i64_cd(2, "i64", typeFromString("Int64")); + cols->push_back(i64_cd); + reload(cols); + + constexpr Int64 num_rows = 128; + { + auto block = DMTestEnv::prepareSimpleWriteBlock(0, num_rows, false); + block.insert(DB::tests::createColumn(createNumbers(0, num_rows), i64_cd.name, i64_cd.id)); + auto stream = std::make_shared(dbContext(), dm_file, *cols); + stream->writePrefix(); + stream->write(block, {}); + stream->writeSuffix(); + } + + const auto filter = toRSFilter(i64_cd, HandleRange{0, num_rows}); + const auto read_all = [&](ReadTag read_tag, const ScanContextPtr & scan_context) { + DMFileBlockInputStreamBuilder builder(dbContext()); + auto stream = builder.setColumnCache(column_cache) + .setRSOperator(filter) + .setReadTag(read_tag) + .build(dm_file, *cols, RowKeyRanges{RowKeyRange::newAll(false, 1)}, scan_context); + while (stream->read()) {} + }; + const auto pack_filter_count = [](const ScanContextPtr & scan_context) { + return scan_context->rs_pack_filter_none.load() + scan_context->rs_pack_filter_some.load() + + scan_context->rs_pack_filter_all.load() + scan_context->rs_pack_filter_all_null.load(); + }; + + auto query_context = std::make_shared(); + read_all(ReadTag::Query, query_context); + EXPECT_EQ(pack_filter_count(query_context), 1); + EXPECT_EQ(query_context->rs_dmfile_read_with_all, 1); + + auto mslm_context = std::make_shared(); + read_all(ReadTag::MSLMPushedFilter, mslm_context); + EXPECT_EQ(pack_filter_count(mslm_context), 1); + EXPECT_EQ(mslm_context->rs_dmfile_read_with_all, 1); + + read_all(ReadTag::MSLMCandidate, mslm_context); + read_all(ReadTag::MSLMFinalRest, mslm_context); + EXPECT_EQ(pack_filter_count(mslm_context), 1); + EXPECT_EQ(mslm_context->rs_dmfile_read_with_all, 1); +} +CATCH + TEST_P(DMFileTest, ReadFilteredByRoughSetFilter) try { diff --git a/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp b/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp index cc51ac457b2..54e72b9afe3 100644 --- a/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp +++ b/dbms/src/Storages/DeltaMerge/tests/gtest_skippable_block_input_stream.cpp @@ -96,6 +96,14 @@ std::vector makeResidualValues(size_t rows, size_t passed_rows) return values; } +std::vector makeTopNAdaptiveValues(size_t rows, size_t candidate_rows) +{ + RUNTIME_CHECK(candidate_rows <= rows); + std::vector values(rows, 3); + std::fill(values.begin(), values.begin() + candidate_rows, 1); + return values; +} + size_t drainMultiStageStream(const BlockInputStreamPtr & stream) { size_t rows = 0; @@ -276,6 +284,30 @@ PushDownFilterPtr makeResidualFilterForMultiStageTest() nullptr); } +PushDownFilterPtr makeResidualFilterWithTopNForMultiStageTest(UInt64 topk) +{ + auto filter_columns = std::make_shared(makeMultiStageFilterColumns()); + auto actions = std::make_shared(toEmptyBlock(*filter_columns).getNamesAndTypes()); + actions->add(ExpressionAction::copyColumn(MULTI_STAGE_FILTER_COL_NAME, MULTI_STAGE_FILTER_TMP_COL_NAME)); + + auto topn = std::make_shared(); + topn->topk = topk; + topn->order_by_columns.push_back(MultiStageLateMaterializationTopNOrderByColumn{ + .column_id = MULTI_STAGE_FILTER_COL_ID, + .direction = 1, + }); + + return std::make_shared( + EMPTY_RS_OPERATOR, + actions, + nullptr, + filter_columns, + MULTI_STAGE_FILTER_TMP_COL_NAME, + nullptr, + nullptr, + topn); +} + void ensureMultiStageTestFunctionsRegistered() { static std::once_flag once; @@ -689,6 +721,55 @@ try } CATCH +TEST_F(SkippableBlockInputStreamTest, MultiStageLateMaterializationRuntimeStatsFlushesByBlockCount) +try +{ + constexpr size_t blocks_before_flush = 64; + + std::vector stage0_filters; + std::vector> filter_values; + std::vector> rest_values; + for (size_t i = 0; i < blocks_before_flush; ++i) + { + stage0_filters.emplace_back(makeFilter({1})); + filter_values.emplace_back(std::vector{1}); + rest_values.emplace_back(std::vector{static_cast(i)}); + } + + auto stage0_stream = std::make_shared(std::move(stage0_filters)); + auto stage1_stream = std::make_shared( + makeMultiStageFilterColumns(), + filter_values, + rest_values); + auto final_rest_stream = std::make_shared( + makeMultiStageRestColumns(), + filter_values, + rest_values); + + auto bitmap_filter = std::make_shared(blocks_before_flush, 1); + auto runtime_stats = std::make_shared(); + auto stream = std::make_shared( + makeMultiStageColumnsToRead(), + stage0_stream, + stage1_stream, + final_rest_stream, + makeResidualFilterForMultiStageTest(), + bitmap_filter, + "test", + runtime_stats); + + for (size_t i = 0; i < blocks_before_flush - 1; ++i) + assertMultiStageRows(stream->read(), {1}, {static_cast(i)}); + ASSERT_EQ(runtime_stats->stage0_output_rows.load(), 0); + ASSERT_EQ(runtime_stats->stage1_output_rows.load(), 0); + + assertMultiStageRows(stream->read(), {1}, {static_cast(blocks_before_flush - 1)}); + ASSERT_EQ(runtime_stats->stage0_output_rows.load(), blocks_before_flush); + ASSERT_EQ(runtime_stats->stage1_output_rows.load(), blocks_before_flush); + ASSERT_EQ(runtime_stats->final_rest_input_rows.load(), blocks_before_flush); +} +CATCH + TEST_F(SkippableBlockInputStreamTest, MultiStageLateMaterializationLateMode) try { @@ -1007,6 +1088,147 @@ try } CATCH +TEST_F(SkippableBlockInputStreamTest, MultiStageLateMaterializationTopNAdaptiveDoesNotDisableDuringWarmup) +try +{ + constexpr UInt64 topk = 2; + constexpr size_t warmup_rows = multi_stage_late_materialization_topn_adaptive_rows_threshold; + + std::vector stage0_filters; + std::vector> filter_values; + std::vector> rest_values; + stage0_filters.emplace_back(warmup_rows, 1); + filter_values.emplace_back(warmup_rows, 1); + rest_values.emplace_back(makeRestValues(0, warmup_rows)); + + auto stage0_stream = std::make_shared(std::move(stage0_filters)); + auto stage1_stream = std::make_shared( + makeMultiStageFilterColumns(), + filter_values, + rest_values); + auto final_rest_stream = std::make_shared( + makeMultiStageRestColumns(), + filter_values, + rest_values); + + auto bitmap_filter = std::make_shared(warmup_rows, 1); + auto runtime_stats = std::make_shared(); + auto stream = std::make_shared( + makeMultiStageColumnsToRead(), + stage0_stream, + stage1_stream, + final_rest_stream, + makeResidualFilterWithTopNForMultiStageTest(topk), + bitmap_filter, + "test", + runtime_stats); + + ASSERT_EQ(drainMultiStageStream(stream), topk); + ASSERT_EQ(runtime_stats->stage0_output_rows.load(), warmup_rows); + ASSERT_EQ(runtime_stats->stage1_output_rows.load(), warmup_rows); + ASSERT_EQ(runtime_stats->topn_candidate_rows.load(), topk); + ASSERT_EQ(runtime_stats->topn_adaptive_disabled_streams.load(), 0); +} +CATCH + +TEST_F(SkippableBlockInputStreamTest, MultiStageLateMaterializationTopNAdaptiveDoesNotDisableBelowThreshold) +try +{ + constexpr UInt64 topk = 2; + constexpr size_t warmup_rows = multi_stage_late_materialization_topn_adaptive_rows_threshold; + constexpr size_t second_block_rows = multi_stage_late_materialization_topn_adaptive_rows_threshold; + constexpr size_t below_threshold_candidate_rows = second_block_rows + * multi_stage_late_materialization_topn_adaptive_disable_ratio_numerator + / multi_stage_late_materialization_topn_adaptive_disable_ratio_denominator + - 1; + + std::vector stage0_filters; + std::vector> filter_values; + std::vector> rest_values; + stage0_filters.emplace_back(warmup_rows, 1); + filter_values.emplace_back(warmup_rows, 2); + rest_values.emplace_back(makeRestValues(0, warmup_rows)); + stage0_filters.emplace_back(second_block_rows, 1); + filter_values.emplace_back(makeTopNAdaptiveValues(second_block_rows, below_threshold_candidate_rows)); + rest_values.emplace_back(makeRestValues(100000, second_block_rows)); + + auto stage0_stream = std::make_shared(std::move(stage0_filters)); + auto stage1_stream = std::make_shared( + makeMultiStageFilterColumns(), + filter_values, + rest_values); + auto final_rest_stream = std::make_shared( + makeMultiStageRestColumns(), + filter_values, + rest_values); + + auto bitmap_filter = std::make_shared(warmup_rows + second_block_rows, 1); + auto runtime_stats = std::make_shared(); + auto stream = std::make_shared( + makeMultiStageColumnsToRead(), + stage0_stream, + stage1_stream, + final_rest_stream, + makeResidualFilterWithTopNForMultiStageTest(topk), + bitmap_filter, + "test", + runtime_stats); + + ASSERT_EQ(drainMultiStageStream(stream), topk * 2); + ASSERT_EQ(runtime_stats->stage0_output_rows.load(), warmup_rows + second_block_rows); + ASSERT_EQ(runtime_stats->stage1_output_rows.load(), warmup_rows + second_block_rows); + ASSERT_EQ(runtime_stats->topn_candidate_rows.load(), topk * 2); + ASSERT_EQ(runtime_stats->topn_adaptive_disabled_streams.load(), 0); +} +CATCH + +TEST_F(SkippableBlockInputStreamTest, MultiStageLateMaterializationTopNTieRowsDoNotExpandCandidates) +try +{ + constexpr UInt64 topk = 2; + constexpr size_t warmup_rows = multi_stage_late_materialization_topn_adaptive_rows_threshold; + constexpr size_t second_block_rows = multi_stage_late_materialization_topn_adaptive_rows_threshold; + + std::vector stage0_filters; + std::vector> filter_values; + std::vector> rest_values; + stage0_filters.emplace_back(warmup_rows, 1); + filter_values.emplace_back(warmup_rows, 1); + rest_values.emplace_back(makeRestValues(0, warmup_rows)); + stage0_filters.emplace_back(second_block_rows, 1); + filter_values.emplace_back(second_block_rows, 1); + rest_values.emplace_back(makeRestValues(100000, second_block_rows)); + + auto stage0_stream = std::make_shared(std::move(stage0_filters)); + auto stage1_stream = std::make_shared( + makeMultiStageFilterColumns(), + filter_values, + rest_values); + auto final_rest_stream = std::make_shared( + makeMultiStageRestColumns(), + filter_values, + rest_values); + + auto bitmap_filter = std::make_shared(warmup_rows + second_block_rows, 1); + auto runtime_stats = std::make_shared(); + auto stream = std::make_shared( + makeMultiStageColumnsToRead(), + stage0_stream, + stage1_stream, + final_rest_stream, + makeResidualFilterWithTopNForMultiStageTest(topk), + bitmap_filter, + "test", + runtime_stats); + + ASSERT_EQ(drainMultiStageStream(stream), topk); + ASSERT_EQ(runtime_stats->stage0_output_rows.load(), warmup_rows + second_block_rows); + ASSERT_EQ(runtime_stats->stage1_output_rows.load(), warmup_rows + second_block_rows); + ASSERT_EQ(runtime_stats->topn_candidate_rows.load(), topk); + ASSERT_EQ(runtime_stats->topn_adaptive_disabled_streams.load(), 0); +} +CATCH + TEST_F(SkippableBlockInputStreamTest, InMemory2) try { diff --git a/dbms/src/Storages/SelectQueryInfo.cpp b/dbms/src/Storages/SelectQueryInfo.cpp index 2ba3b656e08..84325257d1f 100644 --- a/dbms/src/Storages/SelectQueryInfo.cpp +++ b/dbms/src/Storages/SelectQueryInfo.cpp @@ -34,6 +34,7 @@ SelectQueryInfo::SelectQueryInfo(const SelectQueryInfo & rhs) , has_multiple_partitions(rhs.has_multiple_partitions) , enable_multi_stage_late_materialization(rhs.enable_multi_stage_late_materialization) , multi_stage_late_materialization_runtime_stats(rhs.multi_stage_late_materialization_runtime_stats) + , multi_stage_late_materialization_topn(rhs.multi_stage_late_materialization_topn) {} SelectQueryInfo::SelectQueryInfo(SelectQueryInfo && rhs) noexcept @@ -47,6 +48,7 @@ SelectQueryInfo::SelectQueryInfo(SelectQueryInfo && rhs) noexcept , has_multiple_partitions(rhs.has_multiple_partitions) , enable_multi_stage_late_materialization(rhs.enable_multi_stage_late_materialization) , multi_stage_late_materialization_runtime_stats(std::move(rhs.multi_stage_late_materialization_runtime_stats)) + , multi_stage_late_materialization_topn(std::move(rhs.multi_stage_late_materialization_topn)) {} } // namespace DB diff --git a/dbms/src/Storages/SelectQueryInfo.h b/dbms/src/Storages/SelectQueryInfo.h index 2d6b6e318da..ec4090bd330 100644 --- a/dbms/src/Storages/SelectQueryInfo.h +++ b/dbms/src/Storages/SelectQueryInfo.h @@ -37,6 +37,8 @@ namespace DM { struct MultiStageLateMaterializationRuntimeStats; using MultiStageLateMaterializationRuntimeStatsPtr = std::shared_ptr; +struct MultiStageLateMaterializationTopNDescription; +using MultiStageLateMaterializationTopNDescriptionPtr = std::shared_ptr; } // namespace DM @@ -62,6 +64,7 @@ struct SelectQueryInfo bool has_multiple_partitions = false; bool enable_multi_stage_late_materialization = false; DM::MultiStageLateMaterializationRuntimeStatsPtr multi_stage_late_materialization_runtime_stats; + DM::MultiStageLateMaterializationTopNDescriptionPtr multi_stage_late_materialization_topn; SelectQueryInfo(); ~SelectQueryInfo(); diff --git a/dbms/src/Storages/StorageDeltaMerge.cpp b/dbms/src/Storages/StorageDeltaMerge.cpp index 046b11d2fe5..b2b6b408fb6 100644 --- a/dbms/src/Storages/StorageDeltaMerge.cpp +++ b/dbms/src/Storages/StorageDeltaMerge.cpp @@ -22,6 +22,13 @@ #include #include #include +#include +#include +#include +#include +#include +#include +#include #include #include #include @@ -29,6 +36,7 @@ #include #include #include +#include #include #include #include @@ -792,6 +800,158 @@ RuntimeFilteList parseRuntimeFilterList( return runtime_filter_list; } +bool isSupportedMultiStageLateMaterializationTopNOrderByType(const DataTypePtr & type) +{ + const auto type_not_null = removeNullable(type); + return checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()) // + || checkDataType(type_not_null.get()); +} + +bool isSupportedMultiStageLateMaterializationTopNSourceType(TiDB::TP type) +{ + switch (type) + { + case TiDB::TypeTiny: + case TiDB::TypeShort: + case TiDB::TypeLong: + case TiDB::TypeLongLong: + case TiDB::TypeInt24: + case TiDB::TypeFloat: + case TiDB::TypeDouble: + case TiDB::TypeDecimal: + case TiDB::TypeNewDecimal: + case TiDB::TypeDate: + case TiDB::TypeNewDate: + case TiDB::TypeDatetime: + return true; + default: + return false; + } +} + +const TiDB::ColumnInfo * findSourceColumnByID(const TiDB::ColumnInfos & source_columns, ColumnID column_id) +{ + const auto it = std::find_if(source_columns.begin(), source_columns.end(), [&](const auto & column) { + return column.id == column_id; + }); + return it == source_columns.end() ? nullptr : &*it; +} + +DM::MultiStageLateMaterializationTopNDescriptionPtr buildStorageMultiStageLateMaterializationTopN( + const SelectQueryInfo & query_info, + const DM::ColumnDefines & final_columns_to_read, + const DM::ColumnDefinesPtr & residual_filter_columns, + const LoggerPtr & tracing_logger, + DM::ColumnDefinesPtr & merged_stage1_columns) +{ + auto disable = [&](const String & reason) -> DM::MultiStageLateMaterializationTopNDescriptionPtr { + LOG_DEBUG(tracing_logger, "Disable TopN-enhanced multi-stage late materialization, reason={}", reason); + return nullptr; + }; + + if (!query_info.multi_stage_late_materialization_topn) + return nullptr; + if (query_info.multi_stage_late_materialization_topn->topk == 0) + return disable("topk is zero"); + if (query_info.multi_stage_late_materialization_topn->topk > DM::multi_stage_late_materialization_topn_max_topk) + return disable(fmt::format("topk is too large: {}", query_info.multi_stage_late_materialization_topn->topk)); + if (query_info.multi_stage_late_materialization_topn->order_by_columns.empty()) + return disable("empty order by columns"); + if (query_info.multi_stage_late_materialization_topn->order_by_columns.size() + > DM::multi_stage_late_materialization_topn_max_order_by_columns) + return disable(fmt::format( + "too many order by columns: {}", + query_info.multi_stage_late_materialization_topn->order_by_columns.size())); + + std::unordered_map final_column_map; + final_column_map.reserve(final_columns_to_read.size()); + for (const auto & column : final_columns_to_read) + final_column_map.emplace(column.id, column); + + merged_stage1_columns = std::make_shared(*residual_filter_columns); + std::unordered_set stage1_column_ids; + stage1_column_ids.reserve(merged_stage1_columns->size()); + for (const auto & column : *merged_stage1_columns) + stage1_column_ids.insert(column.id); + + auto storage_topn = std::make_shared(); + storage_topn->topk = query_info.multi_stage_late_materialization_topn->topk; + storage_topn->order_by_columns.reserve(query_info.multi_stage_late_materialization_topn->order_by_columns.size()); + + for (const auto & order_by_column : query_info.multi_stage_late_materialization_topn->order_by_columns) + { + const auto column_id = order_by_column.column_id; + if (column_id == ExtraTableIDColumnID || column_id == ExtraCommitTSColumnID) + return disable(fmt::format("order by virtual column is unsupported, column_id={}", column_id)); + + const auto * source_column = findSourceColumnByID(query_info.dag_query->source_columns, column_id); + if (source_column == nullptr) + return disable( + fmt::format("order by column is not found in table scan source columns, column_id={}", column_id)); + if (source_column->hasGeneratedColumnFlag()) + return disable(fmt::format("order by generated column is unsupported, column_id={}", column_id)); + if (!isSupportedMultiStageLateMaterializationTopNSourceType(source_column->tp)) + return disable(fmt::format( + "unsupported order by source type, column_id={}, tp={}", + column_id, + static_cast(source_column->tp))); + + const auto final_column_it = final_column_map.find(column_id); + if (final_column_it == final_column_map.end()) + return disable( + fmt::format("order by column is not found in storage columns_to_read, column_id={}", column_id)); + if (!isSupportedMultiStageLateMaterializationTopNOrderByType(final_column_it->second.type)) + return disable(fmt::format( + "unsupported order by storage type, column_id={}, type={}", + column_id, + final_column_it->second.type->getName())); + + if (stage1_column_ids.insert(column_id).second) + merged_stage1_columns->push_back(final_column_it->second); + storage_topn->order_by_columns.push_back(order_by_column); + } + + const auto final_rest_col_cnt = final_columns_to_read.size() - merged_stage1_columns->size(); + if (final_rest_col_cnt == 0) + return disable("no final rest columns after adding order by columns"); + if (final_rest_col_cnt < 10) + return disable(fmt::format("too few final rest columns after adding order by columns: {}", final_rest_col_cnt)); + if (final_rest_col_cnt < 2 * merged_stage1_columns->size()) + { + return disable(fmt::format( + "final rest columns are not wide enough after adding order by columns, final_rest_col_cnt={}, " + "stage1_col_cnt={}", + final_rest_col_cnt, + merged_stage1_columns->size())); + } + + LOG_DEBUG( + tracing_logger, + "Enable TopN-enhanced multi-stage late materialization, topk={} order_by_columns={} stage1_columns={} " + "final_rest_columns={}", + storage_topn->topk, + storage_topn->order_by_columns.size(), + merged_stage1_columns->size(), + final_rest_col_cnt); + return storage_topn; +} + DM::PushDownFilterPtr buildMultiStageLateMaterializationFilter( const SelectQueryInfo & query_info, const DM::PushDownFilterPtr & stage0_filter, @@ -832,7 +992,27 @@ DM::PushDownFilterPtr buildMultiStageLateMaterializationFilter( "Build multi-stage late materialization residual filter, stage1_filter_columns={} final_columns_to_read={}", residual_filter->filter_columns->size(), final_columns_to_read.size()); - return residual_filter; + DM::ColumnDefinesPtr merged_stage1_columns; + auto storage_topn = buildStorageMultiStageLateMaterializationTopN( + query_info, + final_columns_to_read, + residual_filter->filter_columns, + tracing_logger, + merged_stage1_columns); + if (storage_topn == nullptr) + return residual_filter; + + // This filter object only carries residual filter actions and the stage1 columns. The actual rough-set + // skipping of MSLM streams must use the original query-level rs_operator from the stage0 filter. + return std::make_shared( + DM::EMPTY_RS_OPERATOR, + residual_filter->before_where, + residual_filter->project_after_where, + merged_stage1_columns, + residual_filter->filter_column_name, + residual_filter->extra_cast, + residual_filter->columns_after_cast, + storage_topn); } } // namespace diff --git a/docs/design/2026-07-25-multi-stage-late-materialization.md b/docs/design/2026-07-25-multi-stage-late-materialization.md index e55e90e4174..4422b4b3631 100644 --- a/docs/design/2026-07-25-multi-stage-late-materialization.md +++ b/docs/design/2026-07-25-multi-stage-late-materialization.md @@ -188,7 +188,7 @@ DirectMode 不需要新增 `full_stream`。它本质上还是两次读取： 启用条件： -- TiFlash setting `dt_enable_multi_stage_late_materialization` 为 true。 +- TiFlash setting `dt_enable_multi_stage_late_materialization >= 1`。 - 当前 query 使用 pipeline 执行模式。 - TableScan 已存在 pushed down filters，即可以形成 Stage 0。 - TableScan 上方存在 residual Selection，即 `filter_conditions.hasValue()`。 @@ -458,13 +458,13 @@ if (filter_conditions.hasValue() && multi_stage_lm_consumed_filter_conditions) 改动点： - 在 TiFlash pipeline TableScan 构建路径中识别 candidate： - - `dt_enable_multi_stage_late_materialization` 为 true。 + - `dt_enable_multi_stage_late_materialization >= 1`。 - `filter_conditions.hasValue()`。 - TableScan 已经有 pushed down filters。 - `generated_column_infos` 为空。 - `read_opts.keep_order == false`。 - 列数满足启用规则。 -- 增加 TiFlash setting `dt_enable_multi_stage_late_materialization`，并增加一个内部 flag，例如 `multi_stage_lm_enabled`。 +- 增加 TiFlash setting `dt_enable_multi_stage_late_materialization`，其中 `0` 表示禁用，`1` 表示启用 selection MSLM，`2` 表示启用 TopN-enhanced MSLM；并增加一个内部 flag，例如 `multi_stage_lm_enabled`。 - 将 residual filter conditions 传入 DeltaMerge read 层。 - 当 `multi_stage_lm_enabled` 为 true 时，跳过 `executePushedDownFilter(...)`。 - 同时把当前 `group_builder.getCurProfileInfos()` 注册给 `filter_conditions.executor_id`。 diff --git a/docs/design/2026-07-26-running-local-topn-for-mslm.md b/docs/design/2026-07-26-running-local-topn-for-mslm.md new file mode 100644 index 00000000000..f53fb413bb1 --- /dev/null +++ b/docs/design/2026-07-26-running-local-topn-for-mslm.md @@ -0,0 +1,1042 @@ +# TiFlash MSLM Running Local TopN 设计 + +- Author: TBD +- Discussion PR: TBD +- Tracking Issue: TBD + +## 目录 + +- [背景](#背景) +- [目标](#目标) +- [非目标](#非目标) +- [MVP 范围](#mvp-范围) +- [核心思路](#核心思路) +- [正确性](#正确性) +- [执行流程](#执行流程) +- [ORDER BY 列对齐](#order-by-列对齐) +- [Running Local TopN 状态](#running-local-topn-状态) +- [Filter 坐标](#filter-坐标) +- [启用规则](#启用规则) +- [可观测性](#可观测性) +- [实现计划](#实现计划) +- [测试方案](#测试方案) +- [性能评估](#性能评估) +- [风险](#风险) +- [备选方案](#备选方案) +- [未解决问题](#未解决问题) + +## 背景 + +TiFlash 现有 multi-stage late materialization (MSLM) 通过把 TableScan 上方的 residual Selection 下沉到 DeltaMerge scan 内部，先读取较少的 filter columns，再延迟读取最终输出所需的 final rest columns。 + +对于下面这类 SQL，MSLM 仍然存在进一步优化空间： + +```sql +SELECT payload_cols +FROM t +WHERE selection_pred +ORDER BY order_key +LIMIT k; +``` + +如果 `payload_cols` 很宽，而 `ORDER BY ... LIMIT` 只需要少量 rows，理想情况下 TiFlash 可以先读取 selection columns 和 order-by columns，筛出 TopN 候选 rows，再读取这些候选 rows 的 payload columns。 + +更彻底的方案是让 storage 支持按 candidate row id 回读 rest columns。但当前 DeltaMerge storage 还没有提供这种能力。短期内可以在现有 MSLM 的 block lock-step 模型里做一个更保守的方案：在每个 MSLM stream 内维护 running local TopN，用它生成当前 block 的 candidate filter，只减少 final rest columns 的读取量。 + +## 目标 + +本设计希望在不引入 candidate row id 回读能力的前提下，减少 `TopN + Selection + TableScan` 场景中 final rest columns 的读取。 + +目标： + +- 复用现有 MSLM 框架，不改变 DeltaMerge storage 的 block-level read/skip/readWithFilter 能力。 +- 在 MSLM Stage 1 中读取 residual filter columns 和 order-by columns。 +- 在每个 MSLM stream 内维护 running local TopN 状态。 +- 使用 running local TopN 生成当前 block 的 candidate filter。 +- 将 candidate filter 合并到 final rest columns 的读取 filter 中。 +- 保留上层原有 TopN executor，由上层 TopN 负责最终全局结果。 + +## 非目标 + +第一版不解决以下问题： + +- 不支持按 candidate row id 随机回读 rest columns。 +- 不移除或替换上层 TopN executor。 +- 不做跨 stream 的 storage 内全局 TopN merge。 +- 不支持复杂 ORDER BY expression。 +- 不支持 join、aggregation、window 或其他会改变 row cardinality 的算子下的 TopN 下推。 +- 第一版不支持没有进入现有 MSLM 路径的 `TopN + TableScan`。例如没有 TableScan pushed filter，或者没有 residual Selection 的查询，后续单独扩展。 +- 不对旧的非 pipeline DAG executor 路径做支持。 +- 不追求和未优化路径在 ORDER BY ties 下选择完全相同的物理 rows。 + +## MVP 范围 + +第一版只支持下面范围： + +- 只支持已经满足现有 MSLM 启用条件的查询。也就是说，TableScan 必须已有 pushed-down filters 作为 Stage 0，且 TableScan 上方必须有 residual Selection 作为 Stage 1。 +- 只支持 `ORDER BY` plain columns，即 TiPB `ByItem.expr` 必须直接是 `ColumnRef`。 +- 只支持 TiDB 已经能下发为 constant TopN limit 的查询。SQL 层的 `OFFSET` 会在 TiDB pushdown TopN 时合并进 TiPB `TopN.limit`。 +- 只支持 TiDB 下发的单表 `TopN -> Selection -> TableScan`。在 TiFlash physical plan 中，这一层 `Selection` 必须已经被合并到 `PhysicalTableScan::filter_conditions`，所以第一版实际只识别 `PhysicalTopN -> PhysicalTableScan(with filter_conditions and pushed_down_filters)`。 +- 第一版不穿透 `PhysicalFilter`。如果 TiFlash physical plan 中仍然存在独立的 `PhysicalFilter -> PhysicalTableScan`，说明 residual Selection 没有进入当前 MSLM Stage 1 识别范围，TopN-enhanced MSLM 直接禁用。 +- 不移除上层 TopN。 +- 只在 MSLM 已经满足启用条件时启用。 +- TopN filter 只用于减少 final rest columns 的读取。 + +后续可以分阶段扩展： + +- 阶段二支持有 TableScan pushed filter 但没有 residual Selection 的 `TopN -> TableScan`。这要求把 residual filter 变成 optional，Stage 1 只读取 order-by columns。 +- 阶段三支持完全没有 filter 的 `TopN -> TableScan`。这需要把 MSLM 从当前的 filter-driven 模型泛化为 order-by-column-driven 模型，Stage 1 order-by stream 需要成为 block driver，工程风险更高。 + +其中本地 TopN 使用的 K 为 TiPB `TopN.limit`： + +```text +K = tipb.TopN.limit +``` + +对 SQL 层的 `LIMIT offset, count`，TiDB 下发 pushed-down TopN 前会把 `PhysicalTopN.Count` 改成 `offset + count`，所以 TiFlash 侧不再额外处理 offset。 + +## 核心思路 + +现有 MSLM 的逻辑可以简化为： + +```text +Stage 0: + 读取 pushed filter columns + 执行 pushed filters + 得到 stage0_filter + +Stage 1: + 根据 stage0_filter 读取 residual filter columns + 执行 residual filters + 得到 residual_filter + +Final: + 根据 stage0_filter/residual_filter 读取 final rest columns +``` + +加入 running local TopN 后，逻辑变为： + +```text +Stage 0: + 读取 pushed filter columns + 执行 pushed filters + 得到 stage0_filter + +Stage 1: + 根据 stage0_filter 读取 residual filter columns + order-by columns + 执行 residual filters + 执行 running local TopN + 得到 topn_candidate_filter + +Final: + 根据 stage0_filter && residual_filter && topn_candidate_filter + 读取 final rest columns +``` + +`topn_candidate_filter` 只用于减少 final rest columns 的读取。上层 TopN executor 仍然保留，因此 storage 内输出可以是全局 TopN 的 superset。 + +当前 MSLM 已经不再使用旧的 direct/late adaptive mode。TopN-enhanced MSLM 因此不需要根据 residual 过滤率在 direct/late 两种策略之间做启发式选择，而是只保留下面三个自然分支： + +- `residual_filter && topn_candidate_filter` 全过滤时，skip final rest stream。 +- `residual_filter && topn_candidate_filter` 全通过时，final rest stream 按 Stage 0 filter 直接读取。 +- 其他部分通过场景，把 Stage 1 坐标的 filter compose 回 Stage 0 原始坐标，再用 combined filter 读取 final rest columns。 + +但是 running local TopN 本身有额外代价：Stage 1 需要读取 order-by columns，执行层需要做 comparator、heap 维护和 candidate filter 合并。如果 TopN 几乎不能减少 final rest columns 的读取，就可能带来性能回退。 + +因此第一版增加一个运行时 one-way adaptive disable： + +```text +开始时启用 running local TopN +先用 warmup rows 建立初始 TopN 状态 +warmup 后，持续统计后续 `adaptive_post_warmup_input_rows` 和 `adaptive_post_warmup_candidate_rows` +如果后续累计 input rows 达到最小检查门槛，且 candidate rows / input rows 过高 + 后续 blocks 禁用 TopN filter + 后续 final rest 只使用 residual filter +禁用后不再重新启用 +``` + +这个 adaptive 只做运行时止损，不改变正确性：禁用后输出的是 residual Selection 后的更多 rows，仍然是上层全局 TopN 的 superset。上层 TopN executor 保留，因此最终结果不变。 + +需要注意，这个 runtime disable 不能避免已经加入 Stage 1 的 order-by columns IO，因为 Stage 1 schema 在 stream 构建时已经固定。它主要减少后续 comparator、heap 维护和 TopN filter 处理成本。避免明显不划算的 order-by column 读取仍然依赖静态启发式 gating。 + +## 正确性 + +### Local TopK 覆盖 Global TopK + +如果没有 ties，对于任意一个 stream，如果某一行 `r` 是全局 TopK 结果中的一行，那么 `r` 必然属于它所在 stream 的 local TopK。 + +证明： + +如果 `r` 不属于所在 stream 的 local TopK，说明这个 stream 中至少有 K 行按照相同 comparator 排在 `r` 前面。这 K 行在全局也排在 `r` 前面，因此 `r` 不可能是全局 TopK，矛盾。 + +所以在没有 ties 时： + +```text +global TopK subset union(each stream local TopK) +``` + +有 ties 时，local TopK 不一定保留和全局 TopK 完全相同的一组物理 rows，因为相同 sort key 的 rows 可以在不同执行路径中选择不同 tie winners。此时需要的正确性条件不是固定物理 rows 的严格 subset，而是每个 stream 输出足够多的 rows，且这些 rows 的 sort key 不劣于该 stream 内第 K 个 sort key。这样所有 stream 输出的 union 仍然包含足够多的 rows，其 sort key 不劣于全局第 K 个 sort key，上层 TopN 可以在 SQL ties 语义下得到合法结果。 + +第一版保留上层全局 TopN executor，因此只要 storage 内输出包含上层 TopN 生成合法结果所需的 candidate rows，最终结果就是正确的。 + +### Running Local TopN 的 Streaming Superset + +MSLM stream 是 streaming 输出，不能撤回已经输出的 rows。Running local TopN 处理后续 blocks 时，可能会把之前已经输出的 candidate rows 挤出当前 local TopK。 + +这不会影响正确性，因为这些 rows 只是额外候选。上层 TopN 会再次排序并裁剪。 + +因此 running local TopN 输出的是： + +```text +final local TopK 的 superset +``` + +而不是严格等于 final local TopK。 + +### 被跳过 rows 不会重新进入 TopK + +当 local TopN heap 已满时，如果当前 row 按 comparator 不优于 heap worst row，则该 row 可以被跳过。原因是后续扫描只会增加竞争者，不会让这个 row 的排序位置变得更靠前。 + +对于和 heap worst row 相等的 ties，第一版不会继续输出这些相等 rows 作为 candidate，而是让 heap 中已有 entries 成为本 stream 的 local tie winners。这样可以避免 ties 大量扩展 candidate rows 和 final-rest IO。代价是 ties 下可能选择不同的物理 rows，但 SQL `ORDER BY` 未提供完整 tie-break key 时本来就不保证相同 sort key 内的物理 row 选择，这也属于本设计的非目标。 + +## 执行流程 + +### 计划识别 + +TiFlash 在构建 TableScan pipeline 时识别下面模式： + +```text +TopN + | +Selection + | +TableScan +``` + +识别成功后，生成一个 MSLM TopN 描述： + +```text +struct MSLMTopNDescription +{ + SortDescription storage_sort_description; + UInt64 topk; // tipb.TopN.limit + ColumnDefines order_by_columns; +}; +``` + +第一版 `order_by_columns` 只允许 plain columns，不构造复杂 expression actions。 + +Plan 识别建议在 `PhysicalTopN` 和 `PhysicalTableScan` 之间完成。`PhysicalTopN` 只在 child 直接是 `PhysicalTableScan` 时尝试生成 MSLM TopN 描述；如果 child 是 `PhysicalFilter`、Projection、Join、Aggregation、Window 等其他节点，则不启用该优化。上层 `PhysicalTopN` executor 仍然保留，storage 内描述只用于减少 final rest columns 的读取。 + +对于 TiPB DAG 中的单层 `TopN -> Selection -> TableScan`，TiFlash physical planner 会先尝试把 `Selection` 合并到 `PhysicalTableScan::filter_conditions`。因此第一版实际需要识别的 physical plan 形态是： + +```text +PhysicalTopN + | +PhysicalTableScan(with filter_conditions and pushed_down_filters) +``` + +这里的 `filter_conditions` 是 TiFlash physical planner 内部保存的 residual Selection。对于 TiPB DAG 中形态为 `Selection -> TableScan` 的单个 Selection，TiFlash 构建 physical plan 时会把它放入 `PhysicalTableScan::filter_conditions`，而不是保留一个单独的 `PhysicalFilter`。如果 Selection 的 child 不是 TableScan，或者同一个 TableScan 已经设置过 `filter_conditions`，则会保留普通 `PhysicalFilter`，第一版不处理这种情况。 + +需要注意，`PhysicalTableScan::filter_conditions` 不等于 TiPB `TableScan.pushed_down_filter_conditions`： + +- `TableScan.pushed_down_filter_conditions` 是现有 MSLM 的 Stage 0 pushed filter。 +- `PhysicalTableScan::filter_conditions` 是 Stage 1 residual Selection。 +- 第一版 TopN-enhanced MSLM 要求这两者都存在。 + +### ORDER BY 列对齐 + +即使第一版只支持 plain column，也需要显式处理 TopN order-by column 和 DeltaMerge storage column 的对齐。 + +不要直接使用 `PhysicalTopN` 现有 `SortDescription.column_name` 去匹配 storage column。这个名字来自 TopN child schema，而 DeltaMerge Stage 1 block 使用的是 `ColumnDefine.name`。TableScan 后还有 schema projection、特殊列名和 cast-after-TS 等逻辑，直接按 name 对齐容易和 storage header 不一致。 + +第一版采用基于 column id 的对齐流程： + +```text +tipb::ByItem.expr() 必须是 ColumnRef + -> decode ColumnRef index + -> table_scan.getColumns()[index].id + -> 在 columns_to_read 或 table column defines 中按 column id 找到 ColumnDefine + -> 将该 ColumnDefine 加入 Stage 1 columns + -> 使用 ColumnDefine.name 构造 storage-side SortDescription +``` + +`MSLMTopNDescription::storage_sort_description` 中的 column name 必须是 Stage 1 block 中真实存在的 storage column name，而不是 TopN child schema name。 + +Stage 1 column union 和 final rest column subtraction 都必须按 `ColumnID` 去重： + +```text +stage1_columns = residual_filter_columns union order_by_columns +final_rest_columns = final_columns_to_read - stage1_columns +``` + +如果 order-by column 已经属于 residual filter columns，则 Stage 1 只读一次。如果 order-by column 也是最终输出列，则 final rest columns 中必须排除该列，避免重复读取。 + +第一版不支持没有 residual Selection 的 `TopN -> TableScan`。这种场景需要将 candidate read 从 residual-filter-driven 扩展为 order-by-column-driven；scan context 的 `residual_filter` 已经按 optional 设计，可以兼容后续扩展。 + +plain column 还需要保证 storage 内比较语义和上层 TopN 一致。第一版建议保守禁用下面场景： + +- order-by expr 不是 direct `ColumnRef`。 +- order-by column 是 generated column 或 virtual `_tidb_tid`。 +- order-by column 需要 TableScan 后的 extra cast 才能得到上层 TopN 的比较值。 +- order-by column 是 variable-length string 且暂未确认 collator 和 owned key 内存上界。 +- order-by column 不在第一版支持的类型白名单内。 + +第一版 order-by column 类型白名单为： + +- integer numeric types。 +- decimal types。 +- date/datetime types。 +- float types。 + +第一版明确不支持： + +- `String` / `FixedString` / `Bytes`，因为需要处理 collation 和 owned key 内存上界。 +- `Timestamp`，因为它通常需要 TableScan 后的 timezone extra cast，storage 内底层值不一定等于上层 TopN 的比较值。 +- `Time` / `Duration`，除非后续明确确认它和上层 TopN 使用完全一致的底层比较表示。 +- JSON、Enum/Set、Bit、Vector、Array、Tuple、Map 等其他类型。 + +### Stage 1 列集合 + +开启 TopN-enhanced MSLM 后，Stage 1 需要读取： + +```text +stage1_columns = residual_filter_columns union order_by_columns +``` + +如果某个 order-by column 已经属于 residual filter columns，则去重后只读一次。 + +如果 order-by column 也属于最终输出 columns，它仍然可以作为 Stage 1 column 提前读取，final rest columns 中需要排除该列，避免重复读取。 + +### Running Local TopN 处理 + +每个 `MultiStageLateMaterializationBlockInputStream` 持有一个 local TopN heap。 + +处理一个 block 时： + +1. Stage 0 读取 pushed filter columns，得到 `stage0_filter`。 +2. Stage 1 根据 `stage0_filter` 读取 `stage1_columns`。 +3. 执行 residual filters，得到 `residual_filter`。 +4. 对 residual passed rows 计算 order key。 +5. 使用 order key 更新 local TopN heap。 +6. 为当前 block 生成 `topn_candidate_filter`。 +7. 使用 combined filter 读取 final rest columns。 +8. hstack Stage 1 columns 和 final rest columns。 +9. 上层 TopN executor 继续执行全局 TopN。 + +伪代码： + +```text +for each stage0 block: + stage0_filter = execute_stage0_filter(block) + if stage0_filter passes no rows: + skip stage1 stream + skip final_rest stream + continue + + stage1_block = read stage1 columns with stage0_filter + residual_filter = execute_residual_filter(stage1_block) + if residual_filter passes no rows: + skip final_rest stream + continue + + topn_candidate_filter = running_topn.update(stage1_block, residual_filter) + if topn_candidate_filter passes no rows: + skip final_rest stream + continue + + combined_filter = compose(stage0_filter, residual_filter, topn_candidate_filter) + final_rest_block = read final_rest columns with combined_filter + stage1_block = filter stage1_block by residual_filter && topn_candidate_filter + return hstack(stage1_block, final_rest_block) +``` + +### Current Block 内 Eviction + +如果当前 block 的某个 row 先进入 heap，但随后在同一个 block 内被更优 row 挤出，第一版可以把它从当前 block 的 candidate filter 中清除。 + +如果被挤出的 row 来自之前已经输出的 block，则不能撤回。该 row 会作为额外候选留给上层 TopN 过滤。 + +Heap entry 只能记录轻量信息，不能持有 `Block` 或 column 引用： + +```text +struct HeapEntry +{ + OwnedSortKey key; + UInt64 stream_sequence; + UInt64 block_sequence; + UInt32 row_index_in_stage1_block; +}; +``` + +`block_sequence` 用于判断被 eviction 的 entry 是否属于当前 block。`row_index_in_stage1_block` 只在 entry 属于当前 block 时有效。 + +如果 eviction 的 entry 来自当前 block，可以将 `topn_candidate_filter[row_index_in_stage1_block]` 置为 false，避免对该 row 读取 final rest columns。 + +如果 eviction 的 entry 来自历史 block，则不做任何回撤。历史 block 已经输出给上层 TopN，该 row 只是额外候选。 + +### Heap 内存模型 + +Running local TopN 不能保存 `{Block, row_id}`。如果 heap entry 持有 `Block` 或 column memory 引用，会导致历史 blocks 被 retain，内存随扫描推进累积。 + +Heap 中只保存排序需要的 owned key： + +```text +heap memory = O(topk * order_key_size) +current block filter memory = O(current_block_rows) +``` + +不保存： + +- 历史 block。 +- 历史 row payload。 +- 历史 column 引用。 +- 用于回读的 row locator。 + +历史 candidate rows 一旦输出给上层 TopN，MSLM stream 不再负责保存它们。Heap eviction 只释放 owned sort key。 + +对于第一版支持的 order-by 类型，owned key 只拷贝 fixed-size value 和 NULL flag。对于 variable-length 类型，例如 String，owned key 可能带来较大内存开销，并且需要处理 collation。第一版禁用 variable-length order-by columns；后续如果要支持，可以再增加 per-stream memory guard 或 owned collation key。 + +如果采用 memory guard，触发后应退化为普通 MSLM： + +```text +后续 blocks 不再生成 TopN candidate filter。 +final rest columns 按现有 MSLM 逻辑读取。 +上层 TopN 保留，保证最终结果正确。 +``` + +Ties 也需要避免破坏内存和 IO 上界。第一版 heap 仍最多保存 `topk` 个 entries；与 heap worst 相等的 rows 不插入 heap，也不作为当前 block candidate 输出。这样 ties 较多时 candidate rows 不会被 ties 放大，heap 内存保持 bounded，但不同执行路径可能选择不同的物理 tie winners。 + +## Running Local TopN 状态 + +每个 MSLM stream 维护独立状态： + +```text +class RunningLocalTopN +{ +public: + Filter update(const Block & stage1_block, const FilterPtr & residual_filter); + +private: + UInt64 topk; + SortDescription sort_description; + PriorityQueue heap; // heap top is current worst candidate + UInt64 current_block_sequence; +}; +``` + +Heap 大小最多为 `topk`。如果 `topk` 超过启发式阈值，则禁用该优化。 + +比较器必须和上层 TopN 使用的 comparator 保持一致。第一版只支持 plain order-by columns，减少 comparator 语义不一致风险。 + +第一版不要求 heap entry 能定位历史 block。Heap entry 的 locator 只服务于 current block eviction。历史 entry 被 eviction 时只释放 key，不访问历史 block。 + +### Sort Key 表示 + +第一版不使用 order-preserving bytes key。TiFlash 现有 sort 路径本身也是基于 `IColumn::compareAt()` / typed compare 直接比较 column values，而不是先构造可按 bytes 比较的 sort key。引入 bytes key 需要重新设计 signed integer、float、NaN、DESC、NULL、decimal 等类型的 order-preserving encoding，正确性风险较高。 + +第一版也不直接使用通用 `Field` 作为 heap key。`Field` 虽然已经是 owned variant，但它支持 String/Array/Tuple 等当前不需要的类型，大小和通用抽取路径都偏重；float 的 NaN 比较语义也需要额外绕开。为了降低 hot path 开销，第一版使用一个更窄的 typed sort key： + +```text +static constexpr size_t max_sort_key_columns = 4; + +enum class SortKeyKind +{ + Int64, + UInt64, + Float32, + Float64, + Decimal32, + Decimal64, + Decimal128, + Decimal256, + Date, + DateTime, +}; + +struct SortKeyField +{ + bool is_null; + SortKeyKind kind; + TypedFixedSizeValue value; +}; + +struct OwnedSortKey +{ + std::array fields; + size_t size; +}; +``` + +`TypedFixedSizeValue` 是 conceptual storage，实际实现可以使用 union 或其他 fixed-size typed storage。它只需要覆盖第一版白名单内的 fixed-size values，不支持 String/Array/Tuple。 + +`RunningLocalTopN` 本身不模板化。每个 order-by column 在初始化时生成一个 descriptor，descriptor 中保存 typed extractor 和 comparator： + +```text +struct SortKeyColumnDesc +{ + String column_name; + size_t column_pos; + SortKeyKind kind; + int direction; + + extract_owned(column, row, SortKeyField &); + compare_owned(SortKeyField, SortKeyField); + compare_column_with_owned(column, row, SortKeyField); +}; +``` + +模板只用于生成具体类型的 extractor/comparator。这样一个 `OwnedSortKey` 可以同时包含多个不同类型的 order-by columns，而不需要把整个 `RunningLocalTopN` 或 `SortKeyField` 做成模板。 + +### Sort Key 比较语义 + +NULL 在所有支持类型中统一处理： + +```text +NULL is minimum before applying direction. +ASC -> NULL first +DESC -> NULL last +``` + +非 NULL values 的比较规则： + +- integer numeric、date、datetime 直接比较底层整数值。 +- decimal 直接比较底层 decimal value。因为 lhs/rhs 来自同一个 order-by column，decimal type 和 scale 一致。 +- float 使用 TiFlash 现有 `CompareHelper::compare(..., nan_direction_hint=-1)` 语义，避免 NaN 行为和上层 sort 不一致。 + +每个字段先得到不带 direction 的比较结果，再乘以 `SortKeyColumnDesc::direction`。多列 ORDER BY 按顺序比较，直到第一个非 0 结果。 + +### Lazy Owned Key Materialization + +第一版不对 residual passed 的每一行都构造 `OwnedSortKey`。处理当前 `stage1_block` 时，可以直接使用 order-by columns 参与比较；只有当前 row 真的需要进入 heap 时，才 lazy materialize owned key。 + +处理一个 block 前，先从 `stage1_block` 中收集 order-by column 指针： + +```text +sort_columns[i] = stage1_block.getByPosition(desc[i].column_pos).column.get() +``` + +per-row hot path 不再从 `Block` 里按 name 或 position 查找 column，而是直接使用预先收集好的 `IColumn *`： + +```text +compareRowWithOwnedKey(sort_columns, row, heap.worst().key) +materializeOwnedKey(sort_columns, row) +``` + +heap 中始终只保存 owned key，不保存 `Block`、`IColumn *` 或 historical row reference。当前 row 只有在 `update()` 函数内部比较时可以直接读 `stage1_block` column。这样既避免 retain historical blocks，又避免给明显不可能进入 local TopK 的 rows 做 key copy。 + +当前 block 内如果某个 entry 被后续更优 row eviction，可以通过 `block_sequence` 和 `row_index_in_stage1_block` 将对应的 `topn_candidate_filter` 置为 false。历史 block 的 entry 被 eviction 时不能回撤，因为 rows 已经输出给上层 TopN。 + +## Filter 坐标 + +MSLM 中存在多套 filter 坐标： + +- `stage0_filter`: 原始 Stage 0 block rows 坐标。 +- `residual_filter`: Stage 1 block rows 坐标，即 Stage 0 passed rows 坐标。 +- `topn_candidate_filter`: Stage 1 block rows 坐标。 +- `combined_filter`: 原始 Stage 0 block rows 坐标，用于读取 final rest columns。 + +组合规则： + +```text +stage1_final_filter = residual_filter && topn_candidate_filter +combined_filter = compose(stage0_filter, stage1_final_filter) +``` + +当 `stage0_filter` 为 nullptr，即 Stage 0 全通过时，`combined_filter` 可以直接基于 `stage1_final_filter` 构造。 + +当 `stage1_final_filter` 全通过时，可以退化为现有 MSLM 路径。 + +当 `stage1_final_filter` 全过滤时，需要 skip final rest stream。 + +## 启用规则 + +第一版使用保守 rule-based 判断。 + +必须满足： + +- `dt_enable_multi_stage_late_materialization == 2`。 +- 当前查询已经满足 MSLM 启用条件。 +- 查询包含 TopN，且 TopN 的 child 直接是单表 `PhysicalTableScan`。 +- `PhysicalTableScan` 已经包含 residual `filter_conditions`，并且 TableScan 已经包含 pushed-down filters。 +- TopN 上方仍保留全局 TopN executor。 +- TiPB `TopN.limit` 已经包含 SQL 层 offset，并且未超过阈值。 +- TopN 之前的 Selection predicates 全部在 MSLM Stage 0 或 Stage 1 内执行。 +- ORDER BY 只包含 direct `ColumnRef` plain columns。 +- ORDER BY columns 能通过 `ColumnRef` index 映射到 `table_scan.getColumns()` 中的 column id，并能进一步映射到 DeltaMerge `ColumnDefine`。 +- ORDER BY columns 不需要 TableScan 后 extra cast。 +- ORDER BY columns 类型属于第一版白名单：integer numeric、decimal、date、datetime、float。 +- ORDER BY columns 可加入 Stage 1 column set。 +- ORDER BY columns 数量不超过阈值。 +- final rest columns 数量满足 MSLM 的宽列启发式要求。 + +建议第一版阈值： + +```text +topk <= 2048 +order_by_column_count <= 4 +runtime adaptive warmup rows >= max(topk * 4, 8192) +after warmup, check only when adaptive_post_warmup_input_rows >= max(topk * 4, 8192) +after warmup, disable if adaptive_post_warmup_candidate_rows / adaptive_post_warmup_input_rows >= 0.7 +``` + +实际阈值可以根据 benchmark 调整。 + +禁用场景： + +- ORDER BY expression 不是 plain column。 +- ORDER BY plain column 无法按 column id 映射到 storage column。 +- ORDER BY column 需要 TableScan 后 extra cast。 +- ORDER BY column 类型不在第一版白名单内，例如 String、FixedString、Timestamp、Time/Duration、JSON、Enum/Set、Bit、Vector、Array、Tuple、Map。 +- ORDER BY 依赖 generated column。 +- TopN 包含非常大的 offset。 +- TopN 之前仍有无法进入 MSLM 的 Selection。 +- TopN 的 child 是 `PhysicalFilter`，即 residual Selection 没有被合并到 `PhysicalTableScan::filter_conditions`。 +- 查询没有 residual Selection，或者没有 TableScan pushed-down filters。 +- 查询要求 keep order 且和 MSLM 现有限制冲突。 +- TopN 位于 join、aggregation、window 等算子之上。 + +## 可观测性 + +`EXPLAIN ANALYZE` 中各 executor 的 `actRows` 语义保持现状： + +- TableScan executor 的 `actRows` 展示 Stage 0 pushed filter 之后的 rows。 +- 没有启用 running local TopN 时，Selection executor 的 `actRows` 展示 residual Selection 之后的 rows。 +- 启用 running local TopN 时，Selection executor 的 `actRows` 展示 MSLM 最终需要读取 final rest columns 的 rows。 + +MSLM 内部细节通过 TiFlash scan context 展示。原 single-stage late materialization 继续使用现有字段，不改变语义： + +- `lm_skip_rows` +- `dmfile_lm_filter_scanned_rows` +- `dmfile_lm_filter_skipped_rows` + +MSLM 不复用这些旧 LM 字段，而是在 `TiFlashScanContext` 中新增独立的 `multi_stage_late_materialization` context。这样可以避免 single-stage LM 和 MSLM 的统计混在一起解释。 + +建议 TiPB 结构如下： + +```proto +message TiFlashScanContext { + // Existing fields unchanged. + optional uint64 dmfile_lm_filter_scanned_rows = 33; + optional uint64 dmfile_lm_filter_skipped_rows = 34; + + // Only set when multi-stage late materialization is enabled. + optional TiFlashMultiStageLateMaterializationScanContext multi_stage_late_materialization = 35; + + // Existing vector index fields unchanged. + optional uint64 total_vector_idx_load_from_s3 = 100; + // ... +} + +message TiFlashMultiStageLateMaterializationScanContext { + optional uint64 streams = 1; + optional uint64 late_mode_blocks = 2; + optional uint64 direct_mode_blocks = 3; + + // Physical read stages. These fields describe DMFile/storage read cost only. + optional TiFlashMSLMReadStageContext pushed_filter_read = 10; + optional TiFlashMSLMReadStageContext candidate_read = 11; + optional TiFlashMSLMReadStageContext final_rest_read = 12; + + // Logical row reducers. Each reducer is independently optional. + optional TiFlashMSLMFilterContext pushed_filter = 20; + optional TiFlashMSLMFilterContext residual_filter = 21; + optional TiFlashMSLMRunningTopNContext running_topn = 22; + + // Logical rows that actually need final-rest materialization. + optional uint64 final_rest_input_rows = 30; +} + +message TiFlashMSLMReadStageContext { + optional uint64 dmfile_scanned_rows = 1; + optional uint64 dmfile_skipped_rows = 2; + optional uint64 read_bytes = 3; + optional uint64 read_time_ms = 4; +} + +message TiFlashMSLMFilterContext { + optional uint64 input_rows = 1; + optional uint64 selected_rows = 2; + optional uint64 filtered_rows = 3; +} + +message TiFlashMSLMRunningTopNContext { + optional uint64 input_rows = 1; + optional uint64 selected_rows = 2; + optional uint64 bypass_rows = 3; + optional uint64 filtered_rows = 4; + optional uint64 heap_size_sum = 5; + + optional uint64 adaptive_warmup_rows = 10; + optional uint64 adaptive_post_warmup_input_rows = 11; + optional uint64 adaptive_post_warmup_candidate_rows = 12; + optional uint64 adaptive_disabled_streams = 13; +} +``` + +### Logical Reducer 语义 + +`pushed_filter` 表示 pushed filter 这个逻辑 reducer。只要 MSLM 中存在 pushed filter，就填该字段： + +- `input_rows`: 进入 pushed filter 的逻辑 rows。 +- `selected_rows`: pushed filter 之后留下、会进入后续 MSLM stage 的逻辑 rows。 +- `filtered_rows`: 被 pushed filter 过滤掉的逻辑 rows。 + +`residual_filter` 表示 residual Selection 这个逻辑 reducer。没有 residual filter 时不填该字段，不用 `0` 表示不存在： + +- `input_rows`: 进入 residual filter 的逻辑 rows。当前实现中通常等于 `pushed_filter.selected_rows`。 +- `selected_rows`: residual filter 之后留下、会进入后续 TopN 或 final rest 读取的逻辑 rows。 +- `filtered_rows`: 被 residual filter 过滤掉的逻辑 rows。 + +filter reducer 的 rows 满足： + +```text +filter.input_rows = filter.selected_rows + filter.filtered_rows +``` + +`running_topn` 表示 running local TopN 这个逻辑 reducer。没有 running local TopN 时不填该字段： + +- `input_rows`: 进入 running local TopN 的逻辑 rows。 +- `selected_rows`: 被 running local TopN 选中、需要继续读取 final rest columns 的 rows。 +- `bypass_rows`: runtime adaptive disable 之后，不再经过 TopN filter、直接进入 final rest 读取的 rows。 +- `filtered_rows`: 被 running local TopN 过滤掉、因此不需要读取 final rest columns 的 rows。 +- `heap_size_sum`: 所有 streams 的 heap size 总和。它不是 `selected_rows`。 + +`running_topn` 的 rows 满足： + +```text +running_topn.input_rows + = running_topn.selected_rows + + running_topn.bypass_rows + + running_topn.filtered_rows +``` + +`final_rest_input_rows` 表示最终需要读取 final rest columns 的逻辑 rows。它和物理读行数不同： + +```text +with running_topn: + final_rest_input_rows = running_topn.selected_rows + running_topn.bypass_rows + +without running_topn and with residual_filter: + final_rest_input_rows = residual_filter.selected_rows + +without running_topn and without residual_filter: + final_rest_input_rows = pushed_filter.selected_rows +``` + +如果未来支持没有 pushed filter、直接由 TopN 驱动的 MSLM，`pushed_filter` 和 `residual_filter` 可以都不存在，`running_topn.input_rows` 表示进入 TopN 的候选 rows，`final_rest_input_rows` 仍然等于 `running_topn.selected_rows + running_topn.bypass_rows`。 + +### Physical Read Stage 语义 + +`pushed_filter_read`、`candidate_read`、`final_rest_read` 只表示物理读成本，不承担 logical input/output rows 语义。因此 read stage 中不设置 `input_rows/output_rows` 字段。 + +- `pushed_filter_read`: 为 pushed filter columns 做的物理读取。 +- `candidate_read`: 为产生 final-rest candidate row set 做的物理读取。当前实现中，它读取 residual filter columns 和 order-by columns；未来 topn-only 路径中，它可以只读取 order-by columns。 +- `final_rest_read`: 为最终输出所需 rest columns 做的物理读取。 + +这些 read stage 的字段含义： + +- `dmfile_scanned_rows`: DMFile 层实际 scan 的 rows。 +- `dmfile_skipped_rows`: DMFile 层跳过的 rows。 +- `read_bytes`: 该 read stage 实际读取的 bytes。 +- `read_time_ms`: 该 read stage 的读取耗时。 + +`final_rest_input_rows` 和 `final_rest_read.dmfile_scanned_rows` 不能混用： + +```text +final_rest_input_rows: + logical rows that need final-rest materialization. + +final_rest_read.dmfile_scanned_rows: + physical rows scanned by DMFile for final-rest columns. It may be larger than + final_rest_input_rows because of pack/block-level read amplification. +``` + +### ReadTag 拆分 + +为了让 MSLM 不污染原 single-stage LM 的统计，TiFlash 内部建议拆分 read tag： + +```cpp +ReadTag::LMFilter // existing single-stage LM filter read +ReadTag::MSLMPushedFilter // MSLM pushed filter read +ReadTag::MSLMCandidate // MSLM residual/order-by candidate read +ReadTag::MSLMFinalRest // MSLM final rest read +``` + +其中 `ReadTag::LMFilter` 继续填旧字段；`MSLMPushedFilter`、`MSLMCandidate`、`MSLMFinalRest` 只填新的 `multi_stage_late_materialization` context。RU 统计上这些 MSLM read tags 仍然属于 query read。 + +### Optional 字段规则 + +不存在或无意义的字段不填，不用 `0` 表示不存在。`0` 只表示该阶段或 reducer 真实存在，但计数结果为 0。 + +例如没有 residual filter 时： + +```json +{ + "multi_stage_late_materialization": { + "pushed_filter": { + "input_rows": 1000000, + "selected_rows": 120000, + "filtered_rows": 880000 + }, + "final_rest_input_rows": 120000 + } +} +``` + +而不是： + +```json +{ + "multi_stage_late_materialization": { + "residual_filter": { + "input_rows": 0, + "selected_rows": 0, + "filtered_rows": 0 + } + } +} +``` + +## 实现计划 + +### Step 1: 文档和测试用例设计 + +- 明确 MVP 范围。 +- 增加基于 DAG request 的单测设计。 +- 明确 ties、offset、NULL ordering、float NaN、以及 unsupported string/collation 的禁用测试覆盖。 + +### Step 2: Pattern 识别和元数据传递 + +- 在 `PhysicalTopN` 构建阶段识别 child 直接是 `PhysicalTableScan` 的模式。 +- 只在 `PhysicalTableScan` 同时包含 residual `filter_conditions` 和 TableScan pushed-down filters 时启用。 +- 如果 child 是 `PhysicalFilter` 或其他节点，第一版不尝试穿透或重写。 +- 构造 `MSLMTopNDescription`。 +- 保留上层 TopN executor，不修改最终 DAG 语义。 + +### Step 3: ORDER BY ColumnRef 对齐和 Stage 1 列集合 + +- 要求每个 `ByItem.expr` 都是 direct `ColumnRef`。 +- 按 `ColumnRef` index -> `table_scan.getColumns()[index].id` -> `ColumnDefine` 的顺序对齐 order-by columns。 +- 禁用无法按 column id 映射到 storage column 的 order-by columns。 +- 禁用 generated column、virtual `_tidb_tid`、需要 TableScan 后 extra cast 的 order-by columns。 +- 校验 order-by column 类型属于第一版白名单。 +- 将 order-by columns 按 ColumnID 合并到 MSLM Stage 1 column set。 +- final rest columns 按 ColumnID 排除 Stage 1 columns，避免 order-by columns 重复读取。 +- 使用 Stage 1 storage column names 构造 storage-side order-by descriptor。 + +### Step 4: Typed Sort Key Infrastructure + +- 实现 fixed-size typed `SortKeyField` / `OwnedSortKey`。 +- 实现 `SortKeyColumnDesc`，保存 column position、direction、type kind、extractor 和 comparator。 +- 支持 integer numeric、decimal、date、datetime、float。 +- Nullable wrapper 统一由 extractor/comparator 处理，nested type 必须在白名单内。 +- integer numeric、date、datetime 直接比较底层整数值。 +- decimal 直接比较底层 decimal value。 +- float comparator 使用现有 `CompareHelper` 的 NaN 语义。 +- 不支持 String/FixedString/collation，不使用 order-preserving bytes key，不直接使用通用 `Field` 作为 heap key。 + +### Step 5: RunningLocalTopN 状态和算法 + +- 实现 per-stream bounded TopK heap，heap top 是当前 worst candidate。 +- heap entry 只保存 owned key、stream/block sequence 和 current-block row index。 +- 每个 block 开始时预先收集 order-by column pointers，per-row compare 不从 `Block` 查找 column。 +- 支持 row-to-owned lazy compare，只在 row 需要进入 heap 时 materialize owned key。 +- 支持 ASC/DESC、NULL ordering 和 TiFlash 现有 comparator 语义。 +- 生成当前 block 的 `topn_candidate_filter`。 +- 支持 current block eviction，将被当前 block 后续更优 row 挤出的 candidate filter 清掉。 +- 历史 block entry eviction 不做回撤。 + +### Step 6: 接入 MSLM final rest 读取 + +- 在 residual filter 之后执行 running local TopN。 +- 将 `residual_filter` 和 `topn_candidate_filter` 合并。 +- 使用 combined filter 读取 final rest columns。 +- 当 `residual_filter && topn_candidate_filter` 全过滤时，skip final rest stream。 +- 当 `residual_filter && topn_candidate_filter` 全通过时，final rest stream 按 Stage 0 filter 直接读取。 +- 当 `residual_filter && topn_candidate_filter` 部分通过时，compose 回 Stage 0 原始坐标并读取 final rest columns。 + +### Step 7: Runtime Stats 和可观测性 + +- 保留 logical filter counters：`pushed_filter.input_rows/selected_rows/filtered_rows`、`residual_filter.input_rows/selected_rows/filtered_rows`。 +- 增加 `final_rest_input_rows` internal counter，表示实际需要读取 final rest columns 的逻辑行数。 +- 增加 running TopN counters：`input_rows`、`selected_rows`、`bypass_rows`、`filtered_rows`、`heap_size_sum`。 +- 增加 runtime one-way adaptive disable 统计，表示有多少 streams 因 TopN candidate ratio 过高而禁用后续 TopN filter。 +- 在 running local TopN 启用时，用 `final_rest_input_rows` overwrite residual Selection executor 的 `actRows`，让 `EXPLAIN ANALYZE` 直接展示 final rest columns materialized rows。 +- 未启用 running local TopN 时，Selection executor 的 `actRows` 仍展示 residual Selection 后的逻辑行数。 +- TableScan executor 的 `actRows` 仍然使用 pushed filter output rows。 +- 增加 `TiFlashScanContext.multi_stage_late_materialization`，不要复用原 single-stage LM 的 `lm_*` / `dmfile_lm_filter_*` 字段。 +- 增加 debug log，输出 residual passed rows、TopN selected/bypass/filtered rows、heap size、adaptive disable 状态。 + +## 测试方案 + +### 功能测试 + +- `ORDER BY c ASC LIMIT k`。 +- `ORDER BY c DESC LIMIT k`。 +- `ORDER BY c ASC LIMIT offset, count`。 +- `WHERE + ORDER BY + LIMIT`，其中 WHERE 同时包含 Stage 0 pushed filter 和 Stage 1 residual filter。 +- ORDER BY column 不在最终 projection 中。 +- ORDER BY column 同时属于 residual filter columns。 +- ORDER BY column 同时属于最终输出列，验证 final rest columns 不重复读取该列。 +- 多 block 数据，确保 running local TopN 可以跨 block 收紧 threshold。 +- 多 stream 并发，确保上层 TopN 保证最终结果。 + +### 正确性测试 + +- 和未启用优化的执行结果比较。 +- LIMIT/OFFSET 场景比较。 +- NULL ordering 场景比较。 +- ties 场景比较，验证结果满足 SQL order semantics。 +- Stage 0 全过滤。 +- residual filter 全过滤。 +- TopN filter 全过滤。 +- TopN filter 全通过。 +- ORDER BY plain column 按 column id 对齐到 storage column，避免依赖 child schema name。 +- ORDER BY column 需要 TableScan 后 extra cast 时禁用优化。 +- 没有 residual Selection 时禁用优化。 +- 没有 TableScan pushed-down filters 时禁用优化。 + +### 回归测试 + +- MSLM 原有测试必须全部通过。 +- TopN 上方的 projection 仍然正确。 +- TopN-enhanced MSLM 启用时，EXPLAIN ANALYZE 中 residual Selection `actRows` 展示 `final_rest_input_rows`，即 final rest columns materialized rows。 +- 禁用条件下必须走原路径。 + +## 性能评估 + +重点 benchmark 场景： + +- 宽表，payload columns 多且冷数据。 +- `LIMIT` 小，`topk` 远小于 block rows。 +- `ORDER BY` columns 少且窄。 +- 数据分布有利于 running local TopN threshold 快速收紧。 + +需要对比： + +- MSLM disabled。 +- MSLM enabled but without running local TopN。 +- MSLM enabled with running local TopN。 + +关键指标： + +- Query latency。 +- final rest columns read rows。 +- DMFile read bytes。 +- CPU time。 +- memory usage。 +- `final_rest_input_rows`。 +- `running_topn.selected_rows/bypass_rows/filtered_rows`。 + +## 风险 + +### Comparator 语义不一致 + +如果 storage 内 running local TopN 的 comparator 和上层 TopN 不一致，可能错误过滤掉本应保留的 rows。 + +缓解： + +- MVP 只支持 plain columns。 +- 复用现有排序比较逻辑。 +- 对 ASC/DESC、NULL、float NaN、decimal/date/datetime 做正确性单测。 +- 对 String/FixedString/collation 做禁用路径单测。 + +### ORDER BY 列对齐错误 + +如果直接用 TopN child schema name 匹配 DeltaMerge storage column name，可能因为 TableScan projection、特殊列名或 cast-after-TS 造成错配。 + +缓解： + +- 只支持 direct `ColumnRef`。 +- 使用 `ColumnRef` index 映射到 `table_scan.getColumns()[index].id`。 +- 后续所有 Stage 1 column union、final rest column subtraction 和 storage-side sort description 都基于 `ColumnID` / `ColumnDefine` 构造。 +- 禁用无法按 column id 映射到 storage column 的场景。 + +### Selection 顺序错误 + +如果 TopN filter 在某些 Selection predicates 执行之前生效，可能过滤掉最终应该进入 TopN 的 rows。 + +缓解： + +- 只在 TopN 之前的 Selection 已全部进入 Stage 0/Stage 1 时启用。 +- 无法确认时禁用。 + +### 收益不足或性能回退 + +如果 `topk` 大、order-by columns 宽、ties 很多，running local TopN 可能减少不了 final rest 读取，还增加 Stage 1 读取和 heap 维护成本。 + +缓解： + +- 使用保守启发式启用。 +- `topk` 和 order-by column count 设置上限。 +- 使用 runtime one-way adaptive disable，在 candidate ratio 过高时禁用后续 TopN filter。 +- adaptive disable 不能避免 Stage 1 已经额外读取的 order-by columns，因此静态 gating 仍然是主要防线。 +- 对 ties 多的场景 benchmark。 + +### Heap 持有历史 Block 导致内存累积 + +如果 heap entry 设计成 `{Block, row_id}`，或者持有历史 column memory 引用，会导致已经处理过的 blocks 无法释放，内存随扫描推进累积。 + +缓解： + +- Heap entry 只保存 owned sort key。 +- Heap entry 不持有 `Block`、column 引用或历史 row payload。 +- Row locator 只用于 current block 内 eviction。 +- 历史 block entry 被 eviction 时只释放 key，不访问历史 block。 +- 对 variable-length order-by keys 使用禁用策略或 memory guard。 + +### Streaming Superset 过大 + +Running local TopN 不能撤回之前已经输出的 rows。极端数据分布下，每个 block 都可能输出一批候选 rows。 + +缓解： + +- 上层 TopN 保留，保证正确性。 +- 通过 stats 观察 `final_rest_input_rows` 和 `running_topn.selected_rows/bypass_rows/filtered_rows`。 +- 后续考虑更强的 candidate row id materialization。 + +## 备选方案 + +### Block-local TopN + +每个 block 独立做 TopK，不维护跨 block state。 + +优点： + +- 实现更简单。 +- 正确性证明直接。 +- 不需要处理 previous block candidate eviction。 + +缺点： + +- 最多输出 `K * block_count` 个候选 rows。 +- 无法利用前面 blocks 已经得到的 threshold 剪枝后续 blocks。 + +Running local TopN 是 block-local TopN 的增强版，复杂度增加有限，但剪枝能力更好。 + +### Candidate Row ID 回读 + +先读取 selection/order-by columns，得到全局或局部 TopN row ids，再按 row ids 回读 payload columns。 + +优点： + +- 最接近理想 late materialization。 +- 可以只 materialize 很少 rows。 + +缺点： + +- 需要 DeltaMerge storage 支持按 candidate row id 高效回读 columns。 +- 当前不具备该能力。 +- 设计和实现复杂度明显更高。 + +本设计选择 running local TopN 作为短期方案。 + +## 未解决问题 + +- `topk` 阈值应该如何设置。 +- ties 场景是否需要为了结果稳定性保留所有 equal-worst rows。 +- 第一版是否禁用 variable-length order-by columns，还是实现 per-stream memory guard。 +- 是否需要支持 ORDER BY expression，以及 expression action 如何复用。 +- 是否需要在 `EXPLAIN ANALYZE` 中展示 scan context 之外的更细粒度 MSLM TopN 信息。 +- 是否需要在 cost model 中引入列宽和 TopN selectivity。