Update Kuzu from branch (#51)

Co-authored-by: mewim <14037726+mewim@users.noreply.github.com>

Author: github-actions[bot]

Sources/cxx-kuzu/kuzu/extension/fts/src/function/create_fts_index.cpp

@@ -266,9 +266,10 @@ void LenCompute::vertexCompute(const graph::VertexScanState::Chunk& chunk) {
 static offset_t tableFunc(const TableFuncInput& input, TableFuncOutput&) {
     auto& bindData = *input.bindData->constPtrCast<CreateFTSBindData>();
     auto& context = *input.context;
+    auto transaction = context.clientContext->getTransaction();
+    auto catalog = context.clientContext->getCatalog();
     auto docTableName = FTSUtils::getDocsTableName(bindData.tableID, bindData.indexName);
-    auto docTableEntry = context.clientContext->getCatalog()->getTableCatalogEntry(
-        context.clientContext->getTransaction(), docTableName);
+    auto docTableEntry = catalog->getTableCatalogEntry(transaction, docTableName);
     graph::NativeGraphEntry entry{{docTableEntry}, {} /* relTableEntries */};
     graph::OnDiskGraph graph(context.clientContext, std::move(entry));
     auto sharedState = LenComputeSharedState{};
@@ -283,14 +284,11 @@ static offset_t tableFunc(const TableFuncInput& input, TableFuncOutput&) {
     auto indexEntry = std::make_unique<catalog::IndexCatalogEntry>(FTSIndexCatalogEntry::TYPE_NAME,
         bindData.tableID, bindData.indexName, bindData.propertyIDs,
         std::make_unique<FTSIndexAuxInfo>(ftsConfig));
-    auto catalog = context.clientContext->getCatalog();
-    catalog->createIndex(context.clientContext->getTransaction(), std::move(indexEntry));
+    catalog->createIndex(transaction, std::move(indexEntry));
 
-    auto nodeTable = context.clientContext->getStorageManager()
-                         ->getTable(bindData.tableID)
-                         ->ptrCast<storage::NodeTable>();
-    auto tableEntry =
-        catalog->getTableCatalogEntry(context.clientContext->getTransaction(), bindData.tableID);
+    auto storageManager = context.clientContext->getStorageManager();
+    auto nodeTable = storageManager->getTable(bindData.tableID)->ptrCast<storage::NodeTable>();
+    auto tableEntry = catalog->getTableCatalogEntry(transaction, bindData.tableID);
     std::vector<column_id_t> columnIDs;
     std::vector<PhysicalTypeID> columnTypes;
     for (auto& propertyID : bindData.propertyIDs) {
@@ -302,11 +300,28 @@ static offset_t tableFunc(const TableFuncInput& input, TableFuncOutput&) {
         std::move(columnIDs), std::move(columnTypes),
         ftsIndexType.constraintType == storage::IndexConstraintType::PRIMARY,
         ftsIndexType.definitionType == storage::IndexDefinitionType::BUILTIN};
-    auto storageInfo = std::make_unique<FTSStorageInfo>(numDocs, avgDocLen);
+    auto storageInfo = std::make_unique<FTSStorageInfo>(numDocs, avgDocLen,
+        nodeTable->getNumTotalRows(context.clientContext->getTransaction()));
     auto onDiskIndex = std::make_unique<FTSIndex>(std::move(indexInfo), std::move(storageInfo),
         std::move(ftsConfig), context.clientContext);
+    if (!context.clientContext->isInMemory()) {
+        // Checkpoint internal tables.
+        onDiskIndex->getInternalTableInfo().docTable->checkpoint(context.clientContext,
+            docTableEntry);
+        auto termsTableName = FTSUtils::getTermsTableName(bindData.tableID, bindData.indexName);
+        auto termsTableEntry =
+            context.clientContext->getCatalog()->getTableCatalogEntry(transaction, termsTableName);
+        onDiskIndex->getInternalTableInfo().termsTable->checkpoint(context.clientContext,
+            termsTableEntry);
+        auto appearsInTableName =
+            FTSUtils::getAppearsInTableName(bindData.tableID, bindData.indexName);
+        auto appearsInTableEntry = context.clientContext->getCatalog()->getTableCatalogEntry(
+            transaction, appearsInTableName);
+        onDiskIndex->getInternalTableInfo().appearsInfoTable->checkpoint(context.clientContext,
+            appearsInTableEntry);
+    }
     nodeTable->addIndex(std::move(onDiskIndex));
-    context.clientContext->getTransaction()->setForceCheckpoint();
+    transaction->setForceCheckpoint();
     return 0;
 }
 

Sources/cxx-kuzu/kuzu/extension/fts/src/function/fts_config.cpp

@@ -185,6 +185,9 @@ QueryFTSConfig::QueryFTSConfig(const function::optional_params_t& optionalParams
         } else if (Conjunctive::NAME == lowerCaseName) {
             value.validateType(Conjunctive::TYPE);
             isConjunctive = value.getValue<bool>();
+        } else if (TopK::NAME == lowerCaseName) {
+            value.validateType(TopK::TYPE);
+            topK = value.getValue<uint64_t>();
         } else {
             throw common::BinderException{"Unrecognized optional parameter: " + name};
         }

Sources/cxx-kuzu/kuzu/extension/fts/src/function/query_fts_bind_data.cpp

@@ -27,6 +27,8 @@ QueryFTSOptionalParams::QueryFTSOptionalParams(const binder::expression_vector&
             b = optionalParam;
         } else if (paramName == Conjunctive::NAME) {
             conjunctive = optionalParam;
+        } else if (paramName == TopK::NAME) {
+            topK = optionalParam;
         } else {
             throw common::BinderException{"Unknown optional parameter: " + paramName};
         }
@@ -45,6 +47,9 @@ QueryFTSConfig QueryFTSOptionalParams::getConfig() const {
         config.isConjunctive =
             ExpressionUtil::evaluateLiteral<bool>(*conjunctive, LogicalType::BOOL());
     }
+    if (topK != nullptr) {
+        config.topK = ExpressionUtil::evaluateLiteral<uint64_t>(*topK, LogicalType::UINT64());
+    }
     return config;
 }
 

Sources/cxx-kuzu/kuzu/extension/fts/src/function/query_fts_index.cpp

@@ -1,15 +1,22 @@
 #include "function/query_fts_index.h"
 
+#include <queue>
+
 #include "binder/binder.h"
 #include "binder/expression/expression_util.h"
 #include "binder/expression/literal_expression.h"
+#include "binder/query/reading_clause/bound_table_function_call.h"
 #include "catalog/fts_index_catalog_entry.h"
 #include "common/exception/binder.h"
 #include "common/types/internal_id_util.h"
 #include "function/fts_index_utils.h"
 #include "function/gds/gds_utils.h"
 #include "function/query_fts_bind_data.h"
+#include "graph/on_disk_graph.h"
 #include "index/fts_index.h"
+#include "planner/operator/logical_hash_join.h"
+#include "planner/operator/logical_table_function_call.h"
+#include "planner/planner.h"
 #include "processor/execution_context.h"
 #include "storage/storage_manager.h"
 #include "storage/table/node_table.h"
@@ -22,6 +29,87 @@ using namespace storage;
 using namespace function;
 using namespace binder;
 using namespace common;
+using namespace planner;
+
+struct DocScore {
+    offset_t offset;
+    double score;
+};
+
+struct QFTSSharedState : public GDSFuncSharedState {
+    common::table_id_t outputTableID;
+
+    QFTSSharedState(std::shared_ptr<processor::FactorizedTable> fTable,
+        std::unique_ptr<graph::Graph> graph, common::table_id_t outputTableID)
+        : GDSFuncSharedState{std::move(fTable), std::move(graph)}, outputTableID{outputTableID} {}
+
+    virtual void addDocScore(std::vector<ValueVector*> vectors,
+        processor::FactorizedTable& localTable, DocScore docScore) {
+        KU_ASSERT(vectors[0]->dataType.getLogicalTypeID() == LogicalTypeID::INTERNAL_ID);
+        KU_ASSERT(vectors[1]->dataType.getLogicalTypeID() == LogicalTypeID::DOUBLE);
+        vectors[0]->setValue(0, internalID_t{(common::offset_t)docScore.offset, outputTableID});
+        vectors[1]->setValue(0, docScore.score);
+        localTable.append(vectors);
+    }
+
+    virtual void finalizeResult() { factorizedTablePool.mergeLocalTables(); }
+};
+
+struct ScoreFTInsertState {
+    ValueVector docsVector;
+    ValueVector scoreVector;
+    std::vector<ValueVector*> vectors;
+
+    ScoreFTInsertState();
+};
+
+ScoreFTInsertState::ScoreFTInsertState()
+    : docsVector{LogicalType::INTERNAL_ID()}, scoreVector{LogicalType::DOUBLE()} {
+    auto state = DataChunkState::getSingleValueDataChunkState();
+    docsVector.state = state;
+    scoreVector.state = state;
+    vectors.push_back(&docsVector);
+    vectors.push_back(&scoreVector);
+}
+
+// This sharedState is only used when the user sets the top parameter.
+struct QFTSTopKSharedState : public QFTSSharedState {
+    struct PriorityQueueComp {
+        bool operator()(const DocScore& left, const DocScore& right) {
+            return left.score > right.score;
+        }
+    };
+    std::priority_queue<DocScore, std::vector<DocScore>, PriorityQueueComp> minHeap;
+    std::mutex mtx;
+    uint64_t topK;
+
+    QFTSTopKSharedState(std::shared_ptr<processor::FactorizedTable> fTable,
+        std::unique_ptr<graph::Graph> graph, uint64_t topK, common::table_id_t outputTableID)
+        : QFTSSharedState{std::move(fTable), std::move(graph), outputTableID}, topK{topK} {}
+
+    void addDocScore(std::vector<ValueVector*> /*vectors*/,
+        processor::FactorizedTable& /*localTable*/, DocScore docScore) override {
+        std::lock_guard guard{mtx};
+        if (minHeap.size() < topK) {
+            minHeap.push(docScore);
+        } else if (docScore.score > minHeap.top().score) {
+            minHeap.pop();
+            minHeap.push(docScore);
+        }
+    }
+
+    void finalizeResult() override {
+        ScoreFTInsertState insertState;
+        auto globalTable = factorizedTablePool.getGlobalTable();
+        while (!minHeap.empty()) {
+            auto& docScore = minHeap.top();
+            insertState.docsVector.setValue(0, nodeID_t{(offset_t)docScore.offset, outputTableID});
+            insertState.scoreVector.setValue(0, docScore.score);
+            globalTable->append(insertState.vectors);
+            minHeap.pop();
+        }
+    }
+};
 
 struct ScoreData {
     uint64_t df;
@@ -69,43 +157,32 @@ struct QFTSEdgeCompute final : EdgeCompute {
 class QFTSOutputWriter {
 public:
     QFTSOutputWriter(const node_id_map_t<ScoreInfo>& scores, MemoryManager* mm,
-        QueryFTSConfig config, const QueryFTSBindData& bindData, uint64_t numUniqueTerms);
+        QueryFTSConfig config, const QueryFTSBindData& bindData, uint64_t numUniqueTerms,
+        QFTSSharedState& sharedState);
 
     void write(processor::FactorizedTable& scoreFT, nodeID_t docNodeID, uint64_t len,
         int64_t docsID);
 
     std::unique_ptr<QFTSOutputWriter> copy() {
-        return std::make_unique<QFTSOutputWriter>(scores, mm, config, bindData, numUniqueTerms);
+        return std::make_unique<QFTSOutputWriter>(scores, mm, config, bindData, numUniqueTerms,
+            sharedState);
     }
 
-private:
-    std::unique_ptr<ValueVector> createVector(const LogicalType& type);
-
 private:
     const node_id_map_t<ScoreInfo>& scores;
     MemoryManager* mm;
     QueryFTSConfig config;
     const QueryFTSBindData& bindData;
-
-    std::unique_ptr<ValueVector> docsVector;
-    std::unique_ptr<ValueVector> scoreVector;
-    std::vector<ValueVector*> vectors;
     uint64_t numUniqueTerms;
+    QFTSSharedState& sharedState;
+    ScoreFTInsertState scoreFtInsertState;
 };
 
 QFTSOutputWriter::QFTSOutputWriter(const node_id_map_t<ScoreInfo>& scores, MemoryManager* mm,
-    QueryFTSConfig config, const QueryFTSBindData& bindData, uint64_t numUniqueTerms)
-    : scores{scores}, mm{mm}, config{config}, bindData{bindData}, numUniqueTerms{numUniqueTerms} {
-    docsVector = createVector(LogicalType::INTERNAL_ID());
-    scoreVector = createVector(LogicalType::UINT64());
-}
-
-std::unique_ptr<ValueVector> QFTSOutputWriter::createVector(const LogicalType& type) {
-    auto vector = std::make_unique<ValueVector>(type.copy(), mm);
-    vector->state = DataChunkState::getSingleValueDataChunkState();
-    vectors.push_back(vector.get());
-    return vector;
-}
+    QueryFTSConfig config, const QueryFTSBindData& bindData, uint64_t numUniqueTerms,
+    QFTSSharedState& sharedState)
+    : scores{scores}, mm{mm}, config{config}, bindData{bindData}, numUniqueTerms{numUniqueTerms},
+      sharedState{sharedState}, scoreFtInsertState{} {}
 
 void QFTSOutputWriter::write(processor::FactorizedTable& scoreFT, nodeID_t docNodeID, uint64_t len,
     int64_t docsID) {
@@ -131,9 +208,7 @@ void QFTSOutputWriter::write(processor::FactorizedTable& scoreFT, nodeID_t docNo
         score += log10((numDocs - df + 0.5) / (df + 0.5) + 1) *
                  ((tf * (k + 1) / (tf + k * (1 - b + b * (len / avgDocLen)))));
     }
-    docsVector->setValue(0, nodeID_t{static_cast<offset_t>(docsID), bindData.outputTableID});
-    scoreVector->setValue(0, score);
-    scoreFT.append(vectors);
+    sharedState.addDocScore(scoreFtInsertState.vectors, scoreFT, {(uint64_t)docsID, score});
 }
 
 class QFTSVertexCompute final : public VertexCompute {
@@ -233,7 +308,7 @@ static void initFrontier(FrontierPair& frontierPair, table_id_t termsTableID,
 static offset_t tableFunc(const TableFuncInput& input, TableFuncOutput&) {
     auto clientContext = input.context->clientContext;
     auto transaction = clientContext->getTransaction();
-    auto sharedState = input.sharedState->ptrCast<GDSFuncSharedState>();
+    auto sharedState = input.sharedState->ptrCast<QFTSSharedState>();
     auto graph = sharedState->graph.get();
     auto graphEntry = graph->getGraphEntry();
     auto qFTSBindData = input.bindData->constPtrCast<QueryFTSBindData>();
@@ -265,7 +340,7 @@ static offset_t tableFunc(const TableFuncInput& input, TableFuncOutput&) {
     auto mm = clientContext->getMemoryManager();
     auto numUniqueTerms = getNumUniqueTerms(terms);
     auto writer = std::make_unique<QFTSOutputWriter>(scores, mm, qFTSBindData->getConfig(),
-        *qFTSBindData, numUniqueTerms);
+        *qFTSBindData, numUniqueTerms, *sharedState);
     auto vc = std::make_unique<QFTSVertexCompute>(mm, sharedState, std::move(writer));
     auto vertexPropertiesToScan = std::vector<std::string>{DOC_LEN_PROP_NAME, DOC_ID_PROP_NAME};
     auto docsEntry = graphEntry->nodeInfos[1].entry;
@@ -282,7 +357,7 @@ static offset_t tableFunc(const TableFuncInput& input, TableFuncOutput&) {
         GDSUtils::runVertexCompute(input.context, GDSDensityState::DENSE, graph, *vc, docsEntry,
             vertexPropertiesToScan);
     }
-    sharedState->factorizedTablePool.mergeLocalTables();
+    sharedState->finalizeResult();
     return 0;
 }
 
@@ -342,6 +417,42 @@ static std::unique_ptr<TableFuncBindData> bindFunc(main::ClientContext* context,
     return bindData;
 }
 
+static void getLogicalPlan(Planner* planner, const BoundReadingClause& readingClause,
+    const expression_vector& predicates, LogicalPlan& plan) {
+    auto& call = readingClause.constCast<BoundTableFunctionCall>();
+    auto bindData = call.getBindData()->constPtrCast<GDSBindData>();
+    auto op = std::make_shared<LogicalTableFunctionCall>(call.getTableFunc(), bindData->copy());
+    op->computeFactorizedSchema();
+    planner->planReadOp(std::move(op), predicates, plan);
+
+    auto nodeOutput = bindData->nodeOutput->ptrCast<NodeExpression>();
+    KU_ASSERT(nodeOutput != nullptr);
+    planner->getCardinliatyEstimatorUnsafe().init(*nodeOutput);
+    auto scanPlan = planner->getNodePropertyScanPlan(*nodeOutput);
+    if (scanPlan.isEmpty()) {
+        return;
+    }
+    expression_vector joinConditions;
+    joinConditions.push_back(nodeOutput->getInternalID());
+    planner->appendHashJoin(joinConditions, JoinType::INNER, scanPlan, plan, plan);
+    plan.getLastOperator()->cast<LogicalHashJoin>().getSIPInfoUnsafe().direction =
+        SIPDirection::FORCE_BUILD_TO_PROBE;
+}
+
+std::shared_ptr<TableFuncSharedState> initSharedState(const TableFuncInitSharedStateInput& input) {
+    auto bindData = input.bindData->constPtrCast<QueryFTSBindData>();
+    auto graph = std::make_unique<graph::OnDiskGraph>(input.context->clientContext,
+        bindData->graphEntry.copy());
+    auto topK = bindData->getConfig().topK;
+    if (bindData->getConfig().topK == INVALID_TOP_K) {
+        return std::make_shared<QFTSSharedState>(bindData->getResultTable(), std::move(graph),
+            bindData->outputTableID);
+    } else {
+        return std::make_shared<QFTSTopKSharedState>(bindData->getResultTable(), std::move(graph),
+            topK, bindData->outputTableID);
+    }
+}
+
 function_set QueryFTSFunction::getFunctionSet() {
     function_set result;
     // inputs are tableName, indexName, query
@@ -350,10 +461,10 @@ function_set QueryFTSFunction::getFunctionSet() {
             LogicalTypeID::STRING});
     func->bindFunc = bindFunc;
     func->tableFunc = tableFunc;
-    func->initSharedStateFunc = GDSFunction::initSharedState;
+    func->initSharedStateFunc = initSharedState;
     func->initLocalStateFunc = TableFunction::initEmptyLocalState;
     func->canParallelFunc = [] { return false; };
-    func->getLogicalPlanFunc = GDSFunction::getLogicalPlan;
+    func->getLogicalPlanFunc = getLogicalPlan;
     func->getPhysicalPlanFunc = GDSFunction::getPhysicalPlan;
     result.push_back(std::move(func));
     return result;

Sources/cxx-kuzu/kuzu/extension/fts/src/include/function/fts_config.h

@@ -91,6 +91,14 @@ struct Conjunctive {
     static constexpr bool DEFAULT_VALUE = false;
 };
 
+struct TopK {
+    static constexpr const char* NAME = "top";
+    static constexpr common::LogicalTypeID TYPE = common::LogicalTypeID::UINT64;
+    static constexpr uint64_t DEFAULT_VALUE = UINT64_MAX;
+};
+
+constexpr uint64_t INVALID_TOP_K = UINT64_MAX;
+
 struct QueryFTSConfig {
     // k: parameter controls the influence of term frequency saturation. It limits the effect of
     // additional occurrences of a term within a document.
@@ -99,6 +107,7 @@ struct QueryFTSConfig {
     // document length.
     double b = B::DEFAULT_VALUE;
     bool isConjunctive = Conjunctive::DEFAULT_VALUE;
+    uint64_t topK = TopK::DEFAULT_VALUE;
 
     QueryFTSConfig() = default;
     explicit QueryFTSConfig(const function::optional_params_t& optionalParams);

Sources/cxx-kuzu/kuzu/extension/fts/src/include/function/query_fts_bind_data.h

@@ -12,6 +12,7 @@ struct QueryFTSOptionalParams {
     std::shared_ptr<binder::Expression> k;
     std::shared_ptr<binder::Expression> b;
     std::shared_ptr<binder::Expression> conjunctive;
+    std::shared_ptr<binder::Expression> topK;
 
     explicit QueryFTSOptionalParams(const binder::expression_vector& optionalParams);