From 3938abeb9d829d5f7e7754f41c06bd2bac633f37 Mon Sep 17 00:00:00 2001 From: 3pointer Date: Thu, 13 Aug 2026 13:35:27 +0800 Subject: [PATCH] tici: preview port feature/fts onto master --- .gitmodules | 3 + contrib/CMakeLists.txt | 2 + contrib/client-c | 2 +- contrib/kvproto | 2 +- contrib/tici | 1 + contrib/tici-search-lib/CMakeLists.txt | 39 ++ dbms/CMakeLists.txt | 5 + dbms/src/Common/TiFlashMetrics.h | 3 +- dbms/src/Debug/dbgQueryExecutor.cpp | 26 +- dbms/src/Flash/BatchCoprocessorHandler.cpp | 4 +- dbms/src/Flash/Coprocessor/DAGContext.cpp | 2 + dbms/src/Flash/Coprocessor/DAGContext.h | 4 + dbms/src/Flash/Coprocessor/DAGDriver.cpp | 12 + .../Coprocessor/DAGExpressionAnalyzer.cpp | 26 + .../Flash/Coprocessor/DAGExpressionAnalyzer.h | 6 + .../Flash/Coprocessor/GenSchemaAndColumn.cpp | 13 + .../Flash/Coprocessor/GenSchemaAndColumn.h | 1 + dbms/src/Flash/Coprocessor/RemoteRequest.cpp | 103 ++++ dbms/src/Flash/Coprocessor/RemoteRequest.h | 12 + dbms/src/Flash/Coprocessor/ShardInfo.h | 151 ++++++ .../Coprocessor/StorageTantivyInterpreter.cpp | 213 +++++++++ .../Coprocessor/StorageTantivyInterpreter.h | 71 +++ dbms/src/Flash/Coprocessor/TiCIScan.cpp | 49 ++ dbms/src/Flash/Coprocessor/TiCIScan.h | 76 +++ .../Coprocessor/collectOutputFieldTypes.cpp | 19 + dbms/src/Flash/CoprocessorHandler.cpp | 4 +- dbms/src/Flash/EstimateTiCICountHandler.cpp | 133 ++++++ dbms/src/Flash/EstimateTiCICountHandler.h | 40 ++ dbms/src/Flash/FlashService.cpp | 15 + dbms/src/Flash/FlashService.h | 4 + dbms/src/Flash/Mpp/MPPTask.cpp | 1 + dbms/src/Flash/Planner/PhysicalPlan.cpp | 28 ++ dbms/src/Flash/Planner/PhysicalPlan.h | 1 + dbms/src/Flash/Planner/PhysicalPlanNode.h | 2 + dbms/src/Flash/Planner/PlanType.h | 1 + .../Flash/Planner/Plans/PhysicalAggregation.h | 8 + dbms/src/Flash/Planner/Plans/PhysicalBinary.h | 4 +- dbms/src/Flash/Planner/Plans/PhysicalLeaf.h | 6 +- .../Flash/Planner/Plans/PhysicalTiCIScan.cpp | 151 ++++++ .../Flash/Planner/Plans/PhysicalTiCIScan.h | 73 +++ dbms/src/Flash/Planner/Plans/PhysicalUnary.h | 8 +- dbms/src/Flash/Planner/optimize.cpp | 43 +- .../src/Flash/Statistics/CommonExecutorImpl.h | 12 + .../ExecutorStatisticsCollector.cpp | 4 +- .../Flash/Statistics/traverseExecutors.cpp | 3 +- dbms/src/Operators/TantivyReaderSourceOp.cpp | 93 ++++ dbms/src/Operators/TantivyReaderSourceOp.h | 64 +++ dbms/src/Server/CMakeLists.txt | 8 + dbms/src/Server/MetricsPrometheus.cpp | 47 +- dbms/src/Server/MetricsPrometheus.h | 2 - dbms/src/Server/Server.cpp | 14 + dbms/src/Storages/KVStore/TMTContext.cpp | 19 + dbms/src/Storages/StorageTantivy.cpp | 179 +++++++ dbms/src/Storages/StorageTantivy.h | 90 ++++ .../src/Storages/Tantivy/TantivyInputStream.h | 444 ++++++++++++++++++ dbms/src/Storages/Tantivy/TiCIReadTaskPool.h | 178 +++++++ dbms/src/Storages/Tantivy/TiCIRequestUtils.h | 256 ++++++++++ .../Storages/tests/gtest_filter_parser.cpp | 2 +- .../tests/gtests_parse_push_down_filter.cpp | 2 +- dbms/src/TestUtils/MPPTaskTestUtils.cpp | 2 +- rust-toolchain.toml | 5 +- tests/docker/util.sh | 62 +++ tests/fullstack-test/run.sh | 4 +- tests/fullstack-test2/run.sh | 6 +- tests/tidb-ci/run.sh | 10 +- 65 files changed, 2833 insertions(+), 40 deletions(-) create mode 160000 contrib/tici create mode 100644 contrib/tici-search-lib/CMakeLists.txt create mode 100644 dbms/src/Flash/Coprocessor/ShardInfo.h create mode 100644 dbms/src/Flash/Coprocessor/StorageTantivyInterpreter.cpp create mode 100644 dbms/src/Flash/Coprocessor/StorageTantivyInterpreter.h create mode 100644 dbms/src/Flash/Coprocessor/TiCIScan.cpp create mode 100644 dbms/src/Flash/Coprocessor/TiCIScan.h create mode 100644 dbms/src/Flash/EstimateTiCICountHandler.cpp create mode 100644 dbms/src/Flash/EstimateTiCICountHandler.h create mode 100644 dbms/src/Flash/Planner/Plans/PhysicalTiCIScan.cpp create mode 100644 dbms/src/Flash/Planner/Plans/PhysicalTiCIScan.h create mode 100644 dbms/src/Operators/TantivyReaderSourceOp.cpp create mode 100644 dbms/src/Operators/TantivyReaderSourceOp.h create mode 100644 dbms/src/Storages/StorageTantivy.cpp create mode 100644 dbms/src/Storages/StorageTantivy.h create mode 100644 dbms/src/Storages/Tantivy/TantivyInputStream.h create mode 100644 dbms/src/Storages/Tantivy/TiCIReadTaskPool.h create mode 100644 dbms/src/Storages/Tantivy/TiCIRequestUtils.h mode change 120000 => 100644 rust-toolchain.toml diff --git a/.gitmodules b/.gitmodules index 132b443afd6..45ca53d0002 100644 --- a/.gitmodules +++ b/.gitmodules @@ -162,3 +162,6 @@ [submodule "contrib/highfive"] path = contrib/highfive url = https://github.com/BlueBrain/HighFive +[submodule "contrib/tici"] + path = contrib/tici + url = git@github.com:pingcap-inc/tici.git diff --git a/contrib/CMakeLists.txt b/contrib/CMakeLists.txt index b9d21c95f6f..f9ab57e3c80 100644 --- a/contrib/CMakeLists.txt +++ b/contrib/CMakeLists.txt @@ -207,3 +207,5 @@ if (ENABLE_TESTS AND NOT CMAKE_BUILD_TYPE_UC STREQUAL "DEBUG") add_subdirectory(hdf5-cmake) add_subdirectory(highfive-cmake) endif () + +add_subdirectory(tici-search-lib) \ No newline at end of file diff --git a/contrib/client-c b/contrib/client-c index 55aede24990..2214b0bd30c 160000 --- a/contrib/client-c +++ b/contrib/client-c @@ -1 +1 @@ -Subproject commit 55aede24990ec5f92a1d74768d65a11eacc5f7d5 +Subproject commit 2214b0bd30cffcf832c9c854fcaee2b430a2103d diff --git a/contrib/kvproto b/contrib/kvproto index 2bc1b358068..811b5283cd8 160000 --- a/contrib/kvproto +++ b/contrib/kvproto @@ -1 +1 @@ -Subproject commit 2bc1b358068b5e1c73feed4af966c38c33697bcd +Subproject commit 811b5283cd8a7f8b15be83eb483cb07f1db68639 diff --git a/contrib/tici b/contrib/tici new file mode 160000 index 00000000000..6489830cf3b --- /dev/null +++ b/contrib/tici @@ -0,0 +1 @@ +Subproject commit 6489830cf3b33aacc6e4e5f8285e0d1bf4f790ca diff --git a/contrib/tici-search-lib/CMakeLists.txt b/contrib/tici-search-lib/CMakeLists.txt new file mode 100644 index 00000000000..eddf1e36d2d --- /dev/null +++ b/contrib/tici-search-lib/CMakeLists.txt @@ -0,0 +1,39 @@ +set(TICI_PROJECT_DIR "${CMAKE_SOURCE_DIR}/contrib/tici") +set(TICI_LIB_NAME "tici_search_lib") +set(TICI_LIB "${CMAKE_CURRENT_BINARY_DIR}/release/${CMAKE_STATIC_LIBRARY_PREFIX}${TICI_LIB_NAME}${CMAKE_STATIC_LIBRARY_SUFFIX}") + +file(GLOB LIB_SOURCE_FILES "${TICI_PROJECT_DIR}/src/*") +file(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/cxxbridge) + +add_custom_command( + OUTPUT ${TICI_LIB} + COMMAND cargo build -p tici-search-lib --release --target-dir ${CMAKE_CURRENT_BINARY_DIR} --manifest-path ${TICI_PROJECT_DIR}/Cargo.toml + WORKING_DIRECTORY ${TICI_PROJECT_DIR} + DEPENDS ${LIB_SOURCE_FILES} + COMMENT "Build Rust lib"${CMAKE_CURRENT_BINARY_DIR} +) + +add_custom_target(rustbuild ALL DEPENDS ${TICI_LIB}) +add_library(tici_search_lib_static STATIC IMPORTED GLOBAL) +set_target_properties(tici_search_lib_static PROPERTIES + IMPORTED_LOCATION ${TICI_LIB} +) + +add_dependencies(tici_search_lib_static rustbuild) +target_include_directories(tici_search_lib_static INTERFACE + ${CMAKE_CURRENT_BINARY_DIR}/cxxbridge) + +add_library(tici_search_lib SHARED "${TiFlash_SOURCE_DIR}/libs/libclara-cmake/dummy.cpp") +target_compile_options(tici_search_lib PRIVATE -pthread) +target_link_options(tici_search_lib PRIVATE -pthread) +target_link_libraries(tici_search_lib PRIVATE "$") +if(APPLE) + target_link_libraries(tici_search_lib PRIVATE + "-framework Security" + "-framework CoreFoundation" + "-framework IOKit" + ) +endif() + +target_include_directories(tici_search_lib INTERFACE + ${CMAKE_CURRENT_BINARY_DIR}/cxxbridge) diff --git a/dbms/CMakeLists.txt b/dbms/CMakeLists.txt index a3eec57fa9d..930f0fc94cb 100644 --- a/dbms/CMakeLists.txt +++ b/dbms/CMakeLists.txt @@ -227,6 +227,7 @@ target_link_libraries (dbms absl::synchronization tiflash_contrib::aws_s3 tiflash_vector_search + tici_search_lib etcdpb tiflash_parsers @@ -380,6 +381,10 @@ if (ENABLE_TESTS) DESTINATION ".") endif () + install (TARGETS tici_search_lib + COMPONENT tiflash-gtest + DESTINATION ".") + if (ENABLE_CLARA AND USE_INTERNAL_LIBCLARA) install (TARGETS clara_shared COMPONENT tiflash-gtest diff --git a/dbms/src/Common/TiFlashMetrics.h b/dbms/src/Common/TiFlashMetrics.h index 1d709eb7f30..1b242180f42 100644 --- a/dbms/src/Common/TiFlashMetrics.h +++ b/dbms/src/Common/TiFlashMetrics.h @@ -103,7 +103,8 @@ static_assert(RAFT_REGION_BIG_WRITE_THRES * 4 < RAFT_REGION_BIG_WRITE_MAX, "Inva F(type_cte_source, {"type", "cte_source"}), \ F(type_window, {"type", "window"}), \ F(type_window_sort, {"type", "window_sort"}), \ - F(type_expand, {"type", "expand"})) \ + F(type_expand, {"type", "expand"}), \ + F(type_tici, {"type", "tici_scan"})) \ M(tiflash_memory_exceed_quota_count, "Total number of cases where memory exceeds quota", Counter) \ M(tiflash_coprocessor_request_duration_seconds, \ "Bucketed histogram of request duration", \ diff --git a/dbms/src/Debug/dbgQueryExecutor.cpp b/dbms/src/Debug/dbgQueryExecutor.cpp index d677e7ad09f..6a4eb4c09ec 100644 --- a/dbms/src/Debug/dbgQueryExecutor.cpp +++ b/dbms/src/Debug/dbgQueryExecutor.cpp @@ -403,8 +403,17 @@ tipb::SelectResponse executeDAGRequest( region_id, RegionInfo(region_id, region_version, region_conf_version, std::move(key_ranges), nullptr)); - DAGContext - dag_context(dag_request, std::move(tables_regions_info), NullspaceID, "", DAGRequestKind::Cop, "", 0, "", log); + DAGContext dag_context( + dag_request, + std::move(tables_regions_info), + QueryShardInfos(), + NullspaceID, + "", + DAGRequestKind::Cop, + "", + 0, + "", + log); context.setDAGContext(&dag_context); DAGDriver driver(context, start_ts, DEFAULT_UNSPECIFIED_SCHEMA_VERSION, &dag_response, true); @@ -438,8 +447,17 @@ bool runAndCompareDagReq( region_id, RegionInfo(region_id, region->version(), region->confVer(), std::move(key_ranges), nullptr)); - DAGContext - dag_context(dag_request, std::move(tables_regions_info), NullspaceID, "", DAGRequestKind::Cop, "", 0, "", log); + DAGContext dag_context( + dag_request, + std::move(tables_regions_info), + QueryShardInfos(), + NullspaceID, + "", + DAGRequestKind::Cop, + "", + 0, + "", + log); context.setDAGContext(&dag_context); DAGDriver driver(context, properties.start_ts, DEFAULT_UNSPECIFIED_SCHEMA_VERSION, &dag_response, true); diff --git a/dbms/src/Flash/BatchCoprocessorHandler.cpp b/dbms/src/Flash/BatchCoprocessorHandler.cpp index c1a5c296a3e..c244572e2bd 100644 --- a/dbms/src/Flash/BatchCoprocessorHandler.cpp +++ b/dbms/src/Flash/BatchCoprocessorHandler.cpp @@ -55,8 +55,9 @@ grpc::Status BatchCoprocessorHandler::execute() try { + bool is_tici = cop_request->table_shard_infos_size() != 0; RUNTIME_CHECK_MSG( - !cop_context.db_context.getSharedContextDisagg()->isDisaggregatedComputeMode(), + !cop_context.db_context.getSharedContextDisagg()->isDisaggregatedComputeMode() || is_tici, "cannot run cop or batchCop request on tiflash_compute node"); switch (cop_request->tp()) @@ -96,6 +97,7 @@ grpc::Status BatchCoprocessorHandler::execute() DAGContext dag_context( dag_request, std::move(tables_regions_info), + QueryShardInfos::create(cop_request->table_shard_infos()), RequestUtils::deriveKeyspaceID(cop_request->context()), cop_context.db_context.getClientInfo().current_address.toString(), DAGRequestKind::BatchCop, diff --git a/dbms/src/Flash/Coprocessor/DAGContext.cpp b/dbms/src/Flash/Coprocessor/DAGContext.cpp index 84b81c383f2..4db3013821e 100644 --- a/dbms/src/Flash/Coprocessor/DAGContext.cpp +++ b/dbms/src/Flash/Coprocessor/DAGContext.cpp @@ -49,6 +49,7 @@ bool strictSqlMode(UInt64 sql_mode) DAGContext::DAGContext( tipb::DAGRequest & dag_request_, TablesRegionsInfo && tables_regions_info_, + QueryShardInfos && query_shard_infos_, KeyspaceID keyspace_id_, const String & tidb_host_, DAGRequestKind kind_, @@ -65,6 +66,7 @@ DAGContext::DAGContext( , kind(kind_) , is_root_mpp_task(false) , tables_regions_info(std::move(tables_regions_info_)) + , query_shard_infos(std::move(query_shard_infos_)) , log(std::move(log_)) , operator_spill_contexts(std::make_shared()) , flags(dag_request->flags()) diff --git a/dbms/src/Flash/Coprocessor/DAGContext.h b/dbms/src/Flash/Coprocessor/DAGContext.h index 19216cc4563..8b96a6abf68 100644 --- a/dbms/src/Flash/Coprocessor/DAGContext.h +++ b/dbms/src/Flash/Coprocessor/DAGContext.h @@ -36,6 +36,7 @@ #include #include #include +#include #include #include #include @@ -170,6 +171,7 @@ class DAGContext DAGContext( tipb::DAGRequest & dag_request_, TablesRegionsInfo && tables_regions_info_, + QueryShardInfos && query_shard_infos_, KeyspaceID keyspace_id_, const String & tidb_host_, DAGRequestKind cop_kind_, @@ -473,6 +475,8 @@ class DAGContext std::unique_ptr> bypass_lock_ts; // part of regions_for_local_read + regions_for_remote_read, only used for batch-cop RegionInfoList retry_regions; + QueryShardInfos query_shard_infos; + ShardInfoList retry_shards; LoggerPtr log; diff --git a/dbms/src/Flash/Coprocessor/DAGDriver.cpp b/dbms/src/Flash/Coprocessor/DAGDriver.cpp index daa760cc3c1..61514bb422c 100644 --- a/dbms/src/Flash/Coprocessor/DAGDriver.cpp +++ b/dbms/src/Flash/Coprocessor/DAGDriver.cpp @@ -221,6 +221,18 @@ try batch_cop_writer->Write(response); } + if (!dag_context.retry_shards.empty()) + { + coprocessor::BatchResponse response; + for (const auto & shard : dag_context.retry_shards) + { + auto * retry_shard = response.add_retry_shards(); + retry_shard->set_shard_id(shard.shard_id); + retry_shard->set_shard_epoch(shard.shard_epoch); + } + batch_cop_writer->Write(response); + } + auto streaming_writer = std::make_shared(batch_cop_writer); TiDB::TiDBCollators collators; auto response_writer = std::make_unique>( diff --git a/dbms/src/Flash/Coprocessor/DAGExpressionAnalyzer.cpp b/dbms/src/Flash/Coprocessor/DAGExpressionAnalyzer.cpp index 1d146d16fbc..5dacdd81eb8 100644 --- a/dbms/src/Flash/Coprocessor/DAGExpressionAnalyzer.cpp +++ b/dbms/src/Flash/Coprocessor/DAGExpressionAnalyzer.cpp @@ -1308,6 +1308,32 @@ bool DAGExpressionAnalyzer::appendExtraCastsAfterTS( return true; } +bool DAGExpressionAnalyzer::appendExtraCastsAfterTiCI( + ExpressionActionsChain & chain, + const std::vector & may_need_add_cast_column, + const TiCIScan & tici_scan) +{ + auto & step = initAndGetLastStep(chain); + auto & actions = step.actions; + + auto [has_cast, casted_columns] + = buildExtraCastsAfterTS(actions, may_need_add_cast_column, tici_scan.getReturnColumns()); + + if (!has_cast) + return false; + + // Add a projection to replace the original columns with the casted columns. + NamesWithAliases project_cols; + for (size_t i = 0; i < may_need_add_cast_column.size(); ++i) + project_cols.emplace_back(casted_columns[i], source_columns[i].name); + actions->add(ExpressionAction::project(project_cols)); + + for (auto & col : source_columns) + step.required_output.push_back(col.name); + + return true; +} + String DAGExpressionAnalyzer::appendDurationCast( const String & fsp_expr, const String & dur_expr, diff --git a/dbms/src/Flash/Coprocessor/DAGExpressionAnalyzer.h b/dbms/src/Flash/Coprocessor/DAGExpressionAnalyzer.h index 22251ae853c..5a9e3a424d5 100644 --- a/dbms/src/Flash/Coprocessor/DAGExpressionAnalyzer.h +++ b/dbms/src/Flash/Coprocessor/DAGExpressionAnalyzer.h @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -114,6 +115,11 @@ class DAGExpressionAnalyzer : private boost::noncopyable const std::vector & may_need_add_cast_column, const TiDBTableScan & table_scan); + bool appendExtraCastsAfterTiCI( + ExpressionActionsChain & chain, + const std::vector & may_need_add_cast_column, + const TiCIScan & tici_scan); + /// return true if some actions is needed bool appendJoinKeyAndJoinFilters( ExpressionActionsChain & chain, diff --git a/dbms/src/Flash/Coprocessor/GenSchemaAndColumn.cpp b/dbms/src/Flash/Coprocessor/GenSchemaAndColumn.cpp index 206fc0a01ff..0d4ef5ee6ba 100644 --- a/dbms/src/Flash/Coprocessor/GenSchemaAndColumn.cpp +++ b/dbms/src/Flash/Coprocessor/GenSchemaAndColumn.cpp @@ -100,6 +100,19 @@ NamesAndTypes genNamesAndTypes(const TiDBTableScan & table_scan, const StringRef return genNamesAndTypes(table_scan.getColumns(), column_prefix); } +NamesAndTypes genNamesAndTypesForTiCI(const TiDB::ColumnInfos & column_infos, const StringRef & column_prefix) +{ + NamesAndTypes names_and_types; + names_and_types.reserve(column_infos.size()); + for (const auto & column_info : column_infos) + { + names_and_types.emplace_back( + fmt::format("{}_{}", column_prefix, column_info.id), + getDataTypeByColumnInfoForComputingLayer(column_info)); + } + return names_and_types; +} + std::tuple>> genColumnDefinesForDisaggregatedRead( const TiDBTableScan & table_scan) { diff --git a/dbms/src/Flash/Coprocessor/GenSchemaAndColumn.h b/dbms/src/Flash/Coprocessor/GenSchemaAndColumn.h index 1eb87891ba9..88f5ec3da99 100644 --- a/dbms/src/Flash/Coprocessor/GenSchemaAndColumn.h +++ b/dbms/src/Flash/Coprocessor/GenSchemaAndColumn.h @@ -31,6 +31,7 @@ String genNameForCTESource(Int32 cte_id, Int32 col_index); NamesAndTypes genNamesAndTypes(const TiDBTableScan & table_scan, const StringRef & column_prefix); NamesAndTypes genNamesAndTypes(const TiDB::ColumnInfos & column_infos, const StringRef & column_prefix); +NamesAndTypes genNamesAndTypesForTiCI(const TiDB::ColumnInfos & column_infos, const StringRef & column_prefix); ColumnsWithTypeAndName getColumnWithTypeAndName(const NamesAndTypes & names_and_types); NamesAndTypes toNamesAndTypes(const DAGSchema & dag_schema); diff --git a/dbms/src/Flash/Coprocessor/RemoteRequest.cpp b/dbms/src/Flash/Coprocessor/RemoteRequest.cpp index a056c78cbba..9beddc22ce0 100644 --- a/dbms/src/Flash/Coprocessor/RemoteRequest.cpp +++ b/dbms/src/Flash/Coprocessor/RemoteRequest.cpp @@ -16,7 +16,10 @@ #include #include #include +#include #include +#include +#include #include #include @@ -107,6 +110,80 @@ RemoteRequest RemoteRequest::build( return {std::move(dag_req), std::move(schema), std::move(key_ranges), connection_id, connection_alias}; } +RemoteRequest RemoteRequest::build( + const ShardInfoList & shard_infos, + DAGContext & dag_context, + const TiCIScan & tici_scan, + UInt64 connection_id, + const String & connection_alias, + const LoggerPtr & log) +{ + LOG_INFO(log, "{}", printShards(shard_infos, tici_scan.getTableId(), tici_scan.getIndexId())); + DAGSchema schema; + tipb::DAGRequest dag_req; + { + tipb::Executor * tici_scan_exec; + NamesAndTypes names_and_types; + if (tici_scan.isCount()) + { + auto * root_exec = dag_req.mutable_root_executor(); + root_exec->set_tp(tipb::ExecType::TypeAggregation); + root_exec->set_executor_id(tici_scan.getCountAggExecutorId()); + auto * new_aggregation = root_exec->mutable_aggregation(); + tipb::Expr * agg_func = new_aggregation->add_agg_func(); + agg_func->set_tp(tipb::ExprType::Count); + auto * ft = agg_func->mutable_field_type(); + // "count" always returns a NOT NULL INT64 column + ft->set_tp(TiDB::TypeLongLong); + ft->set_flag(TiDB::ColumnFlagNotNull); + tici_scan_exec = new_aggregation->mutable_child(); + + names_and_types = tici_scan.getNamesAndTypes(); + TiDB::ColumnInfo ci; + // "count" always returns a NOT NULL INT64 column + ci.tp = TiDB::TypeLongLong; + ci.setNotNullFlag(); + schema.emplace_back(std::make_pair(names_and_types[0].name, std::move(ci))); + dag_req.add_output_offsets(0); + } + else + { + const auto & return_columns = tici_scan.getReturnColumns(); + tici_scan_exec = dag_req.mutable_root_executor(); + names_and_types = genNamesAndTypesForTiCI(return_columns, "column"); + for (size_t i = 0; i < return_columns.size(); ++i) + { + const auto & col = return_columns[i]; + schema.emplace_back(std::make_pair(names_and_types[i].name, col)); + dag_req.add_output_offsets(i); + } + } + + tici_scan_exec->set_tp(tipb::ExecType::TypeIndexScan); + tici_scan_exec->set_executor_id(tici_scan.getTiCIScan()->executor_id()); + auto * mutable_tici_scan = tici_scan_exec->mutable_idx_scan(); + tici_scan.constructTiCIScanForRemoteRead(mutable_tici_scan); + + dag_req.set_encode_type(tipb::EncodeType::TypeCHBlock); + dag_req.set_force_encode_type(true); + } + /// do not collect execution summaries because in this case because the execution summaries + /// will be collected by CoprocessorBlockInputStream. + /// Otherwise rows in execution summary of table scan will be double. + dag_req.set_collect_execution_summaries(false); + dag_req.set_flags(dag_context.getFlags()); + dag_req.set_sql_mode(dag_context.getSQLMode()); + dag_req.set_div_precision_increment(dag_context.getDivPrecisionIncrement()); + const auto & original_dag_req = *dag_context.dag_request; + if (original_dag_req.has_time_zone_name() && !original_dag_req.time_zone_name().empty()) + dag_req.set_time_zone_name(original_dag_req.time_zone_name()); + if (original_dag_req.has_time_zone_offset()) + dag_req.set_time_zone_offset(original_dag_req.time_zone_offset()); + + std::vector key_ranges = buildKeyRanges(shard_infos); + return {std::move(dag_req), std::move(schema), std::move(key_ranges), connection_id, connection_alias}; +} + std::vector RemoteRequest::buildKeyRanges(const RegionRetryList & retry_regions) { std::vector key_ranges; @@ -119,6 +196,18 @@ std::vector RemoteRequest::buildKeyRanges(const return key_ranges; } +std::vector RemoteRequest::buildKeyRanges(const ShardInfoList & retry_shards) +{ + std::vector key_ranges; + for (const auto & shard : retry_shards) + { + for (const auto & range : shard.key_ranges) + key_ranges.emplace_back(range.start(), range.end()); + } + sort(key_ranges.begin(), key_ranges.end()); + return key_ranges; +} + std::string RemoteRequest::printRetryRegions(const RegionRetryList & retry_regions, TableID table_id) { FmtBuffer buffer; @@ -132,4 +221,18 @@ std::string RemoteRequest::printRetryRegions(const RegionRetryList & retry_regio return buffer.toString(); } +std::string RemoteRequest::printShards(const ShardInfoList & shards, Int64 table_id, Int64 index_id) +{ + FmtBuffer buffer; + buffer.fmtAppend("Start to build remote request for {} shards (", shards.size()); + buffer.joinStr( + shards.cbegin(), + shards.cend(), + [](const auto & shard, FmtBuffer & fb) { fb.fmtAppend("{}", shard.getID()); }, + ","); + buffer.fmtAppend(") for table {} and index {}", table_id, index_id); + return buffer.toString(); +} + + } // namespace DB diff --git a/dbms/src/Flash/Coprocessor/RemoteRequest.h b/dbms/src/Flash/Coprocessor/RemoteRequest.h index b107d7890a0..13639eebd01 100644 --- a/dbms/src/Flash/Coprocessor/RemoteRequest.h +++ b/dbms/src/Flash/Coprocessor/RemoteRequest.h @@ -16,6 +16,8 @@ #include #include +#include +#include #include #include @@ -59,8 +61,18 @@ struct RemoteRequest UInt64 connection_id, const String & connection_alias, const LoggerPtr & log); + + static RemoteRequest build( + const ShardInfoList & shard_infos, + DAGContext & dag_context, + const TiCIScan & tici_scan, + UInt64 connection_id, + const String & connection_alias, + const LoggerPtr & log); static std::vector buildKeyRanges(const RegionRetryList & retry_regions); + static std::vector buildKeyRanges(const ShardInfoList & retry_shards); static std::string printRetryRegions(const RegionRetryList & retry_regions, TableID table_id); + static std::string printShards(const ShardInfoList & shards, Int64 table_id, Int64 index_id); tipb::DAGRequest dag_request; DAGSchema schema; diff --git a/dbms/src/Flash/Coprocessor/ShardInfo.h b/dbms/src/Flash/Coprocessor/ShardInfo.h new file mode 100644 index 00000000000..b06c6201b08 --- /dev/null +++ b/dbms/src/Flash/Coprocessor/ShardInfo.h @@ -0,0 +1,151 @@ +// Copyright 2025 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 +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wunused-parameter" +#include +#include +#pragma GCC diagnostic pop + +namespace DB +{ + +struct ShardInfo +{ + UInt64 shard_id; + UInt64 shard_epoch; + using KeyRanges = google::protobuf::RepeatedPtrField; + KeyRanges key_ranges; + + explicit ShardInfo(const coprocessor::ShardInfo & info) + : shard_id(info.shard_id()) + , shard_epoch(info.shard_epoch()) + , key_ranges(info.ranges()) + {} + + String toString() const + { + FmtBuffer buf; + buf.fmtAppend("ShardID: {}, ShardEpoch: {}, KeyRanges: ", shard_id, shard_epoch); + buf.joinStr( + key_ranges.begin(), + key_ranges.end(), + [](const coprocessor::KeyRange & range, FmtBuffer & fb) { + fb.fmtAppend( + "[{}, {})", + Redact::keyToDebugString(range.start().data(), range.start().size()), + Redact::keyToDebugString(range.end().data(), range.end().size())); + }, + " "); + return buf.toString(); + } + + UInt64 getID() const { return shard_id; } +}; + +using ShardInfoMap = std::unordered_map; +using ShardInfoList = std::vector; + +class TableShardInfos +{ +public: + TableShardInfos() = default; + + static TableShardInfos create(const coprocessor::TableShardInfos & table_shard_infos) + { + TableShardInfos infos; + infos.executor_id = table_shard_infos.executor_id(); + for (const auto & info : table_shard_infos.shard_infos()) + { + ShardInfo shard_info(info); + infos.shard_info_list.push_back(shard_info); + } + return infos; + } + + String toString() const + { + FmtBuffer buf; + buf.fmtAppend("ExecutorID: {}, ShardInfos: [", executor_id); + buf.joinStr( + shard_info_list.begin(), + shard_info_list.end(), + [](const ShardInfo & shard_info, FmtBuffer & fb) { fb.append(shard_info.toString()); }, + " "); + buf.append("]"); + return buf.toString(); + } + + String executor_id; + ShardInfoList shard_info_list; +}; + +using TableShardInfoMap = std::unordered_map; +using TableShardInfoList = std::vector; + +class QueryShardInfos +{ +public: + QueryShardInfos() = default; + + static QueryShardInfos create(const google::protobuf::RepeatedPtrField & shard_infos) + { + QueryShardInfos query_shard_infos; + for (const auto & shard_info : shard_infos) + { + auto table_shard_info = TableShardInfos::create(shard_info); + query_shard_infos.table_shard_info_list.push_back(table_shard_info); + } + return query_shard_infos; + } + + TableShardInfos getTableShardInfosByExecutorID(String executor_id) const + { + for (const auto & table_shard_info : table_shard_info_list) + { + if (table_shard_info.executor_id == executor_id) + return table_shard_info; + } + throw Exception("No TableShardInfo found for executor ID: " + executor_id, ErrorCodes::LOGICAL_ERROR); + } + + + String toString() const + { + FmtBuffer buf; + buf.fmtAppend("QueryShardInfos: ["); + buf.joinStr( + table_shard_info_list.begin(), + table_shard_info_list.end(), + [](const TableShardInfos & table_shard_info, FmtBuffer & fb) { fb.append(table_shard_info.toString()); }, + ", "); + buf.append("]"); + return buf.toString(); + } + + size_t size() const { return table_shard_info_list.size(); } + + TableShardInfoList table_shard_info_list; +}; + +} // namespace DB diff --git a/dbms/src/Flash/Coprocessor/StorageTantivyInterpreter.cpp b/dbms/src/Flash/Coprocessor/StorageTantivyInterpreter.cpp new file mode 100644 index 00000000000..bbd534d2d6b --- /dev/null +++ b/dbms/src/Flash/Coprocessor/StorageTantivyInterpreter.cpp @@ -0,0 +1,213 @@ +// Copyright 2025 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 +#include + + +namespace DB +{ + +void StorageTantivyIterpreter::execute(PipelineExecutorContext & exec_context, PipelineExecGroupBuilder & group_builder) +{ + storage->splitRemoteReadAndLocalRead(); + // local_read + storage->read(exec_context, group_builder, Names(), SelectQueryInfo(), context, 0, max_streams); + + // remote_read + const auto & executor_id = tici_scan.getTiCIScan()->executor_id(); + auto remote_shard_infos = storage->getRemoteShardInfos(); + const auto total_shards = tici_scan.getShardInfos().shard_info_list.size(); + const auto remote_shards = remote_shard_infos.size(); + const auto local_shards = total_shards >= remote_shards ? (total_shards - remote_shards) : 0; + LOG_INFO( + log, + "tici split shards: executor_id={} total={} local={} remote={} local_concurrency={} start_ts={}", + executor_id, + total_shards, + local_shards, + remote_shards, + group_builder.concurrency(), + context.getSettingsRef().read_tso); + if (!remote_shard_infos.empty()) + { + if (context.getDAGContext()->isCop()) + throw RegionException({}, RegionException::RegionReadStatus::NOT_FOUND, "shard not found"); + context.getDAGContext()->retry_shards.insert( + context.getDAGContext()->retry_shards.end(), + remote_shard_infos.begin(), + remote_shard_infos.end()); + } + + auto remote_request = buildRemoteRequests(remote_shard_infos); + if (!remote_request.empty()) + { + PipelineExecGroupBuilder remote_builder; + buildRemoteExec(exec_context, remote_builder, remote_request); + if (!remote_builder.empty()) + LOG_INFO( + log, + "tici remote sourceOps built: executor_id={} remote_requests={} concurrency={} start_ts={}", + executor_id, + remote_request.size(), + remote_builder.concurrency(), + context.getSettingsRef().read_tso); + group_builder.merge(std::move(remote_builder)); + } +} + +std::vector StorageTantivyIterpreter::buildRemoteRequests(ShardInfoList & remote_shard_infos) +{ + std::vector remote_requests; + if (remote_shard_infos.empty()) + { + return remote_requests; + } + remote_requests.push_back(RemoteRequest::build( + remote_shard_infos, + *context.getDAGContext(), + tici_scan, + context.getDAGContext()->getConnectionID(), + context.getDAGContext()->getConnectionAlias(), + log)); + return remote_requests; +} + +void StorageTantivyIterpreter::buildRemoteExec( + PipelineExecutorContext & exec_context, + PipelineExecGroupBuilder & group_builder, + const std::vector & remote_requests) +{ + auto coprocessor_reader = buildCoprocessorReader(remote_requests); + size_t concurrent_num = coprocessor_reader->enableCopStream() ? context.getSettingsRef().max_threads.get() + : coprocessor_reader->getConcurrency(); + + for (size_t i = 0; i < concurrent_num; ++i) + group_builder.addConcurrency( + std::make_unique(exec_context, log->identifier(), coprocessor_reader)); + + LOG_DEBUG(log, "remote sourceOps built"); +} + +CoprocessorReaderPtr StorageTantivyIterpreter::buildCoprocessorReader( + const std::vector & remote_requests) +{ + std::vector all_tasks = buildCopTasks(remote_requests); + if (all_tasks.empty()) + { + throw TiFlashException( + "No coprocessor tasks built for remote read, please check remote request", + Errors::Coprocessor::BadRequest); + } + const DAGSchema & schema = remote_requests[0].schema; + pingcap::kv::Cluster * cluster = tmt.getKVCluster(); + bool has_enforce_encode_type = true; + pingcap::kv::LabelFilter tiflash_label_filter = pingcap::kv::labelFilterNoTiFlashWriteNode; + + size_t concurrent_num = std::min(context.getSettingsRef().max_threads, all_tasks.size()); + size_t queue_size = context.getSettingsRef().remote_read_queue_size > 0 + ? context.getSettingsRef().remote_read_queue_size.get() + : concurrent_num * 4; + bool enable_cop_stream = context.getSettingsRef().enable_cop_stream_for_remote_read; + UInt64 cop_timeout = context.getSettingsRef().cop_timeout_for_remote_read; + String store_zone_label; + auto kv_store = tmt.getKVStore(); + if likely (kv_store) + { + for (int i = 0; i < kv_store->getStoreMeta().labels_size(); ++i) + { + if (kv_store->getStoreMeta().labels().at(i).key() == "zone") + { + store_zone_label = kv_store->getStoreMeta().labels().at(i).value(); + break; + } + } + } + auto coprocessor_reader = std::make_shared( + schema, + cluster, + std::move(all_tasks), + has_enforce_encode_type, + concurrent_num, + enable_cop_stream, + queue_size, + cop_timeout, + tiflash_label_filter, + log->identifier(), + store_zone_label); + context.getDAGContext()->addCoprocessorReader(coprocessor_reader); + + return coprocessor_reader; +} + +std::vector StorageTantivyIterpreter::buildCopTasks( + const std::vector & remote_requests) +{ + std::vector all_tasks; + + pingcap::kv::Cluster * cluster = tmt.getKVCluster(); + for (const auto & remote_request : remote_requests) + { + pingcap::coprocessor::RequestPtr req = std::make_shared(); + remote_request.dag_request.SerializeToString(&(req->data)); + req->tp = pingcap::coprocessor::ReqType::DAG; + req->start_ts = context.getSettingsRef().read_tso; + req->schema_version = context.getSettingsRef().schema_version; + req->resource_group_name = (*context.getDAGContext()).getResourceGroupName(); + + pingcap::kv::Backoffer bo(pingcap::kv::copBuildTaskMaxBackoff); + pingcap::kv::StoreType store_type = pingcap::kv::StoreType::TiFlash; + std::multimap meta_data; + meta_data.emplace("is_remote_read", "true"); + + auto tasks = pingcap::coprocessor::buildCopTaskForFullText( + bo, + cluster, + remote_request.key_ranges, + req, + store_type, + (*context.getDAGContext()).getKeyspaceID(), + remote_request.connection_id, + remote_request.connection_alias, + &Poco::Logger::get("pingcap/coprocessor"), + std::move(meta_data), + [&] { GET_METRIC(tiflash_coprocessor_request_count, type_remote_read_sent).Increment(); }, + tici_scan.getTableId(), + tici_scan.getIndexId(), + tici_scan.getTiCIScan()->executor_id()); + all_tasks.insert(all_tasks.end(), tasks.begin(), tasks.end()); + } + GET_METRIC(tiflash_coprocessor_request_count, type_remote_read_constructed) + .Increment(static_cast(all_tasks.size())); + return all_tasks; +} + +} // namespace DB diff --git a/dbms/src/Flash/Coprocessor/StorageTantivyInterpreter.h b/dbms/src/Flash/Coprocessor/StorageTantivyInterpreter.h new file mode 100644 index 00000000000..0d4f47c0805 --- /dev/null +++ b/dbms/src/Flash/Coprocessor/StorageTantivyInterpreter.h @@ -0,0 +1,71 @@ +// Copyright 2025 PingCAP, Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace DB +{ + +class StorageTantivyIterpreter +{ +public: + StorageTantivyIterpreter(Context & context_, const TiCIScan & tici_scan_, size_t max_streams_) + : context(context_) + , storage(std::make_unique(context_, tici_scan_)) + , max_streams(max_streams_) + , tmt(context.getTMTContext()) + , log(Logger::get(context.getDAGContext()->log ? context.getDAGContext()->log->identifier() : "")) + , tici_scan(tici_scan_) + {} + + void execute(PipelineExecutorContext & exec_context, PipelineExecGroupBuilder & group_builder); + + std::vector buildRemoteRequests(ShardInfoList & remote_shard_infos); + void buildRemoteExec( + PipelineExecutorContext & exec_context, + PipelineExecGroupBuilder & group_builder, + const std::vector & remote_requests); + CoprocessorReaderPtr buildCoprocessorReader(const std::vector & remote_requests); + std::vector buildCopTasks(const std::vector & remote_requests); + + // Members will be transferred to DAGQueryBlockInterpreter after execute + std::unique_ptr analyzer; + +private: + Context & context; + std::unique_ptr storage; + size_t max_streams; + + TMTContext & tmt; + LoggerPtr log; + const TiCIScan tici_scan; +}; +} // namespace DB diff --git a/dbms/src/Flash/Coprocessor/TiCIScan.cpp b/dbms/src/Flash/Coprocessor/TiCIScan.cpp new file mode 100644 index 00000000000..1e65076e67a --- /dev/null +++ b/dbms/src/Flash/Coprocessor/TiCIScan.cpp @@ -0,0 +1,49 @@ +// Copyright 2025 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 +namespace DB +{ +TiCIScan::TiCIScan(const tipb::Executor * tici_scan_, const String & executor_id_, const DAGContext & dag_context) + : tici_scan(tici_scan_) + , executor_id(executor_id_) + , keyspace_id(dag_context.getKeyspaceID()) + , table_id(tici_scan->idx_scan().table_id()) + , index_id(tici_scan->idx_scan().index_id()) + , return_columns(TiDB::toTiDBColumnInfos(tici_scan->idx_scan().columns())) + , query_type(tici_scan->idx_scan().fts_query_info().query_type()) + , shard_infos(dag_context.query_shard_infos.getTableShardInfosByExecutorID(tici_scan_->executor_id())) + , limit(tici_scan->idx_scan().fts_query_info().top_k()) + , sort_column_ids( + tici_scan->idx_scan().fts_query_info().sort_column_ids().begin(), + tici_scan->idx_scan().fts_query_info().sort_column_ids().end()) + , sort_column_asc( + tici_scan->idx_scan().fts_query_info().sort_column_asc().begin(), + tici_scan->idx_scan().fts_query_info().sort_column_asc().end()) +{} + +void TiCIScan::constructTiCIScanForRemoteRead(tipb::IndexScan * tipb_index_scan) const +{ + assert(tipb_index_scan != nullptr); + *tipb_index_scan = tici_scan->idx_scan(); +} + +} // namespace DB diff --git a/dbms/src/Flash/Coprocessor/TiCIScan.h b/dbms/src/Flash/Coprocessor/TiCIScan.h new file mode 100644 index 00000000000..a9229dcfb98 --- /dev/null +++ b/dbms/src/Flash/Coprocessor/TiCIScan.h @@ -0,0 +1,76 @@ +// Copyright 2025 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 +namespace DB +{ +class DAGContext; + +class TiCIScan +{ +public: + TiCIScan(const tipb::Executor * tici_scan_, const String & executor_id_, const DAGContext & dag_context); + explicit TiCIScan(const tipb::Executor * tici_scan_); + + // Do not use getReturnColumns when is_count_agg is true + const TiDB::ColumnInfos & getReturnColumns() const { return return_columns; } + void setNamesAndTypes(const NamesAndTypes & new_names_and_types) { names_and_types = new_names_and_types; } + const NamesAndTypes getNamesAndTypes() const { return names_and_types; } + const TableShardInfos & getShardInfos() const { return shard_infos; } + const uint & getKeyspaceID() const { return keyspace_id; } + const int & getTableId() const { return table_id; } + const int & getIndexId() const { return index_id; } + const int & getLimit() const { return limit; } + const tipb::Executor * getTiCIScan() const { return tici_scan; } + + void constructTiCIScanForRemoteRead(tipb::IndexScan * tipb_index_scan) const; + + const ::google::protobuf::RepeatedPtrField<::tipb::Expr> & getMatchExpr() const + { + return tici_scan->idx_scan().fts_query_info().match_expr(); + } + + bool isCount() const { return is_count_agg; } + + void setIsCountAgg(bool v) { is_count_agg = v; } + + void setCountAggExecutorId(const String & v) { count_agg_executor_id = v; } + + const String & getCountAggExecutorId() const { return count_agg_executor_id; } + + const std::vector & getSortColumnIds() const { return sort_column_ids; } + + const std::vector & getSortColumnAsc() const { return sort_column_asc; } + +private: + const tipb::Executor * tici_scan; + [[maybe_unused]] String executor_id; + const uint keyspace_id; + const int table_id; + const int index_id; + TiDB::ColumnInfos return_columns; + NamesAndTypes names_and_types; + [[maybe_unused]] tipb::FTSQueryType query_type; + const TableShardInfos shard_infos; + const int limit; + std::vector sort_column_ids; + std::vector sort_column_asc; + bool is_count_agg = false; + String count_agg_executor_id; +}; +} // namespace DB diff --git a/dbms/src/Flash/Coprocessor/collectOutputFieldTypes.cpp b/dbms/src/Flash/Coprocessor/collectOutputFieldTypes.cpp index 6c9ec8d1ed0..f655733d84f 100644 --- a/dbms/src/Flash/Coprocessor/collectOutputFieldTypes.cpp +++ b/dbms/src/Flash/Coprocessor/collectOutputFieldTypes.cpp @@ -17,6 +17,7 @@ #include #include #include +#include namespace DB { @@ -103,6 +104,22 @@ bool collectForTableScan(std::vector & output_field_types, cons return false; } +bool collectForIndexScan(std::vector & output_field_types, const tipb::IndexScan & tici_scan) +{ + for (const auto & ci : tici_scan.columns()) + { + tipb::FieldType field_type; + field_type.set_tp(ci.tp()); + field_type.set_flag(ci.flag()); + field_type.set_flen(ci.columnlen()); + field_type.set_decimal(ci.decimal()); + for (const auto & elem : ci.elems()) + field_type.add_elems(elem); + output_field_types.push_back(field_type); + } + return false; +} + bool collectForExpand2(std::vector & output_field_types, const tipb::Expand2 & expand2) { // just collect from the level one. @@ -256,6 +273,8 @@ bool collectForExecutor(std::vector & output_field_types, const return collectForExpand(output_field_types, executor); case tipb::ExecType::TypeExpand2: return collectForExpand2(output_field_types, executor.expand2()); + case tipb::ExecType::TypeIndexScan: + return collectForIndexScan(output_field_types, executor.idx_scan()); default: return true; } diff --git a/dbms/src/Flash/CoprocessorHandler.cpp b/dbms/src/Flash/CoprocessorHandler.cpp index cfe3b85e4de..6d42c1bab82 100644 --- a/dbms/src/Flash/CoprocessorHandler.cpp +++ b/dbms/src/Flash/CoprocessorHandler.cpp @@ -100,8 +100,9 @@ grpc::Status CoprocessorHandler::execute() try { + bool is_tici = cop_request->table_shard_infos_size() != 0; RUNTIME_CHECK_MSG( - !cop_context.db_context.getSharedContextDisagg()->isDisaggregatedComputeMode(), + !cop_context.db_context.getSharedContextDisagg()->isDisaggregatedComputeMode() || is_tici, "cannot run cop or batchCop request on tiflash_compute node"); switch (cop_request->tp()) @@ -164,6 +165,7 @@ grpc::Status CoprocessorHandler::execute() DAGContext dag_context( dag_request, std::move(tables_regions_info), + QueryShardInfos::create(cop_request->table_shard_infos()), RequestUtils::deriveKeyspaceID(cop_request->context()), cop_context.db_context.getClientInfo().current_address.toString(), kind, diff --git a/dbms/src/Flash/EstimateTiCICountHandler.cpp b/dbms/src/Flash/EstimateTiCICountHandler.cpp new file mode 100644 index 00000000000..8caeddb6491 --- /dev/null +++ b/dbms/src/Flash/EstimateTiCICountHandler.cpp @@ -0,0 +1,133 @@ +// 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 + +namespace DB +{ +namespace +{ +TimezoneInfo buildEstimateTimezoneInfo(const coprocessor::TiCIEstimateCountRequest & request) +{ + TimezoneInfo timezone_info; + if (!request.time_zone_name().empty()) + timezone_info.resetByTimezoneName(request.time_zone_name()); + else + timezone_info.resetByTimezoneOffset(request.time_zone_offset()); + return timezone_info; +} + +tipb::FTSQueryInfo parseEstimateQueryInfo(const coprocessor::TiCIEstimateCountRequest & request) +{ + tipb::FTSQueryInfo query_info; + if (!query_info.ParseFromString(request.fts_query_info())) + throw TiFlashException("Failed to parse fts_query_info", Errors::Coprocessor::BadRequest); + if (query_info.match_expr_size() == 0) + throw TiFlashException("Empty TiCI estimate query expression", Errors::Coprocessor::BadRequest); + return query_info; +} + +rust::Vec<::ShardWithRange> buildEstimateShardRanges(const coprocessor::TiCIEstimateCountRequest & request) +{ + rust::Vec<::ShardWithRange> shards; + for (const auto & shard_info : request.shard_infos()) + { + shards.push_back({ + .shard_id = shard_info.shard_id(), + .ranges = TS::getKeyRanges(shard_info.ranges()), + }); + } + return shards; +} +} // namespace + +EstimateTiCICountHandler::EstimateTiCICountHandler( + const coprocessor::TiCIEstimateCountRequest * request_, + coprocessor::TiCIEstimateCountResponse * response_, + const String & identifier) + : request(request_) + , response(response_) + , log(Logger::get(identifier)) +{} + +grpc::Status EstimateTiCICountHandler::execute() +{ + try + { + if (request->shard_infos_size() == 0) + return grpc::Status::OK; + + const auto keyspace_id = RequestUtils::deriveKeyspaceID(request->context()); + const auto fts_query_info = parseEstimateQueryInfo(*request); + const auto timezone_info = buildEstimateTimezoneInfo(*request); + auto [query, column_ids] = TS::tipbToTiCIExpr(fts_query_info.match_expr(), timezone_info); + (void)column_ids; + + auto shard_ranges = buildEstimateShardRanges(*request); + const auto estimate_result = estimate_count(keyspace_id, shard_ranges, query); + response->set_est_count(estimate_result.estimated_total_count); + LOG_DEBUG( + log, + "GetEstimateTiCICount done, est_count={}, input_shards={}, available_shards={}, sampled_shards={}", + response->est_count(), + request->shard_infos_size(), + estimate_result.available_shards, + estimate_result.sampled_shards); + } + catch (const TiFlashException & e) + { + LOG_WARNING( + log, + "GetEstimateTiCICount failed with TiFlash exception: {}\n{}", + e.displayText(), + e.getStackTrace().toString()); + response->set_other_error(e.standardText()); + } + catch (const Exception & e) + { + LOG_WARNING( + log, + "GetEstimateTiCICount failed with DB exception: {}\n{}", + e.message(), + e.getStackTrace().toString()); + response->set_other_error(e.message()); + } + catch (const pingcap::Exception & e) + { + LOG_WARNING(log, "GetEstimateTiCICount failed with KV exception: {}", e.message()); + response->set_other_error(e.message()); + } + catch (const std::exception & e) + { + LOG_WARNING(log, "GetEstimateTiCICount failed: {}", e.what()); + response->set_other_error(e.what()); + } + catch (...) + { + LOG_WARNING(log, "GetEstimateTiCICount failed with unknown exception"); + response->set_other_error("other exception"); + } + + return grpc::Status::OK; +} + +} // namespace DB diff --git a/dbms/src/Flash/EstimateTiCICountHandler.h b/dbms/src/Flash/EstimateTiCICountHandler.h new file mode 100644 index 00000000000..6d116401f5f --- /dev/null +++ b/dbms/src/Flash/EstimateTiCICountHandler.h @@ -0,0 +1,40 @@ +// 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 + +namespace DB +{ +class EstimateTiCICountHandler +{ +public: + EstimateTiCICountHandler( + const coprocessor::TiCIEstimateCountRequest * request_, + coprocessor::TiCIEstimateCountResponse * response_, + const String & identifier); + + grpc::Status execute(); + +private: + const coprocessor::TiCIEstimateCountRequest * request; + coprocessor::TiCIEstimateCountResponse * response; + const LoggerPtr log; +}; + +} // namespace DB diff --git a/dbms/src/Flash/FlashService.cpp b/dbms/src/Flash/FlashService.cpp index f99aa66f140..e0e75753900 100644 --- a/dbms/src/Flash/FlashService.cpp +++ b/dbms/src/Flash/FlashService.cpp @@ -27,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -887,6 +888,20 @@ grpc::Status FlashService::Compact( return manual_compact_manager->handleRequest(request, response); } +grpc::Status FlashService::GetEstimateTiCICount( + grpc::ServerContext * grpc_context, + const coprocessor::TiCIEstimateCountRequest * request, + coprocessor::TiCIEstimateCountResponse * response) +{ + CPUAffinityManager::getInstance().bindSelfGrpcThread(); + auto check_result = checkGrpcContext(grpc_context); + if (!check_result.ok()) + return check_result; + + EstimateTiCICountHandler handler(request, response, log->identifier()); + return handler.execute(); +} + grpc::Status FlashService::tryAddLock( grpc::ServerContext * grpc_context, const disaggregated::TryAddLockRequest * request, diff --git a/dbms/src/Flash/FlashService.h b/dbms/src/Flash/FlashService.h index 0c87fa5493a..d4323131a6e 100644 --- a/dbms/src/Flash/FlashService.h +++ b/dbms/src/Flash/FlashService.h @@ -100,6 +100,10 @@ class FlashService const kvrpcpb::CompactRequest * request, kvrpcpb::CompactResponse * response) override; + grpc::Status GetEstimateTiCICount( + grpc::ServerContext * grpc_context, + const coprocessor::TiCIEstimateCountRequest * request, + coprocessor::TiCIEstimateCountResponse * response) override; // For S3 Lock Service grpc::Status tryAddLock( diff --git a/dbms/src/Flash/Mpp/MPPTask.cpp b/dbms/src/Flash/Mpp/MPPTask.cpp index 67e65f3b40a..ac15d4f859d 100644 --- a/dbms/src/Flash/Mpp/MPPTask.cpp +++ b/dbms/src/Flash/Mpp/MPPTask.cpp @@ -545,6 +545,7 @@ void MPPTask::prepare(const mpp::DispatchTaskRequest & task_request) dag_context->log = log; dag_context->tables_regions_info = std::move(tables_regions_info); dag_context->setBypassLockTs(std::move(bypass_lock_ts)); + dag_context->query_shard_infos = QueryShardInfos::create(task_request.table_shard_infos()); dag_context->tidb_host = context->getClientInfo().current_address.toString(); context->setDAGContext(dag_context.get()); diff --git a/dbms/src/Flash/Planner/PhysicalPlan.cpp b/dbms/src/Flash/Planner/PhysicalPlan.cpp index e050a01cf52..bb4f2158a38 100644 --- a/dbms/src/Flash/Planner/PhysicalPlan.cpp +++ b/dbms/src/Flash/Planner/PhysicalPlan.cpp @@ -12,10 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include #include #include +#include #include #include #include @@ -36,12 +38,14 @@ #include #include #include +#include #include #include #include #include #include #include +#include namespace DB { @@ -92,6 +96,26 @@ void PhysicalPlan::buildTableScan(const String & executor_id, const tipb::Execut dagContext().table_scan_executor_id = executor_id; } +void PhysicalPlan::buildTiCIScan(const String & executor_id, const tipb::Executor * executor) +{ + RUNTIME_ASSERT(executor->idx_scan().has_fts_query_info()); + TiCIScan tici_scan(executor, executor_id, dagContext()); + LOG_INFO( + log, + "tici scan: keyspace_id={} table_id={} index_id={} limit={} shard_count={} match_expr_size={} query_type={} " + "start_ts={}", + tici_scan.getKeyspaceID(), + tici_scan.getTableId(), + tici_scan.getIndexId(), + tici_scan.getLimit(), + tici_scan.getShardInfos().shard_info_list.size(), + tici_scan.getMatchExpr().size(), + tipb::FTSQueryType_Name(executor->idx_scan().fts_query_info().query_type()), + context.getSettingsRef().read_tso); + pushBack(PhysicalTiCIScan::build(executor_id, log, tici_scan)); + dagContext().table_scan_executor_id = executor_id; +} + void PhysicalPlan::build(const tipb::Executor * executor) { RUNTIME_CHECK(executor); @@ -201,6 +225,10 @@ void PhysicalPlan::build(const tipb::Executor * executor) GET_METRIC(tiflash_coprocessor_executor_count, type_partition_ts).Increment(); buildTableScan(executor_id, executor); break; + case tipb::ExecType::TypeIndexScan: + GET_METRIC(tiflash_coprocessor_executor_count, type_tici).Increment(); + buildTiCIScan(executor_id, executor); + break; case tipb::ExecType::TypeJoin: { GET_METRIC(tiflash_coprocessor_executor_count, type_join).Increment(); diff --git a/dbms/src/Flash/Planner/PhysicalPlan.h b/dbms/src/Flash/Planner/PhysicalPlan.h index 1668595eab2..c148c3e5661 100644 --- a/dbms/src/Flash/Planner/PhysicalPlan.h +++ b/dbms/src/Flash/Planner/PhysicalPlan.h @@ -63,6 +63,7 @@ class PhysicalPlan DAGContext & dagContext() const; void buildTableScan(const String & executor_id, const tipb::Executor * executor); + void buildTiCIScan(const String & executor_id, const tipb::Executor * executor); private: std::vector cur_plan_nodes; diff --git a/dbms/src/Flash/Planner/PhysicalPlanNode.h b/dbms/src/Flash/Planner/PhysicalPlanNode.h index 19f3f68cf9b..aeeacc94687 100644 --- a/dbms/src/Flash/Planner/PhysicalPlanNode.h +++ b/dbms/src/Flash/Planner/PhysicalPlanNode.h @@ -57,6 +57,8 @@ class PhysicalPlanNode : public std::enable_shared_from_this virtual PhysicalPlanNodePtr children(size_t /*i*/) const = 0; + virtual void setChild(size_t i, PhysicalPlanNodePtr new_child) = 0; + const PlanType & tp() const { return type; } const String & execId() const { return executor_id; } diff --git a/dbms/src/Flash/Planner/PlanType.h b/dbms/src/Flash/Planner/PlanType.h index 26010abdd8b..34a03b114db 100644 --- a/dbms/src/Flash/Planner/PlanType.h +++ b/dbms/src/Flash/Planner/PlanType.h @@ -44,6 +44,7 @@ struct PlanType GetResult = 19, CTESource = 20, CTESink = 21, + TiCiScan = 22, }; PlanTypeEnum enum_value; diff --git a/dbms/src/Flash/Planner/Plans/PhysicalAggregation.h b/dbms/src/Flash/Planner/Plans/PhysicalAggregation.h index 6be21ded21b..b15f0a28e53 100644 --- a/dbms/src/Flash/Planner/Plans/PhysicalAggregation.h +++ b/dbms/src/Flash/Planner/Plans/PhysicalAggregation.h @@ -66,6 +66,14 @@ class PhysicalAggregation : public PhysicalUnary const Block & getSampleBlock() const override; + bool isCountNotNullableColumnWithoutGroupbyKey() const + { + // Only single count(const/not-null column) without any group by key pattern can be optimized by TiCI + // TiFlash will convert count(const/not-null column) to count(), so the argument_names is empty. + return aggregate_descriptions.size() == 1 && aggregate_descriptions[0].function->getName() == "count" + && aggregate_descriptions[0].argument_names.empty() && aggregation_keys.empty(); + } + private: void buildBlockInputStreamImpl(DAGPipeline & pipeline, Context & context, size_t max_streams) override; diff --git a/dbms/src/Flash/Planner/Plans/PhysicalBinary.h b/dbms/src/Flash/Planner/Plans/PhysicalBinary.h index 6a78cfa6660..1886a3b39d1 100644 --- a/dbms/src/Flash/Planner/Plans/PhysicalBinary.h +++ b/dbms/src/Flash/Planner/Plans/PhysicalBinary.h @@ -48,7 +48,9 @@ class PhysicalBinary : public PhysicalPlanNode return i == 0 ? left : right; } - size_t childrenSize() const override { return 2; }; + void setChild(size_t, PhysicalPlanNodePtr) override { throw Exception("Unsupport"); } + + size_t childrenSize() const override { return 2; } protected: PhysicalPlanNodePtr left; diff --git a/dbms/src/Flash/Planner/Plans/PhysicalLeaf.h b/dbms/src/Flash/Planner/Plans/PhysicalLeaf.h index 549953930d2..ade58afbdcd 100644 --- a/dbms/src/Flash/Planner/Plans/PhysicalLeaf.h +++ b/dbms/src/Flash/Planner/Plans/PhysicalLeaf.h @@ -39,6 +39,8 @@ class PhysicalLeaf : public PhysicalPlanNode throw TiFlashException("the children size of PhysicalLeaf is zero", Errors::Planner::Internal); } - size_t childrenSize() const override { return 0; }; + void setChild(size_t, PhysicalPlanNodePtr) override { throw Exception("Unsupport"); } + + size_t childrenSize() const override { return 0; } }; -} // namespace DB \ No newline at end of file +} // namespace DB diff --git a/dbms/src/Flash/Planner/Plans/PhysicalTiCIScan.cpp b/dbms/src/Flash/Planner/Plans/PhysicalTiCIScan.cpp new file mode 100644 index 00000000000..55288c74391 --- /dev/null +++ b/dbms/src/Flash/Planner/Plans/PhysicalTiCIScan.cpp @@ -0,0 +1,151 @@ +// Copyright 2025 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 + + +namespace DB +{ +namespace +{ +NamesWithAliases buildTableScanProjectionCols( + Int64 logical_table_id, + const NamesAndTypes & schema, + const Block & storage_header) +{ + if (unlikely(schema.size() != storage_header.columns())) + throw TiFlashException( + fmt::format( + "The tidb table scan schema size {} is different from the tiflash storage schema size {}, table id is " + "{}", + schema.size(), + storage_header.columns(), + logical_table_id), + Errors::Planner::BadRequest); + NamesWithAliases schema_project_cols; + for (size_t i = 0; i < schema.size(); ++i) + { + const auto & table_scan_col_name = schema[i].name; + const auto & table_scan_col_type = schema[i].type; + const auto & storage_col_name = storage_header.getColumnsWithTypeAndName()[i].name; + const auto & storage_col_type = storage_header.getColumnsWithTypeAndName()[i].type; + if (unlikely(!table_scan_col_type->equals(*storage_col_type))) + throw TiFlashException( + fmt::format( + R"(The data type {} from tidb table scan schema is different from the data type {} from tiflash storage schema, + table id is {}, + column index is {}, + column name from tidb table scan is {}, + column name from tiflash storage is {})", + table_scan_col_type->getName(), + storage_col_type->getName(), + logical_table_id, + i, + table_scan_col_name, + storage_col_name), + Errors::Planner::BadRequest); + schema_project_cols.emplace_back(storage_col_name, table_scan_col_name); + } + return schema_project_cols; +} +} // namespace + +PhysicalTiCIScan::PhysicalTiCIScan( + const String & executor_id_, + const NamesAndTypes & schema_, + const String & req_id, + const TiCIScan & tici_scan_, + const Block & sample_block_) + : PhysicalLeaf(executor_id_, PlanType::TiCiScan, schema_, FineGrainedShuffle{}, req_id) + , tici_scan(tici_scan_) + , sample_block(sample_block_) +{} + +PhysicalPlanNodePtr PhysicalTiCIScan::build( + const String & executor_id, + const LoggerPtr & log, + const TiCIScan & tici_scan) +{ + auto schema = genNamesAndTypes(tici_scan.getReturnColumns(), "tici_scan"); + auto physical_tici_scan + = std::make_shared(executor_id, schema, log->identifier(), tici_scan, Block(schema)); + return physical_tici_scan; +} + +void PhysicalTiCIScan::buildPipeline( + PipelineBuilder & builder, + Context & context, + PipelineExecutorContext & exec_context) +{ + LOG_INFO( + log, + "build tici pipeline: executor_id={} is_count={} max_streams={} start_ts={}", + executor_id, + tici_scan.isCount(), + context.getMaxStreams(), + context.getSettingsRef().read_tso); + // For building PipelineExec in compile time. + StorageTantivyIterpreter storage_interpreter(context, tici_scan, context.getMaxStreams()); + storage_interpreter.execute(exec_context, pipeline_exec_builder); + if (!tici_scan.isCount()) + buildProjection(exec_context, pipeline_exec_builder); + + PhysicalPlanNode::buildPipeline(builder, context, exec_context); +} + +void PhysicalTiCIScan::finalizeImpl(const Names & parent_require) +{ + FinalizeHelper::checkSchemaContainsParentRequire(schema, parent_require); +} + +const Block & PhysicalTiCIScan::getSampleBlock() const +{ + return sample_block; +} + +void PhysicalTiCIScan::buildPipelineExecGroupImpl( + PipelineExecutorContext & /*exec_status*/, + PipelineExecGroupBuilder & group_builder, + Context & /*context*/, + size_t /*concurrency*/) +{ + assert(group_builder.empty()); + group_builder = std::move(pipeline_exec_builder); +} + +void PhysicalTiCIScan::buildProjection(PipelineExecutorContext & exec_context, PipelineExecGroupBuilder & group_builder) +{ + auto header = group_builder.getCurrentHeader(); + const auto & schema_project_cols = buildTableScanProjectionCols(tici_scan.getTableId(), schema, header); + + /// In order to keep TransformOp's schema consistent with PhysicalPlan's schema. + /// It is worth noting that the column uses the name as the unique identifier in the Block, so the column name must also be consistent. + ExpressionActionsPtr schema_actions = PhysicalPlanHelper::newActions(header); + schema_actions->add(ExpressionAction::project(schema_project_cols)); + executeExpression(exec_context, group_builder, schema_actions, log); +} +} // namespace DB diff --git a/dbms/src/Flash/Planner/Plans/PhysicalTiCIScan.h b/dbms/src/Flash/Planner/Plans/PhysicalTiCIScan.h new file mode 100644 index 00000000000..10d161611cc --- /dev/null +++ b/dbms/src/Flash/Planner/Plans/PhysicalTiCIScan.h @@ -0,0 +1,73 @@ +// Copyright 2025 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 + +namespace DB +{ +class PhysicalTiCIScan : public PhysicalLeaf +{ +public: + static PhysicalPlanNodePtr build(const String & executor_id, const LoggerPtr & log, const TiCIScan & tici_scan); + + PhysicalTiCIScan( + const String & executor_id_, + const NamesAndTypes & schema_, + const String & req_id, + const TiCIScan & tici_scan_, + const Block & sample_block_); + + void finalizeImpl(const Names & parent_require) override; + + const Block & getSampleBlock() const override; + + bool setFilterConditions(const String & filter_executor_id, const tipb::Selection & selection); + + bool hasFilterConditions() const; + + const String & getFilterConditionsId() const; + + void buildPipeline(PipelineBuilder & builder, Context & context, PipelineExecutorContext & exec_context) override; + + void setCountAgg(std::shared_ptr agg) + { + tici_scan.setIsCountAgg(true); + schema = agg->getSchema(); + tici_scan.setNamesAndTypes(agg->getSchema()); + tici_scan.setCountAggExecutorId(agg->execId()); + } + +private: + void buildPipelineExecGroupImpl( + PipelineExecutorContext & /*exec_status*/, + PipelineExecGroupBuilder & group_builder, + Context & /*context*/, + size_t /*concurrency*/) override; + void buildProjection(PipelineExecutorContext & exec_context, PipelineExecGroupBuilder & group_builder); + FilterConditions filter_conditions; + + TiCIScan tici_scan; + + Block sample_block; + + PipelineExecGroupBuilder pipeline_exec_builder; +}; +} // namespace DB diff --git a/dbms/src/Flash/Planner/Plans/PhysicalUnary.h b/dbms/src/Flash/Planner/Plans/PhysicalUnary.h index fada29a645c..6695e4e947f 100644 --- a/dbms/src/Flash/Planner/Plans/PhysicalUnary.h +++ b/dbms/src/Flash/Planner/Plans/PhysicalUnary.h @@ -46,7 +46,13 @@ class PhysicalUnary : public PhysicalPlanNode return child; } - size_t childrenSize() const override { return 1; }; + void setChild(size_t i, PhysicalPlanNodePtr new_child) override + { + RUNTIME_CHECK_MSG(i == 0, "child_index({}) shouldn't >= childrenSize({})", i, childrenSize()); + child = std::move(new_child); + } + + size_t childrenSize() const override { return 1; } protected: PhysicalPlanNodePtr child; diff --git a/dbms/src/Flash/Planner/optimize.cpp b/dbms/src/Flash/Planner/optimize.cpp index 216dab04296..8978f7d69e7 100644 --- a/dbms/src/Flash/Planner/optimize.cpp +++ b/dbms/src/Flash/Planner/optimize.cpp @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include #include #include @@ -38,10 +40,49 @@ class FinalizeRule : public Rule static RulePtr create() { return std::make_shared(); } }; +class TiCICountAggOptimizeRule : public Rule +{ +public: + PhysicalPlanNodePtr apply(const Context & context, PhysicalPlanNodePtr plan, const LoggerPtr & logger) override + { + if (plan->tp() == PlanType::Aggregation) + { + auto agg = std::static_pointer_cast(plan); + if (agg->isCountNotNullableColumnWithoutGroupbyKey()) + { + auto child = agg->children(0); + if (child->tp() == PlanType::TiCiScan) + { + auto tici_scan = std::static_pointer_cast(child); + tici_scan->setCountAgg(agg); + // return tici_scan directly to replace agg + return tici_scan; + } + } + } + if (plan->childrenSize() != 0) + { + for (size_t i = 0; i < plan->childrenSize(); ++i) + { + auto new_child = apply(context, plan->children(i), logger); + if (new_child != plan->children(i)) + { + // replace child with the PhysicalPlanNodePtr returned by apply + plan->setChild(i, new_child); + } + } + } + + return plan; + } + + static RulePtr create() { return std::make_shared(); } +}; + PhysicalPlanNodePtr optimize(const Context & context, PhysicalPlanNodePtr plan, const LoggerPtr & log) { RUNTIME_CHECK(plan); - static std::vector rules{FinalizeRule::create()}; + static std::vector rules{TiCICountAggOptimizeRule::create(), FinalizeRule::create()}; for (const auto & rule : rules) { plan = rule->apply(context, plan, log); diff --git a/dbms/src/Flash/Statistics/CommonExecutorImpl.h b/dbms/src/Flash/Statistics/CommonExecutorImpl.h index 91adf116a92..849b6f98c74 100644 --- a/dbms/src/Flash/Statistics/CommonExecutorImpl.h +++ b/dbms/src/Flash/Statistics/CommonExecutorImpl.h @@ -138,4 +138,16 @@ struct CTESourceImpl static bool isSourceExecutor() { return true; } }; using CTESourceStatistics = ExecutorStatistics; + +struct TiCIStatisticsImpl +{ + static constexpr bool has_extra_info = false; + + static constexpr auto type = "IndexRangeScan"; + + static bool isMatch(const tipb::Executor * executor) { return executor->has_idx_scan(); } + + static bool isSourceExecutor() { return true; } +}; +using TiCIStatistics = ExecutorStatistics; } // namespace DB diff --git a/dbms/src/Flash/Statistics/ExecutorStatisticsCollector.cpp b/dbms/src/Flash/Statistics/ExecutorStatisticsCollector.cpp index a6aed782ee3..32fb6a3f050 100644 --- a/dbms/src/Flash/Statistics/ExecutorStatisticsCollector.cpp +++ b/dbms/src/Flash/Statistics/ExecutorStatisticsCollector.cpp @@ -98,7 +98,8 @@ void ExecutorStatisticsCollector::initialize(DAGContext * dag_context_) WindowStatistics, ExpandStatistics, CTESinkStatistics, - CTESourceStatistics>(&executor)) + CTESourceStatistics, + TiCIStatistics>(&executor)) { throw TiFlashException( fmt::format("Unknown executor type, executor_id: {}", executor.executor_id()), @@ -250,6 +251,7 @@ void ExecutorStatisticsCollector::fillLocalExecutionSummaries(tipb::SelectRespon } } } + return; RUNTIME_CHECK_MSG(target_executor_summary, "cannot find executor summary to put ru consumption"); RUNTIME_CHECK_MSG( local_ru->SerializeToString(target_executor_summary->mutable_ru_consumption()), diff --git a/dbms/src/Flash/Statistics/traverseExecutors.cpp b/dbms/src/Flash/Statistics/traverseExecutors.cpp index 3af2bb960c8..9f0c8e95399 100644 --- a/dbms/src/Flash/Statistics/traverseExecutors.cpp +++ b/dbms/src/Flash/Statistics/traverseExecutors.cpp @@ -24,11 +24,10 @@ Children getChildren(const tipb::Executor & executor) { case tipb::ExecType::TypeTableScan: case tipb::ExecType::TypePartitionTableScan: + case tipb::ExecType::TypeIndexScan: return {}; case tipb::ExecType::TypeJoin: return {&executor.join().children(0), &executor.join().children(1)}; - case tipb::ExecType::TypeIndexScan: - throw TiFlashException("IndexScan is not supported", Errors::Coprocessor::Unimplemented); case tipb::ExecType::TypeSelection: return Children{&executor.selection().child()}; case tipb::ExecType::TypeAggregation: diff --git a/dbms/src/Operators/TantivyReaderSourceOp.cpp b/dbms/src/Operators/TantivyReaderSourceOp.cpp new file mode 100644 index 00000000000..3073aeb5709 --- /dev/null +++ b/dbms/src/Operators/TantivyReaderSourceOp.cpp @@ -0,0 +1,93 @@ +// Copyright 2025 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 + +namespace DB +{ +TantivyReaderSourceOp::TantivyReaderSourceOp( + PipelineExecutorContext & exec_context_, + const String & req_id, + const TS::TiCIReadTaskPoolPtr task_pool_, + const NamesAndTypes & return_columns) + : SourceOp(exec_context_, req_id) + , task_pool(task_pool_) +{ + setHeader(Block(return_columns)); +} + +String TantivyReaderSourceOp::getName() const +{ + return "TantivyReaderSourceOp"; +} + +void TantivyReaderSourceOp::operatePrefixImpl() +{ + watcher.start(); +} + +void TantivyReaderSourceOp::operateSuffixImpl() +{ + LOG_DEBUG(log, "TantivyReaderSourceOp read {} rows, took {} ms", total_rows, watcher.elapsedMilliseconds()); +} + +Block TantivyReaderSourceOp::popFromBlockQueue() +{ + assert(!block_queue.empty()); + Block block = std::move(block_queue.front()); + block_queue.pop(); + return block; +} + +OperatorStatus TantivyReaderSourceOp::readImpl(Block & block) +{ + if (unlikely(done)) + { + return OperatorStatus::HAS_OUTPUT; + } + + while (true) + { + while (!cur_stream) + { + auto task = task_pool->getNextTask(); + if (!task) + { + done = true; + return OperatorStatus::HAS_OUTPUT; + } + cur_stream = task_pool->buildInputStream(task); + } + Block res = cur_stream->read(); + if (res) + { + total_rows += res.rows(); + block.swap(res); + return OperatorStatus::HAS_OUTPUT; + } + else + { + cur_stream = {}; + continue; + } + } +} + +} // namespace DB diff --git a/dbms/src/Operators/TantivyReaderSourceOp.h b/dbms/src/Operators/TantivyReaderSourceOp.h new file mode 100644 index 00000000000..bd704515a01 --- /dev/null +++ b/dbms/src/Operators/TantivyReaderSourceOp.h @@ -0,0 +1,64 @@ +// Copyright 2023 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 + +namespace DB +{ + +class TantivyReaderSourceOp : public SourceOp +{ +public: + TantivyReaderSourceOp( + PipelineExecutorContext & exec_context_, + const String & req_id, + TS::TiCIReadTaskPoolPtr task_pool_, + const NamesAndTypes & return_columns); + + String getName() const override; + +protected: + void operatePrefixImpl() override; + void operateSuffixImpl() override; + + OperatorStatus readImpl(Block & block) override; + + IOProfileInfoPtr getIOProfileInfo() const override { return io_profile_info; } + +private: + Block popFromBlockQueue(); + +private: + std::queue block_queue; + + UInt64 total_rows{}; + IOProfileInfoPtr io_profile_info; + TS::TiCIReadTaskPoolPtr task_pool; + bool done = false; + + BlockInputStreamPtr cur_stream; + Stopwatch watcher; +}; + +} // namespace DB diff --git a/dbms/src/Server/CMakeLists.txt b/dbms/src/Server/CMakeLists.txt index c19374f3598..974c7018c7e 100644 --- a/dbms/src/Server/CMakeLists.txt +++ b/dbms/src/Server/CMakeLists.txt @@ -126,6 +126,10 @@ if (OS_LINUX) DESTINATION ".") endif() + install (TARGETS tici_search_lib + COMPONENT tiflash-release + DESTINATION ".") + if (ENABLE_CLARA AND USE_INTERNAL_LIBCLARA) install (TARGETS clara_shared COMPONENT tiflash-release @@ -156,6 +160,10 @@ elseif(APPLE) COMPONENT tiflash-release DESTINATION "." PRE_INCLUDE_REGEXES ".*proxy.*|.*clara.*") + + install (TARGETS tici_search_lib + COMPONENT tiflash-release + DESTINATION ".") endif () include(CMakePrintHelpers) diff --git a/dbms/src/Server/MetricsPrometheus.cpp b/dbms/src/Server/MetricsPrometheus.cpp index 23cab7ba987..b40e3f130c4 100644 --- a/dbms/src/Server/MetricsPrometheus.cpp +++ b/dbms/src/Server/MetricsPrometheus.cpp @@ -37,9 +37,9 @@ #include #include #include -#include #include #include +#include namespace DB { @@ -101,11 +101,29 @@ class MetricHandler : public Poco::Net::HTTPRequestHandler { auto metrics = collectMetrics(); auto serializer = std::unique_ptr{new prometheus::TextSerializer()}; - String body = serializer->Serialize(metrics); + auto body = concatTextMetrics(serializer->Serialize(metrics), gather_prometheus_metrics()); response.sendBuffer(body.data(), body.size()); } private: + static String concatTextMetrics(String && tiflash_metrics, ::rust::Vec<::std::uint8_t> && tici_metrics) + { + if (tiflash_metrics.empty()) + return String{tici_metrics.begin(), tici_metrics.end()}; + if (tici_metrics.empty()) + return tiflash_metrics; + + if (tiflash_metrics.back() != '\n') + tiflash_metrics.push_back('\n'); + + tiflash_metrics.insert(tiflash_metrics.end(), tici_metrics.begin(), tici_metrics.end()); + + if (tiflash_metrics.back() != '\n') + tiflash_metrics.push_back('\n'); + + return tiflash_metrics; + } + std::vector collectMetrics() const { auto collected_metrics = std::vector{}; @@ -155,7 +173,7 @@ class MetricHandlerFactory : public Poco::Net::HTTPRequestHandlerFactory std::vector> collectables; }; -std::shared_ptr getHTTPServer( +std::shared_ptr getSecureHTTPServer( Context & global_context, std::vector> collectables, const String & address) @@ -193,6 +211,19 @@ std::shared_ptr getHTTPServer( return server; } +std::shared_ptr getHTTPServer( + std::vector> collectables, + const String & address) +{ + Poco::Net::ServerSocket socket; + Poco::Net::HTTPServerParams::Ptr http_params = new Poco::Net::HTTPServerParams; + Poco::Net::SocketAddress addr = Poco::Net::SocketAddress(address); + socket.bind(addr, true); + socket.listen(); + auto server = std::make_shared(new MetricHandlerFactory(collectables), socket, http_params); + return server; +} + constexpr Int64 MILLISECOND = 1000; constexpr Int64 INIT_DELAY = 5; @@ -281,7 +312,7 @@ MetricsPrometheus::MetricsPrometheus(Context & context, const AsynchronousMetric std::vector> collectables{ tiflash_metrics.registry, tiflash_metrics.process_collector}; - server = getHTTPServer(context, collectables, addr); + server = getSecureHTTPServer(context, collectables, addr); server->start(); LOG_INFO( log, @@ -291,9 +322,11 @@ MetricsPrometheus::MetricsPrometheus(Context & context, const AsynchronousMetric } else { - exposer = std::make_shared(addr); - exposer->RegisterCollectable(tiflash_metrics.registry); - exposer->RegisterCollectable(tiflash_metrics.process_collector); + std::vector> collectables{ + tiflash_metrics.registry, + tiflash_metrics.process_collector}; + server = getHTTPServer(collectables, addr); + server->start(); LOG_INFO( log, "Enable prometheus pull mode; Listen Host = {}, Metrics Port = {}", diff --git a/dbms/src/Server/MetricsPrometheus.h b/dbms/src/Server/MetricsPrometheus.h index 29aaa7154cd..7b743aefff8 100644 --- a/dbms/src/Server/MetricsPrometheus.h +++ b/dbms/src/Server/MetricsPrometheus.h @@ -17,7 +17,6 @@ #include #include #include -#include #include namespace DB @@ -54,7 +53,6 @@ class MetricsPrometheus int metrics_interval; std::shared_ptr gateway; - std::shared_ptr exposer; std::shared_ptr server; }; diff --git a/dbms/src/Server/Server.cpp b/dbms/src/Server/Server.cpp index 3d4769d7f34..a70876fd882 100644 --- a/dbms/src/Server/Server.cpp +++ b/dbms/src/Server/Server.cpp @@ -104,6 +104,7 @@ #include #include #include +#include #include #include @@ -1162,6 +1163,7 @@ try wn_ps->waitUntilInitedFromRemoteStore(); } + { TCPServersHolder tcp_http_servers_holder( *this, @@ -1317,6 +1319,18 @@ try proxy_machine.runKVStore(tmt_context); + auto tici_reader_addr = config().getString("tici.reader-node.addr", ""); + auto tici_reader_port = config().getInt("tici.reader-node.port", 0); + if (!tici_reader_addr.empty() || tici_reader_port > 0) + { + Stopwatch watch; + auto service_addr = config().getString("flash.service_addr"); + auto pd_addr = config().getString("raft.pd_addr"); + LOG_INFO(log, "TiCI starting, addr={}, port={}", tici_reader_addr, tici_reader_port); + start_reader_server(config_path, tici_reader_addr, service_addr, pd_addr); + LOG_INFO(log, "TiCI started, cost={}s", watch.elapsedSeconds()); + } + try { // Bind CPU affinity after all threads started. diff --git a/dbms/src/Storages/KVStore/TMTContext.cpp b/dbms/src/Storages/KVStore/TMTContext.cpp index e78f26340a1..0f2fcc7ba8f 100644 --- a/dbms/src/Storages/KVStore/TMTContext.cpp +++ b/dbms/src/Storages/KVStore/TMTContext.cpp @@ -160,6 +160,25 @@ TMTContext::TMTContext( { startMonitorMPPTaskThread(mpp_task_manager); etcd_client = Etcd::Client::create(cluster->pd_client, cluster_config); + + if (!cluster->pd_client->isMock()) + { + auto [kv, status] = etcd_client.get()->leader("/tici/metaservice/election"); + if (!status.ok()) + { + LOG_WARNING( + Logger::get(), + "Failed to get tici meta service leader from etcd, code={} msg={}", + magic_enum::enum_name(status.error_code()), + status.error_message()); + } + else + { + LOG_INFO(Logger::get(), "Get tici meta service leader from etcd: {}", kv.value()); + cluster->shard_cache = std::make_unique( + kv.value()); // Use the leader address as shard cache address + } + } } void TMTContext::initS3GCManager(const TiFlashRaftProxyHelper * proxy_helper) diff --git a/dbms/src/Storages/StorageTantivy.cpp b/dbms/src/Storages/StorageTantivy.cpp new file mode 100644 index 00000000000..84883892680 --- /dev/null +++ b/dbms/src/Storages/StorageTantivy.cpp @@ -0,0 +1,179 @@ +// Copyright 2025 PingCAP, Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace DB +{ + +StorageTantivy::StorageTantivy(Context & context_, const TiCIScan & tici_scan_) + : tici_scan(tici_scan_) + , context(context_) + , log(Logger::get(context_.getDAGContext()->log ? context_.getDAGContext()->log->identifier() : "")) +{} + +BlockInputStreams StorageTantivy::read( + const Names &, + const SelectQueryInfo &, + const Context &, + QueryProcessingStage::Enum &, + size_t, + unsigned int) +{ + return {}; +} + +void StorageTantivy::read( + PipelineExecutorContext & exec_status, + PipelineExecGroupBuilder & group_builder, + [[maybe_unused]] const Names & column_names, + [[maybe_unused]] const SelectQueryInfo & info, + const Context & context, + [[maybe_unused]] size_t max_block_size, + unsigned num_streams) +{ + NamesAndTypes return_columns; + if (tici_scan.isCount()) + { + RUNTIME_CHECK_MSG(tici_scan.getNamesAndTypes().size() == 1, "count search should return one column"); + return_columns = tici_scan.getNamesAndTypes(); + } + else + { + return_columns = genNamesAndTypesForTiCI(tici_scan.getReturnColumns(), "column"); + } + + RUNTIME_CHECK(local_shards_snapshot.has_value()); + auto shards_snapshot = std::move(*local_shards_snapshot); + local_shards_snapshot.reset(); + + auto tici_task_pool = std::make_shared( + log, + tici_scan.getKeyspaceID(), + tici_scan.getTableId(), + tici_scan.getIndexId(), + local_read, + return_columns, + tici_scan.getLimit(), + tici_scan.getSortColumnIds(), + tici_scan.getSortColumnAsc(), + context.getSettingsRef().read_tso, + tici_scan.getMatchExpr(), + tici_scan.isCount(), + context.getTimezoneInfo(), + std::move(shards_snapshot)); + + num_streams = std::max(1, std::min(num_streams, local_read.size())); + // local read + for (size_t i = 0; i < num_streams; ++i) + { + group_builder.addConcurrency( + std::make_unique(exec_status, log->identifier(), tici_task_pool, return_columns)); + } + if (!tici_scan.isCount()) + executeCastAfterTiCIScan(exec_status, group_builder); +} + +void StorageTantivy::executeCastAfterTiCIScan( + PipelineExecutorContext & exec_status, + PipelineExecGroupBuilder & group_builder) +{ + if (group_builder.concurrency() == 0) + return; + // execute timezone cast or duration cast if needed for local tici scan + DAGExpressionAnalyzer analyzer{group_builder.getCurrentHeader(), context}; + ExpressionActionsChain chain; + std::vector may_need_add_cast_column; + for (size_t i = 0; i < tici_scan.getReturnColumns().size(); ++i) + may_need_add_cast_column.push_back(true); + if (analyzer.appendExtraCastsAfterTiCI(chain, may_need_add_cast_column, tici_scan)) + { + ExpressionActionsPtr extra_cast = chain.getLastActions(); + assert(extra_cast); + chain.finalize(); + chain.clear(); + for (size_t i = 0; i < group_builder.concurrency(); ++i) + { + auto & builder = group_builder.getCurBuilder(i); + builder.appendTransformOp( + std::make_unique(exec_status, log->identifier(), extra_cast)); + } + } +} + +void StorageTantivy::splitRemoteReadAndLocalRead() +{ + const auto & all = tici_scan.getShardInfos().shard_info_list; + ShardInfoList local_shard_infos; + ShardInfoList remote_shard_infos; + ::rust::Vec<::Shard> shards; + auto index_id = tici_scan.getIndexId(); + for (const auto & shard_info : all) + { + shards.push_back(::Shard{ + .keyspace_id = tici_scan.getKeyspaceID(), + .index_id = index_id, + .shard_id = shard_info.shard_id, + .shard_epoch = shard_info.shard_epoch, + }); + } + + ::rust::Vec local_results; + local_shards_snapshot = check_shards_and_acquire_snapshot(shards, local_results); + RUNTIME_CHECK(local_shards_snapshot.has_value()); + RUNTIME_CHECK(local_results.size() == all.size()); + for (size_t i = 0; i < all.size(); ++i) + { + const auto is_local = local_results[i]; + if (is_local) + { + local_shard_infos.push_back(all[i]); + } + else + { + remote_shard_infos.push_back(all[i]); + } + } + remote_read = remote_shard_infos; + local_read = local_shard_infos; +} + +} // namespace DB diff --git a/dbms/src/Storages/StorageTantivy.h b/dbms/src/Storages/StorageTantivy.h new file mode 100644 index 00000000000..c019d9f0fa5 --- /dev/null +++ b/dbms/src/Storages/StorageTantivy.h @@ -0,0 +1,90 @@ +// Copyright 2022 PingCAP, Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wunused-parameter" +#include +#include +#pragma GCC diagnostic pop + +#include +#include +namespace DB +{ +class StorageTantivy : public IStorage +{ +public: + StorageTantivy(Context & context_, const TiCIScan & tici_scan_); + + std::string getName() const override { return "StorageTantivy"; } + + std::string getTableName() const override { return "StorageTantivy"; } + + BlockInputStreams read( + const Names & column_names, + const SelectQueryInfo & query_info, + const Context & context, + QueryProcessingStage::Enum & processed_stage, + size_t max_block_size, + unsigned num_streams) override; + + + void read( + PipelineExecutorContext & exec_status, + PipelineExecGroupBuilder & group_builder, + [[maybe_unused]] const Names & column_names, + [[maybe_unused]] const SelectQueryInfo & info, + [[maybe_unused]] const Context & context, + [[maybe_unused]] size_t max_block_size, + unsigned num_streams) override; + + void executeCastAfterTiCIScan(PipelineExecutorContext & exec_status, PipelineExecGroupBuilder & group_builder); + // Members will be transferred to DAGQueryBlockInterpreter after execute + std::unique_ptr analyzer; + + void splitRemoteReadAndLocalRead(); + ShardInfoList getRemoteShardInfos() const { return remote_read; } + +private: + const TiCIScan tici_scan; + [[maybe_unused]] Context & context; + LoggerPtr log; + + ShardInfoList remote_read; + ShardInfoList local_read; + std::optional> local_shards_snapshot; +}; +} // namespace DB diff --git a/dbms/src/Storages/Tantivy/TantivyInputStream.h b/dbms/src/Storages/Tantivy/TantivyInputStream.h new file mode 100644 index 00000000000..5cc61e1e98f --- /dev/null +++ b/dbms/src/Storages/Tantivy/TantivyInputStream.h @@ -0,0 +1,444 @@ +// Copyright 2025 PingCAP, Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace DB::TS +{ + +class TantivyInputStream : public IProfilingBlockInputStream +{ + static constexpr auto NAME = "TantivyInputStream"; + +public: + TantivyInputStream( + LoggerPtr log_, + UInt32 keyspace_id_, + Int64 table_id_, + Int64 index_id_, + ShardInfo query_shard_info_, + NamesAndTypes return_columns_, + UInt64 limit_, + std::vector sort_column_ids_, + std::vector sort_column_asc_, + UInt64 read_ts_, + ::Expr match_expr_, + bool is_count, + std::shared_ptr> shards_snapshot_) + : log(log_) + , keyspace_id(keyspace_id_) + , table_id(table_id_) + , index_id(index_id_) + , query_shard_info(query_shard_info_) + , return_columns(return_columns_) + , limit(limit_) + , sort_column_ids(sort_column_ids_) + , sort_column_asc(sort_column_asc_) + , read_ts(read_ts_) + , match_expr(match_expr_) + , is_count(is_count) + , shards_snapshot(std::move(shards_snapshot_)) + {} + + String getName() const override { return NAME; } + + Block getHeader() const override { return header; } + + Block readImpl() override + { + if (done) + { + return {}; + } + Block ret = readFromS3(is_count); + done = true; + return ret; + } + +protected: + Block readFromS3(bool is_count) + { + auto shard_info = query_shard_info; + LOG_DEBUG(log, "shard info: {}", shard_info.toString()); + auto key_ranges = getKeyRanges(shard_info.key_ranges); + + const auto shard = ::Shard{ + .keyspace_id = keyspace_id, + .index_id = index_id, + .shard_id = shard_info.shard_id, + .shard_epoch = shard_info.shard_epoch, + }; + + RUNTIME_CHECK(shards_snapshot != nullptr); + + Block res(return_columns); + if (is_count) + { + CountResult count_result = count(**shards_snapshot, shard, key_ranges, match_expr, read_ts); + + RUNTIME_CHECK_MSG(return_columns.size() == 1, "count search should return one column"); + auto & column = res.getByPosition(0).column->assumeMutableRef(); + column.insert(Int64(count_result.count)); + return res; + } + + auto return_fields = getFields(return_columns); + + rust::Vec tici_sort_column_names; + for (const auto & sort_column_id : sort_column_ids) + { + tici_sort_column_names.push_back(rust::String("column_" + std::to_string(sort_column_id))); + } + + rust::Vec tici_sort_column_asc; + for (const auto & asc : sort_column_asc) + { + tici_sort_column_asc.push_back(asc); + } + + SearchParam search_param{ + .limit = static_cast(limit), + .sort_field_names = std::move(tici_sort_column_names), + .is_asc = std::move(tici_sort_column_asc), + }; + + SearchResult search_result + = search(**shards_snapshot, shard, key_ranges, return_fields, match_expr, search_param, read_ts); + + const auto row_count = static_cast(search_result.row_count); + if (row_count == 0) + return res; + + auto name_to_pos = buildColumnPositionMap(return_columns); + std::vector filled(return_columns.size(), false); + + for (const auto & column_data : search_result.i64_columns) + { + installIntegerLikeColumn( + res, + name_to_pos, + filled, + return_columns, + column_data.col_name, + column_data.values, + column_data.null_map); + } + for (const auto & column_data : search_result.u64_columns) + { + installIntegerLikeColumn( + res, + name_to_pos, + filled, + return_columns, + column_data.col_name, + column_data.values, + column_data.null_map); + } + for (const auto & column_data : search_result.f64_columns) + { + installFloatColumn(res, name_to_pos, filled, return_columns, column_data); + } + for (const auto & column_data : search_result.bytes_columns) + { + installBytesColumn(res, name_to_pos, filled, return_columns, column_data); + } + + for (size_t i = 0; i < return_columns.size(); ++i) + { + if (!filled[i]) + fillDefaultColumn(res.getByPosition(i), row_count); + } + + return res; + } + +private: + Block header; + bool done = false; + LoggerPtr log; + UInt32 keyspace_id; + [[maybe_unused]] Int64 table_id; + Int64 index_id; + ShardInfo query_shard_info; + NamesAndTypes return_columns; + UInt64 limit; + std::vector sort_column_ids; + std::vector sort_column_asc; + UInt64 read_ts; + ::Expr match_expr; + bool is_count; + std::shared_ptr> shards_snapshot; + + static std::unordered_map buildColumnPositionMap(const NamesAndTypes & columns) + { + std::unordered_map positions; + positions.reserve(columns.size()); + for (size_t i = 0; i < columns.size(); ++i) + positions.emplace(columns[i].name, i); + return positions; + } + + static ColumnUInt8::MutablePtr buildNullMapColumn(size_t row_count, const rust::Vec<::std::uint8_t> & null_map) + { + auto null_map_column = ColumnUInt8::create(row_count, 0); + if (null_map.size() == 0) + return null_map_column; + + RUNTIME_CHECK_MSG( + null_map.size() == row_count, + "null map size mismatch, expect {}, got {}", + row_count, + null_map.size()); + auto & dst = null_map_column->getData(); + if (row_count != 0) + std::memcpy(dst.data(), null_map.data(), row_count * sizeof(UInt8)); + return null_map_column; + } + + static void fillDefaultColumn(ColumnWithTypeAndName & column, size_t row_count) + { + auto mutable_column = column.type->createColumn(); + mutable_column->insertManyDefaults(row_count); + column.column = std::move(mutable_column); + } + + template + static MutableColumnPtr buildNumericColumn( + const NameAndTypePair & name_and_type, + const rust::Vec & values, + const rust::Vec<::std::uint8_t> & null_map) + { + const auto row_count = values.size(); + auto nested_column = ColumnVector::create(row_count); + auto & data = nested_column->getData(); + const bool has_null_map = null_map.size() != 0; + + if (has_null_map) + { + RUNTIME_CHECK_MSG( + null_map.size() == row_count, + "null map size mismatch for column {}, expect {}, got {}", + name_and_type.name, + row_count, + null_map.size()); + } + + if constexpr (std::is_same_v) + { + if (row_count != 0) + std::memcpy(data.data(), values.data(), row_count * sizeof(TargetType)); + } + else + { + for (size_t i = 0; i < row_count; ++i) + data[i] = static_cast(values[i]); + } + + if (name_and_type.type->isNullable()) + return ColumnNullable::create(std::move(nested_column), buildNullMapColumn(row_count, null_map)); + return nested_column; + } + + template + static MutableColumnPtr buildIntegerLikeColumn( + const NameAndTypePair & name_and_type, + const rust::Vec & values, + const rust::Vec<::std::uint8_t> & null_map) + { + const auto nested_type = removeNullable(name_and_type.type); + if (typeid_cast(nested_type.get())) + return buildNumericColumn(name_and_type, values, null_map); + if (typeid_cast(nested_type.get())) + return buildNumericColumn(name_and_type, values, null_map); + if (typeid_cast(nested_type.get())) + return buildNumericColumn(name_and_type, values, null_map); + if (typeid_cast(nested_type.get())) + return buildNumericColumn(name_and_type, values, null_map); + if (typeid_cast(nested_type.get())) + return buildNumericColumn(name_and_type, values, null_map); + if (typeid_cast(nested_type.get())) + return buildNumericColumn(name_and_type, values, null_map); + if (typeid_cast(nested_type.get())) + return buildNumericColumn(name_and_type, values, null_map); + if (typeid_cast(nested_type.get())) + return buildNumericColumn(name_and_type, values, null_map); + if (typeid_cast(nested_type.get())) + return buildNumericColumn(name_and_type, values, null_map); + if (typeid_cast(nested_type.get())) + return buildNumericColumn(name_and_type, values, null_map); + if (typeid_cast(nested_type.get())) + return buildNumericColumn(name_and_type, values, null_map); + if (typeid_cast(nested_type.get())) + return buildNumericColumn(name_and_type, values, null_map); + + throw Exception( + ErrorCodes::LOGICAL_ERROR, + "unsupported integer-like target type {} for column {}", + nested_type->getName(), + name_and_type.name); + } + + static MutableColumnPtr buildFloatColumn( + const NameAndTypePair & name_and_type, + const rust::Vec & values, + const rust::Vec<::std::uint8_t> & null_map) + { + const auto nested_type = removeNullable(name_and_type.type); + if (typeid_cast(nested_type.get())) + return buildNumericColumn(name_and_type, values, null_map); + if (typeid_cast(nested_type.get())) + return buildNumericColumn(name_and_type, values, null_map); + + throw Exception( + ErrorCodes::LOGICAL_ERROR, + "unsupported float target type {} for column {}", + nested_type->getName(), + name_and_type.name); + } + + static MutableColumnPtr buildBytesColumn(const NameAndTypePair & name_and_type, const BytesColumnData & column_data) + { + const auto row_count = column_data.offsets.size(); + if (column_data.offsets.size() != 0) + { + RUNTIME_CHECK_MSG( + column_data.offsets[column_data.offsets.size() - 1] == column_data.chars.size(), + "string offsets and chars size mismatch for column {}", + name_and_type.name); + } + + if (removeNullable(name_and_type.type)->isString()) + { + auto nested_column = ColumnString::create(); + auto & chars = nested_column->getChars(); + auto & offsets = nested_column->getOffsets(); + + chars.resize(column_data.chars.size()); + if (!column_data.chars.empty()) + std::memcpy(chars.data(), column_data.chars.data(), column_data.chars.size()); + + offsets.resize(row_count); + for (size_t i = 0; i < row_count; ++i) + offsets[i] = static_cast(column_data.offsets[i]); + + if (name_and_type.type->isNullable()) + return ColumnNullable::create( + std::move(nested_column), + buildNullMapColumn(row_count, column_data.null_map)); + return nested_column; + } + throw Exception( + ErrorCodes::LOGICAL_ERROR, + "unsupported bytes target type {} for column {}", + name_and_type.type->getName(), + name_and_type.name); + } + + template + static void installIntegerLikeColumn( + Block & res, + const std::unordered_map & name_to_pos, + std::vector & filled, + const NamesAndTypes & columns, + const rust::String & col_name, + const rust::Vec & values, + const rust::Vec<::std::uint8_t> & null_map) + { + const String column_name(col_name); + const auto it = name_to_pos.find(column_name); + RUNTIME_CHECK_MSG(it != name_to_pos.end(), "unexpected column {} returned from TiCI", column_name); + + const auto pos = it->second; + RUNTIME_CHECK_MSG(!filled[pos], "duplicate column {} returned from TiCI", column_name); + auto & result_column = res.getByPosition(pos); + result_column.column = buildIntegerLikeColumn(columns[pos], values, null_map); + filled[pos] = true; + } + + static void installFloatColumn( + Block & res, + const std::unordered_map & name_to_pos, + std::vector & filled, + const NamesAndTypes & columns, + const F64ColumnData & column_data) + { + const String column_name(column_data.col_name); + const auto it = name_to_pos.find(column_name); + RUNTIME_CHECK_MSG(it != name_to_pos.end(), "unexpected column {} returned from TiCI", column_name); + + const auto pos = it->second; + RUNTIME_CHECK_MSG(!filled[pos], "duplicate column {} returned from TiCI", column_name); + auto & result_column = res.getByPosition(pos); + result_column.column = buildFloatColumn(columns[pos], column_data.values, column_data.null_map); + filled[pos] = true; + } + + static void installBytesColumn( + Block & res, + const std::unordered_map & name_to_pos, + std::vector & filled, + const NamesAndTypes & columns, + const BytesColumnData & column_data) + { + const String column_name(column_data.col_name); + const auto it = name_to_pos.find(column_name); + RUNTIME_CHECK_MSG(it != name_to_pos.end(), "unexpected column {} returned from TiCI", column_name); + + const auto pos = it->second; + RUNTIME_CHECK_MSG(!filled[pos], "duplicate column {} returned from TiCI", column_name); + auto & result_column = res.getByPosition(pos); + result_column.column = buildBytesColumn(columns[pos], column_data); + filled[pos] = true; + } + + static rust::Vec getFields(NamesAndTypes & columns) + { + rust::Vec fields; + for (auto & name_and_type : columns) + { + fields.push_back(name_and_type.name); + } + return fields; + } +}; +} // namespace DB::TS diff --git a/dbms/src/Storages/Tantivy/TiCIReadTaskPool.h b/dbms/src/Storages/Tantivy/TiCIReadTaskPool.h new file mode 100644 index 00000000000..11174dd45b1 --- /dev/null +++ b/dbms/src/Storages/Tantivy/TiCIReadTaskPool.h @@ -0,0 +1,178 @@ +// Copyright 2025 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 + +namespace DB::TS +{ +struct TiCIReadTask +{ +public: + explicit TiCIReadTask(const ShardInfo & shard_info_) + : shard_info(shard_info_) + {} + + void initInputStream( + LoggerPtr log_, + UInt32 keyspace_id_, + Int64 table_id_, + Int64 index_id_, + ShardInfo query_shard_info_, + NamesAndTypes return_columns_, + UInt64 limit_, + std::vector sort_column_ids_, + std::vector sort_column_asc_, + UInt64 read_ts_, + ::Expr match_expr_, + bool is_count, + const std::shared_ptr> & shards_snapshot_) + { + input_stream = std::make_shared( + log_, + keyspace_id_, + table_id_, + index_id_, + query_shard_info_, + return_columns_, + limit_, + sort_column_ids_, + sort_column_asc_, + read_ts_, + match_expr_, + is_count, + shards_snapshot_); + } + + bool isInitialized() const { return input_stream != nullptr; } + + BlockInputStreamPtr getInputStream() const + { + RUNTIME_CHECK(input_stream != nullptr); + return input_stream; + } + ShardInfo getShardInfo() const { return shard_info; } + +private: + ShardInfo shard_info; + BlockInputStreamPtr input_stream; +}; + +using TiCIReadTaskPtr = std::shared_ptr; + +struct TiCIReadTaskPool +{ +public: + using TiCIReadTasks = std::vector>; + + TiCIReadTaskPool( + LoggerPtr log_, + UInt32 keyspace_id_, + Int64 table_id_, + Int64 index_id_, + const ShardInfoList & shard_infos, + NamesAndTypes return_columns_, + UInt64 limit_, + std::vector sort_column_ids_, + std::vector sort_column_asc_, + UInt64 read_ts_, + google::protobuf::RepeatedPtrField match_expr_, + bool is_count, + const TimezoneInfo & timezone_info_, + rust::Box shards_snapshot_) + : log(log_) + , keyspace_id(keyspace_id_) + , table_id(table_id_) + , index_id(index_id_) + , return_columns(return_columns_) + , limit(limit_) + , sort_column_ids(sort_column_ids_) + , sort_column_asc(sort_column_asc_) + , read_ts(read_ts_) + , is_count(is_count) + , shards_snapshot(std::make_shared>(std::move(shards_snapshot_))) + { + for (const auto & shard_info : shard_infos) + { + tasks.emplace_back(std::make_shared(shard_info)); + } + FmtBuffer buf; + buf.joinStr( + return_columns.begin(), + return_columns.end(), + [](const auto & nt, FmtBuffer & fb) { fb.fmtAppend("{}:{}", nt.name, nt.type->getName()); }, + ", "); + auto [expr, cids] = TS::tipbToTiCIExpr(match_expr_, timezone_info_); + match_expr = std::move(expr); + LOG_DEBUG(log, "columns: [{}], match columns: {}", buf.toString(), cids); + } + + TiCIReadTaskPtr getNextTask() + { + std::lock_guard lock(mutex); + if (tasks.empty()) + return nullptr; + TiCIReadTaskPtr task = tasks.back(); + tasks.pop_back(); + return task; + } + + BlockInputStreamPtr buildInputStream(TiCIReadTaskPtr & task) + { + RUNTIME_CHECK(task != nullptr); + if (!task->isInitialized()) + { + task->initInputStream( + log, + keyspace_id, + table_id, + index_id, + task->getShardInfo(), + return_columns, + limit, + sort_column_ids, + sort_column_asc, + read_ts, + match_expr, + is_count, + shards_snapshot); + } + return task->getInputStream(); + } + +private: + mutable std::mutex mutex; + TiCIReadTasks tasks; + + LoggerPtr log; + UInt32 keyspace_id; + Int64 table_id; + Int64 index_id; + NamesAndTypes return_columns; + UInt64 limit; + std::vector sort_column_ids; + std::vector sort_column_asc; + UInt64 read_ts; + ::Expr match_expr; + bool is_count; + std::shared_ptr> shards_snapshot; +}; + +using TiCIReadTaskPoolPtr = std::shared_ptr; + +} // namespace DB::TS diff --git a/dbms/src/Storages/Tantivy/TiCIRequestUtils.h b/dbms/src/Storages/Tantivy/TiCIRequestUtils.h new file mode 100644 index 00000000000..07553dea0b4 --- /dev/null +++ b/dbms/src/Storages/Tantivy/TiCIRequestUtils.h @@ -0,0 +1,256 @@ +// 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 +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace DB::TS +{ + +// Convert literal value from timezone specified in cop request to UTC in-place. +inline UInt64 convertPackedU64WithTimezone(UInt64 from_time, const TimezoneInfo & timezone_info) +{ + static const auto & time_zone_utc = DateLUT::instance("UTC"); + UInt64 result_time = from_time; + if (timezone_info.is_name_based) + convertTimeZone(from_time, result_time, *timezone_info.timezone, time_zone_utc); + else if (timezone_info.timezone_offset != 0) + convertTimeZoneByOffset(from_time, result_time, false, timezone_info.timezone_offset); + return result_time; +} + +inline rust::Vec<::Range> getKeyRanges(const ShardInfo::KeyRanges & key_ranges) +{ + rust::Vec<::Range> res; + for (const auto & range : key_ranges) + { + rust::Slice<::std::uint8_t const> start( + reinterpret_cast(range.start().c_str()), + range.start().size()); + rust::Slice<::std::uint8_t const> end( + reinterpret_cast(range.end().c_str()), + range.end().size()); + res.push_back({ + .start = std::move(start), + .end = std::move(end), + }); + } + return res; +} + +inline std::tuple<::Expr, std::vector> tipbToTiCIExpr( + const tipb::Expr & expr, + const TimezoneInfo & timezone_info) +{ + ::Expr ret; + switch (expr.tp()) + { + case tipb::ExprType::ScalarFunc: + { + std::vector children_cids; + ret.tp = tipb::ExprType::ScalarFunc; + for (const auto & child : expr.children()) + { + auto [child_expr, child_cids] = tipbToTiCIExpr(child, timezone_info); + ret.children.push_back(child_expr); + children_cids.insert(children_cids.end(), child_cids.begin(), child_cids.end()); + } + switch (expr.sig()) + { + case tipb::ScalarFuncSig::FTSMatchWord: + case tipb::ScalarFuncSig::FTSMatchPrefix: + case tipb::ScalarFuncSig::FTSMatchPhrase: + case tipb::ScalarFuncSig::LogicalAnd: + case tipb::ScalarFuncSig::LogicalOr: + case tipb::ScalarFuncSig::UnaryNotInt: + case tipb::ScalarFuncSig::UnaryNotReal: + case tipb::ScalarFuncSig::EQInt: + case tipb::ScalarFuncSig::NEInt: + case tipb::ScalarFuncSig::LTInt: + case tipb::ScalarFuncSig::LEInt: + case tipb::ScalarFuncSig::GTInt: + case tipb::ScalarFuncSig::GEInt: + case tipb::ScalarFuncSig::EQString: + case tipb::ScalarFuncSig::NEString: + case tipb::ScalarFuncSig::LTString: + case tipb::ScalarFuncSig::LEString: + case tipb::ScalarFuncSig::GTString: + case tipb::ScalarFuncSig::GEString: + case tipb::ScalarFuncSig::EQReal: + case tipb::ScalarFuncSig::NEReal: + case tipb::ScalarFuncSig::LTReal: + case tipb::ScalarFuncSig::LEReal: + case tipb::ScalarFuncSig::GTReal: + case tipb::ScalarFuncSig::GEReal: + case tipb::ScalarFuncSig::EQDecimal: + case tipb::ScalarFuncSig::NEDecimal: + case tipb::ScalarFuncSig::LTDecimal: + case tipb::ScalarFuncSig::LEDecimal: + case tipb::ScalarFuncSig::GTDecimal: + case tipb::ScalarFuncSig::GEDecimal: + case tipb::ScalarFuncSig::InInt: + case tipb::ScalarFuncSig::InString: + case tipb::ScalarFuncSig::InReal: + case tipb::ScalarFuncSig::InDecimal: + ret.sig = expr.sig(); + break; + case tipb::ScalarFuncSig::EQTime: + case tipb::ScalarFuncSig::NETime: + case tipb::ScalarFuncSig::LTTime: + case tipb::ScalarFuncSig::LETime: + case tipb::ScalarFuncSig::GTTime: + case tipb::ScalarFuncSig::GETime: + { + ret.sig = expr.sig(); + size_t col_idx = 0, val_idx = 1; + if (isColumnExpr(expr.children(1))) + std::swap(col_idx, val_idx); + if (expr.children(col_idx).field_type().tp() == TiDB::TypeTimestamp) + { + const auto & child_expr = expr.children(val_idx); + if (isLiteralExpr(child_expr)) + { + UInt64 val = decodeDAGUInt64(child_expr.val()); + val = convertPackedU64WithTimezone(val, timezone_info); + WriteBufferFromOwnString ss; + encodeDAGUInt64(val, ss); + ret.children[val_idx].val.clear(); + auto str = ss.releaseStr(); + std::copy(str.begin(), str.end(), std::back_inserter(ret.children[val_idx].val)); + } + } + break; + } + case tipb::ScalarFuncSig::InTime: + { + ret.sig = expr.sig(); + if (expr.children(0).field_type().tp() == TiDB::TypeTimestamp) + { + for (int val_idx = 1; val_idx < expr.children_size(); ++val_idx) + { + const auto & child_expr = expr.children(val_idx); + if (isLiteralExpr(child_expr)) + { + UInt64 val = decodeDAGUInt64(child_expr.val()); + val = convertPackedU64WithTimezone(val, timezone_info); + WriteBufferFromOwnString ss; + encodeDAGUInt64(val, ss); + ret.children[val_idx].val.clear(); + auto str = ss.releaseStr(); + std::copy(str.begin(), str.end(), std::back_inserter(ret.children[val_idx].val)); + } + else + { + throw TiFlashException("InTime only support literal values", Errors::Coprocessor::BadRequest); + } + } + } + break; + } + default: + throw std::runtime_error(fmt::format( + "Unsupported expression sig: tp={}, sig={}, expr={}", + static_cast(expr.tp()), + static_cast(expr.sig()), + expr.DebugString())); + } + + return {ret, children_cids}; + } + case tipb::ExprType::ColumnRef: + { + ret.tp = expr.tp(); + auto id = decodeDAGInt64(expr.val()); + auto str = fmt::format("column_{}", id); + std::copy(str.begin(), str.end(), std::back_inserter(ret.val)); + return {ret, {id}}; + } + case tipb::ExprType::String: + case tipb::ExprType::Int64: + case tipb::ExprType::Uint64: + case tipb::ExprType::Float32: + case tipb::ExprType::Float64: + case tipb::ExprType::MysqlTime: + { + ret.tp = expr.tp(); + std::copy(expr.val().begin(), expr.val().end(), std::back_inserter(ret.val)); + return {ret, {}}; + } + case tipb::ExprType::MysqlDecimal: + { + ret.tp = expr.tp(); + auto field = decodeDAGDecimal(expr.val()); + String str; + if (field.getType() == Field::Types::Decimal32) + str = field.get>().toString(); + else if (field.getType() == Field::Types::Decimal64) + str = field.get>().toString(); + else if (field.getType() == Field::Types::Decimal128) + str = field.get>().toString(); + else if (field.getType() == Field::Types::Decimal256) + str = field.get>().toString(); + else + throw TiFlashException("Not decimal literal: " + expr.DebugString(), Errors::Coprocessor::BadRequest); + std::copy(str.begin(), str.end(), std::back_inserter(ret.val)); + return {ret, {}}; + } + default: + throw std::runtime_error(fmt::format( + "Unsupported expression type: tp={}, expr={}", + static_cast(expr.tp()), + expr.DebugString())); + } +} + +inline std::tuple<::Expr, std::vector> tipbToTiCIExpr( + const google::protobuf::RepeatedPtrField & exprs, + const TimezoneInfo & tz) +{ + if (exprs.empty()) + throw std::runtime_error("Empty match expression"); + auto [ret, cids] = tipbToTiCIExpr(exprs[0], tz); + for (auto i = 1; i < exprs.size(); ++i) + { + auto [child_expr, child_cids] = tipbToTiCIExpr(exprs[i], tz); + ret = { + .tp = tipb::ExprType::ScalarFunc, + .children = {ret, child_expr}, + .sig = tipb::ScalarFuncSig::LogicalAnd, + }; + cids.insert(cids.end(), child_cids.begin(), child_cids.end()); + } + return {ret, cids}; +} + +} // namespace DB::TS diff --git a/dbms/src/Storages/tests/gtest_filter_parser.cpp b/dbms/src/Storages/tests/gtest_filter_parser.cpp index bc2882b1f47..63733ee39b6 100644 --- a/dbms/src/Storages/tests/gtest_filter_parser.cpp +++ b/dbms/src/Storages/tests/gtest_filter_parser.cpp @@ -84,7 +84,7 @@ DM::RSOperatorPtr FilterParserTest::generateRsOperator( [&](const String &, const String &) { return table_info; }, getDAGProperties("")); auto & dag_request = *query_tasks[0].dag_request; - DAGContext dag_context(dag_request, {}, NullspaceID, "", DAGRequestKind::Cop, "", 0, "", log); + DAGContext dag_context(dag_request, {}, QueryShardInfos(), NullspaceID, "", DAGRequestKind::Cop, "", 0, "", log); ctx->setDAGContext(&dag_context); // Don't care about regions information in this test google::protobuf::RepeatedPtrField conditions; diff --git a/dbms/src/Storages/tests/gtests_parse_push_down_filter.cpp b/dbms/src/Storages/tests/gtests_parse_push_down_filter.cpp index 34426b335f3..518277067b0 100644 --- a/dbms/src/Storages/tests/gtests_parse_push_down_filter.cpp +++ b/dbms/src/Storages/tests/gtests_parse_push_down_filter.cpp @@ -80,7 +80,7 @@ DM::PushDownExecutorPtr generatePushDownExecutor( getDAGProperties("")); auto & dag_request = *query_tasks[0].dag_request; auto log = Logger::get(); - DAGContext dag_context(dag_request, {}, NullspaceID, "", DAGRequestKind::Cop, "", 0, "", log); + DAGContext dag_context(dag_request, {}, QueryShardInfos(), NullspaceID, "", DAGRequestKind::Cop, "", 0, "", log); ctx.setDAGContext(&dag_context); // Don't care about regions information in this test google::protobuf::RepeatedPtrField empty_condition; diff --git a/dbms/src/TestUtils/MPPTaskTestUtils.cpp b/dbms/src/TestUtils/MPPTaskTestUtils.cpp index a8cdecacb79..9e5db6027da 100644 --- a/dbms/src/TestUtils/MPPTaskTestUtils.cpp +++ b/dbms/src/TestUtils/MPPTaskTestUtils.cpp @@ -175,7 +175,7 @@ ColumnsWithTypeAndName MPPTaskTestUtils::executeCoprocessorTask(std::shared_ptr< auto * data = req->mutable_data(); dag_request->AppendToString(data); - DAGContext dag_context(*dag_request, {}, NullspaceID, "", DAGRequestKind::Cop, "", 0, "", Logger::get()); + DAGContext dag_context(*dag_request, {}, {}, NullspaceID, "", DAGRequestKind::Cop, "", 0, "", Logger::get()); TiFlashTestEnv::getGlobalContext(test_meta.context_idx).setDAGContext(&dag_context); TiFlashTestEnv::getGlobalContext(test_meta.context_idx).setCopTest(); diff --git a/rust-toolchain.toml b/rust-toolchain.toml deleted file mode 120000 index 82e224aad22..00000000000 --- a/rust-toolchain.toml +++ /dev/null @@ -1 +0,0 @@ -contrib/tiflash-proxy/rust-toolchain.toml \ No newline at end of file diff --git a/rust-toolchain.toml b/rust-toolchain.toml new file mode 100644 index 00000000000..14ba487c8e7 --- /dev/null +++ b/rust-toolchain.toml @@ -0,0 +1,4 @@ +[toolchain] +channel = "nightly-2025-02-28" +components = ["rustfmt", "clippy", "rust-src", "rust-analyzer"] +profile = "minimal" diff --git a/tests/docker/util.sh b/tests/docker/util.sh index 7247661a743..efb0d0e0a06 100644 --- a/tests/docker/util.sh +++ b/tests/docker/util.sh @@ -182,6 +182,66 @@ function prepare_next_gen_columnar_data_dirs() { ./data/tiflash-cn0 ./log/tiflash-cn0 } +# feature/fts TiFlash is 8.5.x, while CI may pair it with master PD/TiKV/TiDB. +# Pin PD's compatibility gate before TiFlash registers itself to PD. +function set_pd_cluster_version_for_tiflash() { + local cluster_version="${TIFLASH_TEST_PD_CLUSTER_VERSION:-8.5.6}" + if [[ -z "${cluster_version}" || "${cluster_version}" == "0" || "${cluster_version}" == "false" ]]; then + return + fi + + cluster_version="${cluster_version#v}" + if [[ ! "${cluster_version}" =~ ^[0-9]+\.[0-9]+\.[0-9]+([.-][0-9A-Za-z.-]+)?$ ]]; then + echo "Invalid TIFLASH_TEST_PD_CLUSTER_VERSION: ${TIFLASH_TEST_PD_CLUSTER_VERSION}" >&2 + exit 1 + fi + + local timeout="${TIFLASH_TEST_PD_CLUSTER_VERSION_TIMEOUT:-60}" + local failed='true' + + echo "=> set PD cluster-version to ${cluster_version} for TiFlash test" + + for (( i = 0; i < "${timeout}"; i++ )); do + local output='' + if output=$(${COMPOSE} "$@" exec -T pd0 /pd-ctl -u http://127.0.0.1:2379 config set cluster-version "${cluster_version}" 2>&1); then + echo "${output}" + if echo "${output}" | grep -q "Success!"; then + local current_version='' + current_version=$(${COMPOSE} "$@" exec -T pd0 /pd-ctl -u http://127.0.0.1:2379 config show cluster-version 2>&1 || true) + echo "${current_version}" + if echo "${current_version}" | grep -Fq "${cluster_version}"; then + failed='false' + break + fi + fi + else + echo "${output}" + fi + + if [ $((${i} % 10)) = 0 ] && [ ${i} -ge 10 ]; then + echo " #${i} waiting for PD to set cluster-version" + fi + + sleep 1 + done + + if [ "${failed}" == 'true' ]; then + echo " can not set PD cluster-version" >&2 + exit 1 + fi +} + +function start_cluster_with_tiflash() { + local compose_args=() + for compose_file in "$@"; do + compose_args+=("-f" "${compose_file}") + done + + ${COMPOSE} "${compose_args[@]}" up -d pd0 tikv0 tidb0 + set_pd_cluster_version_for_tiflash "${compose_args[@]}" + ${COMPOSE} "${compose_args[@]}" up -d tiflash0 +} + function check_env() { local cur_dir=$(pwd) local tiflash_bin_dir="${LOCAL_TiFLASH_BIN_DIR:-${cur_dir}/../.build/tiflash}" @@ -318,6 +378,8 @@ export -f set_branch export -f clean_data_log export -f prepare_next_gen_data_dirs export -f prepare_next_gen_columnar_data_dirs +export -f set_pd_cluster_version_for_tiflash +export -f start_cluster_with_tiflash export -f check_env export -f check_docker_compose export -f validate_local_binary diff --git a/tests/fullstack-test/run.sh b/tests/fullstack-test/run.sh index fbccef881c6..f9430c6bef0 100755 --- a/tests/fullstack-test/run.sh +++ b/tests/fullstack-test/run.sh @@ -35,7 +35,7 @@ ${COMPOSE} -f cluster.yaml -f tiflash-dt.yaml down clean_data_log # run fullstack-tests -${COMPOSE} -f cluster.yaml -f tiflash-dt.yaml up -d +start_cluster_with_tiflash cluster.yaml tiflash-dt.yaml wait_env echo "PD version:" @@ -49,7 +49,7 @@ ${COMPOSE} -f cluster.yaml -f tiflash-dt.yaml exec -T tiflash0 bash -c 'cd /test ${COMPOSE} -f cluster.yaml -f tiflash-dt.yaml down clean_data_log -${COMPOSE} -f cluster.yaml -f tiflash-dt-sync-grpc.yaml up -d +start_cluster_with_tiflash cluster.yaml tiflash-dt-sync-grpc.yaml wait_env ${COMPOSE} -f cluster.yaml -f tiflash-dt-sync-grpc.yaml exec -T tiflash0 bash -c 'cd /tests ; ./run-test.sh fullstack-test/mpp' ${COMPOSE} -f cluster.yaml -f tiflash-dt-sync-grpc.yaml down diff --git a/tests/fullstack-test2/run.sh b/tests/fullstack-test2/run.sh index 77041ca62d9..69a0db811a4 100755 --- a/tests/fullstack-test2/run.sh +++ b/tests/fullstack-test2/run.sh @@ -36,7 +36,7 @@ ${COMPOSE} -f cluster.yaml -f tiflash-dt.yaml down clean_data_log # FIXME: now vector does not support run with encryption-at-rest enabled -${COMPOSE} -f cluster.yaml -f tiflash-dt-disable-encrypt.yaml up -d +start_cluster_with_tiflash cluster.yaml tiflash-dt-disable-encrypt.yaml wait_env echo "PD version:" @@ -51,13 +51,13 @@ ${COMPOSE} -f cluster.yaml -f tiflash-dt-disable-encrypt.yaml down clean_data_log -${COMPOSE} -f cluster.yaml -f tiflash-dt.yaml up -d +start_cluster_with_tiflash cluster.yaml tiflash-dt.yaml wait_env ${COMPOSE} -f cluster.yaml -f tiflash-dt.yaml exec -T tiflash0 bash -c 'cd /tests ; ./run-test.sh fullstack-test2 true' ${COMPOSE} -f cluster.yaml -f tiflash-dt.yaml down clean_data_log -${COMPOSE} -f cluster.yaml -f tiflash-dt-disable-local-tunnel.yaml up -d +start_cluster_with_tiflash cluster.yaml tiflash-dt-disable-local-tunnel.yaml wait_env ${COMPOSE} -f cluster.yaml -f tiflash-dt-disable-local-tunnel.yaml exec -T tiflash0 bash -c 'cd /tests ; ./run-test.sh fullstack-test/mpp' ${COMPOSE} -f cluster.yaml -f tiflash-dt-disable-local-tunnel.yaml down diff --git a/tests/tidb-ci/run.sh b/tests/tidb-ci/run.sh index 5b917698cca..96b4cbb08a2 100755 --- a/tests/tidb-ci/run.sh +++ b/tests/tidb-ci/run.sh @@ -35,7 +35,7 @@ echo "Running fullstack test on classic TiFlash" ${COMPOSE} -f cluster.yaml -f tiflash-dt.yaml down clean_data_log -${COMPOSE} -f cluster.yaml -f tiflash-dt.yaml up -d +start_cluster_with_tiflash cluster.yaml tiflash-dt.yaml wait_env echo "PD version:" @@ -54,7 +54,7 @@ clean_data_log ${COMPOSE} -f cluster_new_collation.yaml -f tiflash-dt.yaml down clean_data_log -${COMPOSE} -f cluster_new_collation.yaml -f tiflash-dt.yaml up -d +start_cluster_with_tiflash cluster_new_collation.yaml tiflash-dt.yaml wait_env ${COMPOSE} -f cluster_new_collation.yaml -f tiflash-dt.yaml exec -T tiflash0 bash -c 'cd /tests ; ./run-test.sh tidb-ci/new_collation_fullstack' @@ -65,7 +65,7 @@ clean_data_log ${COMPOSE} -f cluster_disable_new_collation.yaml -f tiflash-dt.yaml down clean_data_log -${COMPOSE} -f cluster_disable_new_collation.yaml -f tiflash-dt.yaml up -d +start_cluster_with_tiflash cluster_disable_new_collation.yaml tiflash-dt.yaml wait_env ${COMPOSE} -f cluster_disable_new_collation.yaml -f tiflash-dt.yaml exec -T tiflash0 bash -c 'cd /tests ; ./run-test.sh tidb-ci/disable_new_collation_fullstack' @@ -73,7 +73,7 @@ ${COMPOSE} -f cluster_disable_new_collation.yaml -f tiflash-dt.yaml down clean_data_log # run force_enable_lm tests -${COMPOSE} -f cluster.yaml -f tiflash-dt-force-enable-lm.yaml up -d +start_cluster_with_tiflash cluster.yaml tiflash-dt-force-enable-lm.yaml wait_env ${COMPOSE} -f cluster.yaml -f tiflash-dt-force-enable-lm.yaml exec -T tiflash0 bash -c 'cd /tests ; ./run-test.sh tidb-ci/force_enable_lm' @@ -81,7 +81,7 @@ ${COMPOSE} -f cluster.yaml -f tiflash-dt-force-enable-lm.yaml down clean_data_log # run lightweight compression tests -${COMPOSE} -f cluster.yaml -f tiflash-dt-lightweight-compression.yaml up -d +start_cluster_with_tiflash cluster.yaml tiflash-dt-lightweight-compression.yaml wait_env ${COMPOSE} -f cluster.yaml -f tiflash-dt-lightweight-compression.yaml exec -T tiflash0 bash -c 'cd /tests ; ./run-test.sh tidb-ci/lightweight_compression'