Merge pull request #1839 from kuzudb/ftable-scan-refactor

ResultCollector, FTableScan, UnionAllScan refactor
kuzudb · Jul 20, 2023 · 60b21ca · 60b21ca
2 parents 0262e82 + a7ef541
commit 60b21ca
Show file tree

Hide file tree

Showing 26 changed files with 328 additions and 395 deletions.
diff --git a/src/include/planner/logical_plan/logical_operator/base_logical_operator.h b/src/include/planner/logical_plan/logical_operator/base_logical_operator.h
@@ -26,7 +26,6 @@ enum class LogicalOperatorType : uint8_t {
     EXTEND,
     FILTER,
     FLATTEN,
-    FTABLE_SCAN,
     HASH_JOIN,
     IN_QUERY_CALL,
     INDEX_SCAN_NODE,

diff --git a/src/include/planner/logical_plan/logical_operator/logical_ftable_scan.h b/src/include/planner/logical_plan/logical_operator/logical_ftable_scan.h
diff --git a/src/include/processor/mapper/plan_mapper.h b/src/include/processor/mapper/plan_mapper.h
@@ -80,8 +80,6 @@ class PlanMapper {
         planner::LogicalOperator* logicalOperator);
     std::unique_ptr<PhysicalOperator> mapLogicalExpressionsScanToPhysical(
         planner::LogicalOperator* logicalOperator);
-    std::unique_ptr<PhysicalOperator> mapLogicalFTableScanToPhysical(
-        planner::LogicalOperator* logicalOperator);
     std::unique_ptr<PhysicalOperator> mapLogicalCreateNodeToPhysical(
         planner::LogicalOperator* logicalOperator);
     std::unique_ptr<PhysicalOperator> mapLogicalCreateRelToPhysical(
@@ -120,9 +118,13 @@ class PlanMapper {
         planner::LogicalOperator* logicalOperator);
     std::unique_ptr<PhysicalOperator> mapLogicalCreateMacroToPhysical(
         planner::LogicalOperator* logicalOperator);
-    std::unique_ptr<ResultCollector> appendResultCollector(
-        const binder::expression_vector& expressionsToCollect, planner::Schema* schema,
+
+    std::unique_ptr<ResultCollector> createResultCollector(
+        const binder::expression_vector& expressions, planner::Schema* schema,
         std::unique_ptr<PhysicalOperator> prevOperator);
+    std::unique_ptr<PhysicalOperator> createFactorizedTableScan(
+        const binder::expression_vector& expressions, planner::Schema* schema,
+        std::shared_ptr<FactorizedTable> table, std::unique_ptr<PhysicalOperator> prevOperator);
 
     inline uint32_t getOperatorID() { return physicalOperatorID++; }
 

diff --git a/src/include/processor/operator/result_collector.h b/src/include/processor/operator/result_collector.h
@@ -6,77 +6,65 @@
 namespace kuzu {
 namespace processor {
 
-struct FTableScanMorsel {
-    FTableScanMorsel(FactorizedTable* table, uint64_t startTupleIdx, uint64_t numTuples)
-        : table{table}, startTupleIdx{startTupleIdx}, numTuples{numTuples} {}
-
-    FactorizedTable* table;
-    uint64_t startTupleIdx;
-    uint64_t numTuples;
-};
-
-class FTableSharedState {
+class ResultCollectorSharedState {
 public:
-    FTableSharedState(std::shared_ptr<FactorizedTable> table, uint64_t maxMorselSize)
-        : table{std::move(table)}, maxMorselSize{maxMorselSize} {}
-    FTableSharedState(storage::MemoryManager* memoryManager,
-        std::unique_ptr<FactorizedTableSchema> tableSchema, uint64_t maxMorselSize)
-        : maxMorselSize{maxMorselSize}, nextTupleIdxToScan{0} {
-        table = std::make_shared<FactorizedTable>(memoryManager, std::move(tableSchema));
-    }
-
-    // We want to control the granularity of morsel, e.g. in recursive join pipeline, we always want
-    // to scan 1 src at a time.
-    inline void setMaxMorselSize(uint64_t size) { maxMorselSize = size; }
-    inline uint64_t getMaxMorselSize() const { return maxMorselSize; }
+    explicit ResultCollectorSharedState(std::shared_ptr<FactorizedTable> table)
+        : table{std::move(table)} {}
 
     inline void mergeLocalTable(FactorizedTable& localTable) {
-        std::lock_guard<std::mutex> lck{mtx};
+        std::unique_lock lck{mtx};
         table->merge(localTable);
     }
 
     inline std::shared_ptr<FactorizedTable> getTable() { return table; }
 
-    std::unique_ptr<FTableScanMorsel> getMorsel();
-
 private:
     std::mutex mtx;
     std::shared_ptr<FactorizedTable> table;
-    uint64_t maxMorselSize;
+};
+
+struct ResultCollectorInfo {
+    std::unique_ptr<FactorizedTableSchema> tableSchema;
+    std::vector<DataPos> payloadPositions;
 
-    uint64_t nextTupleIdxToScan = 0u;
+    ResultCollectorInfo(
+        std::unique_ptr<FactorizedTableSchema> tableSchema, std::vector<DataPos> payloadPositions)
+        : tableSchema{std::move(tableSchema)}, payloadPositions{std::move(payloadPositions)} {}
+    ResultCollectorInfo(const ResultCollectorInfo& other)
+        : tableSchema{other.tableSchema->copy()}, payloadPositions{other.payloadPositions} {}
+
+    inline std::unique_ptr<ResultCollectorInfo> copy() const {
+        return std::make_unique<ResultCollectorInfo>(*this);
+    }
 };
 
 class ResultCollector : public Sink {
 public:
     ResultCollector(std::unique_ptr<ResultSetDescriptor> resultSetDescriptor,
-        std::unique_ptr<FactorizedTableSchema> tableSchema, std::vector<DataPos> payloadsPos,
-        std::shared_ptr<FTableSharedState> sharedState, std::unique_ptr<PhysicalOperator> child,
-        uint32_t id, const std::string& paramsString)
+        std::unique_ptr<ResultCollectorInfo> info,
+        std::shared_ptr<ResultCollectorSharedState> sharedState,
+        std::unique_ptr<PhysicalOperator> child, uint32_t id, const std::string& paramsString)
         : Sink{std::move(resultSetDescriptor), PhysicalOperatorType::RESULT_COLLECTOR,
               std::move(child), id, paramsString},
-          tableSchema{std::move(tableSchema)}, payloadsPos{std::move(payloadsPos)},
-          sharedState{std::move(sharedState)} {}
+          info{std::move(info)}, sharedState{std::move(sharedState)} {}
 
     void initLocalStateInternal(ResultSet* resultSet, ExecutionContext* context) override;
 
     void executeInternal(ExecutionContext* context) override;
 
-    inline std::shared_ptr<FTableSharedState> getSharedState() { return sharedState; }
     inline std::shared_ptr<FactorizedTable> getResultFactorizedTable() {
         return sharedState->getTable();
     }
 
     std::unique_ptr<PhysicalOperator> clone() override {
-        return make_unique<ResultCollector>(resultSetDescriptor->copy(), tableSchema->copy(),
-            payloadsPos, sharedState, children[0]->clone(), id, paramsString);
+        return make_unique<ResultCollector>(resultSetDescriptor->copy(), info->copy(), sharedState,
+            children[0]->clone(), id, paramsString);
     }
 
 private:
-    std::unique_ptr<FactorizedTableSchema> tableSchema;
-    std::vector<DataPos> payloadsPos;
+    std::unique_ptr<ResultCollectorInfo> info;
+    std::shared_ptr<ResultCollectorSharedState> sharedState;
     std::vector<common::ValueVector*> payloadVectors;
-    std::shared_ptr<FTableSharedState> sharedState;
     std::unique_ptr<FactorizedTable> localTable;
 };
 

diff --git a/src/include/processor/operator/table_scan/base_table_scan.h b/src/include/processor/operator/table_scan/base_table_scan.h
diff --git a/src/include/processor/operator/table_scan/factorized_table_scan.h b/src/include/processor/operator/table_scan/factorized_table_scan.h
@@ -1,45 +1,81 @@
 #pragma once
 
-#include "processor/operator/table_scan/base_table_scan.h"
+#include "common/join_type.h"
+#include "processor/operator/result_collector.h"
 
 namespace kuzu {
 namespace processor {
 
-class FactorizedTableScan : public BaseTableScan {
+struct FactorizedTableScanInfo {
+    std::vector<DataPos> outputPositions;
+    std::vector<ft_col_idx_t> columnIndices;
+
+    FactorizedTableScanInfo(
+        std::vector<DataPos> outputPositions, std::vector<ft_col_idx_t> columnIndices)
+        : outputPositions{std::move(outputPositions)}, columnIndices{std::move(columnIndices)} {}
+    FactorizedTableScanInfo(const FactorizedTableScanInfo& other)
+        : outputPositions{other.outputPositions}, columnIndices{other.columnIndices} {}
+
+    inline std::unique_ptr<FactorizedTableScanInfo> copy() const {
+        return std::make_unique<FactorizedTableScanInfo>(*this);
+    }
+};
+
+class FactorizedTableScan;
+
+struct FactorizedTableScanMorsel {
+    uint64_t startTupleIdx;
+    uint64_t numTuples;
+
+    FactorizedTableScanMorsel(uint64_t startTupleIdx, uint64_t numTuples)
+        : startTupleIdx{startTupleIdx}, numTuples{numTuples} {}
+};
+
+class FactorizedTableScanSharedState {
+    friend class FactorizedTableScan;
+
 public:
-    // Scan all columns.
-    FactorizedTableScan(std::vector<DataPos> outVecPositions,
-        std::vector<uint32_t> colIndicesToScan, std::shared_ptr<FTableSharedState> sharedState,
+    FactorizedTableScanSharedState(std::shared_ptr<FactorizedTable> table, uint64_t maxMorselSize)
+        : table{std::move(table)}, maxMorselSize{maxMorselSize}, nextTupleIdxToScan{0} {}
+
+    std::unique_ptr<FactorizedTableScanMorsel> getMorsel();
+
+private:
+    std::mutex mtx;
+    std::shared_ptr<FactorizedTable> table;
+    uint64_t maxMorselSize;
+    uint64_t nextTupleIdxToScan;
+};
+
+class FactorizedTableScan : public PhysicalOperator {
+public:
+    FactorizedTableScan(std::unique_ptr<FactorizedTableScanInfo> info,
+        std::shared_ptr<FactorizedTableScanSharedState> sharedState,
         std::unique_ptr<PhysicalOperator> child, uint32_t id, const std::string& paramsString)
-        : BaseTableScan{PhysicalOperatorType::FACTORIZED_TABLE_SCAN, std::move(outVecPositions),
-              std::move(colIndicesToScan), std::move(child), id, paramsString},
-          sharedState{std::move(sharedState)} {}
-
-    // Scan some columns.
-    FactorizedTableScan(std::vector<DataPos> outVecPositions,
-        std::vector<uint32_t> colIndicesToScan, uint32_t id, const std::string& paramsString)
-        : BaseTableScan{PhysicalOperatorType::FACTORIZED_TABLE_SCAN, std::move(outVecPositions),
-              std::move(colIndicesToScan), id, paramsString} {}
-
-    // For clone only.
-    FactorizedTableScan(std::vector<DataPos> outVecPositions,
-        std::vector<uint32_t> colIndicesToScan, std::shared_ptr<FTableSharedState> sharedState,
-        uint32_t id, const std::string& paramsString)
-        : BaseTableScan{PhysicalOperatorType::FACTORIZED_TABLE_SCAN, std::move(outVecPositions),
-              std::move(colIndicesToScan), id, paramsString},
-          sharedState{std::move(sharedState)} {}
-
-    inline std::unique_ptr<FTableScanMorsel> getMorsel() override {
-        return sharedState->getMorsel();
-    }
+        : PhysicalOperator{PhysicalOperatorType::FACTORIZED_TABLE_SCAN, std::move(child), id,
+              paramsString},
+          info{std::move(info)}, sharedState{std::move(sharedState)} {}
+
+    FactorizedTableScan(std::unique_ptr<FactorizedTableScanInfo> info,
+        std::shared_ptr<FactorizedTableScanSharedState> sharedState, uint32_t id,
+        const std::string& paramsString)
+        : PhysicalOperator{PhysicalOperatorType::FACTORIZED_TABLE_SCAN, id, paramsString},
+          info{std::move(info)}, sharedState{std::move(sharedState)} {}
+
+    bool isSource() const final { return true; }
+
+    void initLocalStateInternal(ResultSet* resultSet_, ExecutionContext* context) final;
+
+    bool getNextTuplesInternal(ExecutionContext* context) final;
 
-    inline std::unique_ptr<PhysicalOperator> clone() override {
-        return make_unique<FactorizedTableScan>(
-            outVecPositions, colIndicesToScan, sharedState, id, paramsString);
+    inline std::unique_ptr<PhysicalOperator> clone() final {
+        return make_unique<FactorizedTableScan>(info->copy(), sharedState, id, paramsString);
     }
 
 private:
-    std::shared_ptr<FTableSharedState> sharedState;
+    std::unique_ptr<FactorizedTableScanInfo> info;
+    std::shared_ptr<FactorizedTableScanSharedState> sharedState;
+    std::vector<common::ValueVector*> vectors;
 };
 
 } // namespace processor