tici: preview port feature/fts onto master

.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

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)

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 "$<LINK_LIBRARY:WHOLE_ARCHIVE,tici_search_lib_static>")
+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)

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

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",                                                                                     \

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<DAGRequestKind::Cop> 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<DAGRequestKind::Cop>
         driver(context, properties.start_ts, DEFAULT_UNSPECIFIED_SCHEMA_VERSION, &dag_response, true);

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,

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<TaskOperatorSpillContexts>())
     , flags(dag_request->flags())

dbms/src/Flash/Coprocessor/DAGContext.h

@@ -36,6 +36,7 @@
 #include <Flash/Coprocessor/DAGRequest.h>
 #include <Flash/Coprocessor/FineGrainedShuffle.h>
 #include <Flash/Coprocessor/RuntimeFilterMgr.h>
+#include <Flash/Coprocessor/ShardInfo.h>
 #include <Flash/Coprocessor/TablesRegionsInfo.h>
 #include <Flash/Executor/toRU.h>
 #include <Flash/Mpp/MPPTaskId.h>
@@ -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<std::unordered_set<UInt64>> 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;
 

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<BatchCopStreamWriter>(batch_cop_writer);
         TiDB::TiDBCollators collators;
         auto response_writer = std::make_unique<StreamingDAGResponseWriter<BatchCopStreamWriterPtr>>(

dbms/src/Flash/Coprocessor/DAGExpressionAnalyzer.cpp

@@ -1308,6 +1308,32 @@ bool DAGExpressionAnalyzer::appendExtraCastsAfterTS(
     return true;
 }
 
+bool DAGExpressionAnalyzer::appendExtraCastsAfterTiCI(
+    ExpressionActionsChain & chain,
+    const std::vector<UInt8> & 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,

dbms/src/Flash/Coprocessor/DAGExpressionAnalyzer.h

@@ -22,6 +22,7 @@
 #include <Flash/Coprocessor/DAGSet.h>
 #include <Flash/Coprocessor/DAGUtils.h>
 #include <Flash/Coprocessor/RuntimeFilterMgr.h>
+#include <Flash/Coprocessor/TiCIScan.h>
 #include <Flash/Coprocessor/TiDBTableScan.h>
 #include <Interpreters/AggregateDescription.h>
 #include <Interpreters/ExpressionActions.h>
@@ -114,6 +115,11 @@ class DAGExpressionAnalyzer : private boost::noncopyable
         const std::vector<UInt8> & may_need_add_cast_column,
         const TiDBTableScan & table_scan);
 
+    bool appendExtraCastsAfterTiCI(
+        ExpressionActionsChain & chain,
+        const std::vector<UInt8> & may_need_add_cast_column,
+        const TiCIScan & tici_scan);
+
     /// return true if some actions is needed
     bool appendJoinKeyAndJoinFilters(
         ExpressionActionsChain & chain,

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<DM::ColumnDefinesPtr, int, std::vector<std::tuple<UInt64, String, DataTypePtr>>> genColumnDefinesForDisaggregatedRead(
     const TiDBTableScan & table_scan)
 {

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);