Parallel Hash Index

CMakeLists.txt

@@ -1,6 +1,6 @@
 cmake_minimum_required(VERSION 3.15)
 
-project(Kuzu VERSION 0.1.0.3 LANGUAGES CXX C)
+project(Kuzu VERSION 0.1.0.4 LANGUAGES CXX C)
 
 find_package(Threads REQUIRED)
 

src/common/file_system/local_file_system.cpp

@@ -1,5 +1,7 @@
 #include "common/file_system/local_file_system.h"
 
+#include <cstring>
+
 #if defined(_WIN32)
 #include <fileapi.h>
 #include <io.h>
@@ -57,7 +59,7 @@
 #else
     int fd = open(path.c_str(), flags, 0644);
     if (fd == -1) {
-        throw Exception("Cannot open file: " + path);
+        throw Exception(stringFormat("Cannot open file {}: {}", path, posixErrMessage()));
     }
     if (lock_type != FileLockType::NO_LOCK) {
         struct flock fl;

src/include/common/copy_constructors.h

@@ -2,6 +2,8 @@
 
 // This file defines many macros for controlling copy constructors and move constructors on classes.
 
+// NOLINTBEGIN(bugprone-macro-parentheses): Although this is a good check in general, here, we
+// cannot add parantheses around the arguments, for it would be invalid syntax.
 #define DELETE_COPY_CONSTRUCT(Object) Object(const Object& other) = delete
 #define DELETE_COPY_ASSN(Object) Object& operator=(const Object& other) = delete
 // NOLINTBEGIN
@@ -60,3 +62,5 @@
 #define DELETE_COPY_AND_MOVE(Object)                                                               \
     DELETE_BOTH_COPY(Object);                                                                      \
     DELETE_BOTH_MOVE(Object)
+
+// NOLINTEND(bugprone-macro-parentheses)

src/include/function/hash/hash_functions.h

@@ -16,7 +16,13 @@ namespace function {
 constexpr const uint64_t NULL_HASH = UINT64_MAX;
 
 inline common::hash_t murmurhash64(uint64_t x) {
-    return x * UINT64_C(0xbf58476d1ce4e5b9);
+    // taken from https://nullprogram.com/blog/2018/07/31.
+    x ^= x >> 32;
+    x *= 0xd6e8feb86659fd93U;
+    x ^= x >> 32;
+    x *= 0xd6e8feb86659fd93U;
+    x ^= x >> 32;
+    return x;
 }
 
 inline common::hash_t combineHashScalar(common::hash_t a, common::hash_t b) {

src/include/main/database.h

@@ -44,6 +44,7 @@ struct KUZU_API SystemConfig {
  */
 class Database {
     friend class EmbeddedShell;
+    friend class ClientContext;
     friend class Connection;
     friend class StorageDriver;
     friend class kuzu::testing::BaseGraphTest;

src/include/processor/operator/persistent/copy_node.h

@@ -4,6 +4,7 @@
 
 #include "processor/operator/aggregate/hash_aggregate.h"
 #include "processor/operator/call/in_query_call.h"
+#include "processor/operator/persistent/index_builder.h"
 #include "processor/operator/sink.h"
 #include "storage/store/node_group.h"
 #include "storage/store/node_table.h"
@@ -18,19 +19,22 @@ class CopyNodeSharedState {
 
 public:
     CopyNodeSharedState()
-        : indexBuilder{nullptr}, readerSharedState{nullptr}, distinctSharedState{nullptr},
-          currentNodeGroupIdx{0}, sharedNodeGroup{nullptr} {};
+        : readerSharedState{nullptr}, distinctSharedState{nullptr}, currentNodeGroupIdx{0},
+          sharedNodeGroup{nullptr} {};
 
-    void init(common::VirtualFileSystem* vfs);
+    void init();
 
     inline common::offset_t getNextNodeGroupIdx() {
-        std::unique_lock<std::mutex> lck{mtx};
+        std::unique_lock lck{mtx};
         return getNextNodeGroupIdxWithoutLock();
     }
 
     inline uint64_t getCurNodeGroupIdx() const { return currentNodeGroupIdx; }
 
-    void appendLocalNodeGroup(std::unique_ptr<storage::NodeGroup> localNodeGroup);
+    void appendIncompleteNodeGroup(std::unique_ptr<storage::NodeGroup> localNodeGroup,
+        std::optional<IndexBuilder>& indexBuilder);
+
+    void addLastNodeGroup(std::optional<IndexBuilder>& indexBuilder);
 
 private:
     inline common::offset_t getNextNodeGroupIdxWithoutLock() { return currentNodeGroupIdx++; }
@@ -89,10 +93,12 @@ class CopyNode : public Sink {
 public:
     CopyNode(std::shared_ptr<CopyNodeSharedState> sharedState, std::unique_ptr<CopyNodeInfo> info,
         std::unique_ptr<ResultSetDescriptor> resultSetDescriptor,
-        std::unique_ptr<PhysicalOperator> child, uint32_t id, const std::string& paramsString)
+        std::unique_ptr<PhysicalOperator> child, uint32_t id, const std::string& paramsString,
+        std::optional<IndexBuilder> indexBuilder = std::nullopt)
         : Sink{std::move(resultSetDescriptor), PhysicalOperatorType::COPY_NODE, std::move(child),
               id, paramsString},
-          sharedState{std::move(sharedState)}, info{std::move(info)} {}
+          sharedState{std::move(sharedState)}, info{std::move(info)},
+          indexBuilder(std::move(indexBuilder)) {}
 
     inline std::shared_ptr<CopyNodeSharedState> getSharedState() const { return sharedState; }
 
@@ -108,41 +114,30 @@ class CopyNode : public Sink {
 
     inline std::unique_ptr<PhysicalOperator> clone() final {
         return std::make_unique<CopyNode>(sharedState, info->copy(), resultSetDescriptor->copy(),
-            children[0]->clone(), id, paramsString);
+            children[0]->clone(), id, paramsString,
+            indexBuilder ? std::make_optional<IndexBuilder>(indexBuilder->clone()) : std::nullopt);
     }
 
     static void writeAndResetNodeGroup(common::node_group_idx_t nodeGroupIdx,
-        storage::PrimaryKeyIndexBuilder* pkIndex, common::column_id_t pkColumnID,
+        std::optional<IndexBuilder>& indexBuilder, common::column_id_t pkColumnID,
         storage::NodeTable* table, storage::NodeGroup* nodeGroup);
 
 private:
-    static void populatePKIndex(storage::PrimaryKeyIndexBuilder* pkIndex,
-        storage::ColumnChunk* chunk, common::offset_t startNodeOffset, common::offset_t numNodes);
-    static void checkNonNullConstraint(
-        storage::NullColumnChunk* nullChunk, common::offset_t numNodes);
-
-    template<typename T>
-    static uint64_t appendToPKIndex(storage::PrimaryKeyIndexBuilder* pkIndex,
-        storage::ColumnChunk* chunk, common::offset_t startOffset, common::offset_t numNodes);
-
     void copyToNodeGroup();
+    void initGlobalIndexBuilder(ExecutionContext* context);
+    void initLocalIndexBuilder(ExecutionContext* context);
 
 protected:
     std::shared_ptr<CopyNodeSharedState> sharedState;
     std::unique_ptr<CopyNodeInfo> info;
 
+    std::optional<IndexBuilder> indexBuilder;
+
     common::DataChunkState* columnState;
     std::vector<std::shared_ptr<common::ValueVector>> nullColumnVectors;
     std::vector<common::ValueVector*> columnVectors;
     std::unique_ptr<storage::NodeGroup> localNodeGroup;
 };
 
-template<>
-uint64_t CopyNode::appendToPKIndex<int64_t>(storage::PrimaryKeyIndexBuilder* pkIndex,
-    storage::ColumnChunk* chunk, common::offset_t startOffset, common::offset_t numNodes);
-template<>
-uint64_t CopyNode::appendToPKIndex<std::string>(storage::PrimaryKeyIndexBuilder* pkIndex,
-    storage::ColumnChunk* chunk, common::offset_t startOffset, common::offset_t numNodes);
-
 } // namespace processor
 } // namespace kuzu

src/include/processor/operator/persistent/index_builder.h

@@ -0,0 +1,135 @@
+#pragma once
+
+#include <variant>
+
+#include "common/copy_constructors.h"
+#include "common/mpsc_queue.h"
+#include "common/static_vector.h"
+#include "common/types/internal_id_t.h"
+#include "common/types/types.h"
+#include "processor/execution_context.h"
+#include "storage/index/hash_index_builder.h"
+#include "storage/store/column_chunk.h"
+
+namespace kuzu {
+namespace processor {
+
+constexpr size_t BUFFER_SIZE = 1024;
+using IntBuffer = common::StaticVector<std::pair<int64_t, common::offset_t>, BUFFER_SIZE>;
+using StringBuffer = common::StaticVector<std::pair<std::string, common::offset_t>, BUFFER_SIZE>;
+
+class IndexBuilderGlobalQueues {
+public:
+    explicit IndexBuilderGlobalQueues(std::unique_ptr<storage::PrimaryKeyIndexBuilder> pkIndex);
+
+    void flushToDisk() const;
+
+    void insert(size_t index, StringBuffer elem);
+    void insert(size_t index, IntBuffer elem);
+
+    void consume();
+
+    common::LogicalTypeID pkTypeID() const { return pkIndex->keyTypeID(); }
+
+private:
+    void maybeConsumeIndex(size_t index);
+
+    std::array<std::mutex, storage::NUM_HASH_INDEXES> mutexes;
+    std::unique_ptr<storage::PrimaryKeyIndexBuilder> pkIndex;
+
+    using StringQueues = std::array<common::MPSCQueue<StringBuffer>, storage::NUM_HASH_INDEXES>;
+    using IntQueues = std::array<common::MPSCQueue<IntBuffer>, storage::NUM_HASH_INDEXES>;
+
+    // Queues for distributing primary keys.
+    std::variant<StringQueues, IntQueues> queues;
+};
+
+class IndexBuilderLocalBuffers {
+public:
+    explicit IndexBuilderLocalBuffers(IndexBuilderGlobalQueues& globalQueues);
+
+    void insert(std::string key, common::offset_t value);
+    void insert(int64_t key, common::offset_t value);
+
+    void flush();
+
+private:
+    IndexBuilderGlobalQueues* globalQueues;
+
+    // These arrays are much too large to be inline.
+    using StringBuffers = std::array<StringBuffer, storage::NUM_HASH_INDEXES>;
+    using IntBuffers = std::array<IntBuffer, storage::NUM_HASH_INDEXES>;
+    std::unique_ptr<StringBuffers> stringBuffers;
+    std::unique_ptr<IntBuffers> intBuffers;
+};
+
+class IndexBuilderSharedState {
+    friend class IndexBuilder;
+
+public:
+    explicit IndexBuilderSharedState(std::unique_ptr<storage::PrimaryKeyIndexBuilder> pkIndex);
+    void consume() { globalQueues.consume(); }
+    void flush() { globalQueues.flushToDisk(); }
+
+    void addProducer() { producers.fetch_add(1, std::memory_order_relaxed); }
+    void quitProducer();
+    bool isDone() { return done.load(std::memory_order_relaxed); }
+
+private:
+    IndexBuilderGlobalQueues globalQueues;
+
+    std::atomic<size_t> producers;
+    std::atomic<bool> done;
+};
+
+// RAII for producer counting.
+class ProducerToken {
+public:
+    explicit ProducerToken(std::shared_ptr<IndexBuilderSharedState> sharedState)
+        : sharedState(std::move(sharedState)) {
+        this->sharedState->addProducer();
+    }
+    DELETE_COPY_DEFAULT_MOVE(ProducerToken);
+
+    void quit() {
+        sharedState->quitProducer();
+        sharedState.reset();
+    }
+    ~ProducerToken() {
+        if (sharedState) {
+            quit();
+        }
+    }
+
+private:
+    std::shared_ptr<IndexBuilderSharedState> sharedState;
+};
+
+class IndexBuilder {
+    explicit IndexBuilder(std::shared_ptr<IndexBuilderSharedState> sharedState);
+
+public:
+    DELETE_COPY_DEFAULT_MOVE(IndexBuilder);
+    explicit IndexBuilder(std::unique_ptr<storage::PrimaryKeyIndexBuilder> pkIndex);
+
+    IndexBuilder clone() { return IndexBuilder(sharedState); }
+
+    void initGlobalStateInternal(ExecutionContext* /*context*/) {}
+    void initLocalStateInternal(ExecutionContext* /*context*/) {}
+    void insert(
+        storage::ColumnChunk* chunk, common::offset_t nodeOffset, common::offset_t numNodes);
+
+    ProducerToken getProducerToken() const { return ProducerToken(sharedState); }
+
+    void finishedProducing();
+    void finalize(ExecutionContext* context);
+
+private:
+    void checkNonNullConstraint(storage::NullColumnChunk* nullChunk, common::offset_t numNodes);
+    std::shared_ptr<IndexBuilderSharedState> sharedState;
+
+    IndexBuilderLocalBuffers localBuffers;
+};
+
+} // namespace processor
+} // namespace kuzu