lbug 0.15.4

An in-process property graph database management system built for query speed and scalability
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

#pragma once

#include <memory>
#include <string>
#include <vector>

#include "catalog/catalog_entry/node_table_catalog_entry.h"
#include "common/arrow/arrow.h"
#include "common/exception/runtime.h"
#include "function/table/table_function.h"
#include "storage/table/columnar_node_table_base.h"

namespace lbug {
namespace storage {

struct ArrowNodeTableScanState final : ColumnarNodeTableScanState {
    size_t currentBatchIdx = static_cast<size_t>(common::INVALID_NODE_GROUP_IDX);
    size_t currentMorselStartOffset = 0; // Start of current morsel within batch
    size_t currentMorselEndOffset = 0;   // End of current morsel within batch
    std::vector<int64_t> outputToArrowColumnIdx;
    bool initialized = false;
    bool scanCompleted = false;

    ArrowNodeTableScanState(MemoryManager& mm, common::ValueVector* nodeIDVector,
        std::vector<common::ValueVector*> outputVectors,
        std::shared_ptr<common::DataChunkState> outChunkState)
        : ColumnarNodeTableScanState{mm, nodeIDVector, std::move(outputVectors),
              std::move(outChunkState)} {}

    void setToTable(const transaction::Transaction* transaction, Table* table_,
        std::vector<common::column_id_t> columnIDs_,
        std::vector<ColumnPredicateSet> columnPredicateSets_ = {},
        common::RelDataDirection direction = common::RelDataDirection::INVALID) override;
};

struct ArrowNodeTableScanSharedState final : ColumnarNodeTableScanSharedState {
private:
    std::mutex mtx;
    std::vector<size_t> batchSizes;
    common::node_group_idx_t currentBatchIdx = 0;
    size_t currentMorselStartOffset = 0;
    const size_t morselSize;

public:
    ArrowNodeTableScanSharedState(const size_t morselSize)
        : ColumnarNodeTableScanSharedState(), morselSize(morselSize) {}

    void reset(std::vector<size_t> batchSizes) {
        std::lock_guard<std::mutex> lock(mtx);
        this->batchSizes = batchSizes;
        this->currentBatchIdx = 0;
        this->currentMorselStartOffset = 0;
    }

    bool getNextMorsel(ColumnarNodeTableScanState* scanState) override {
        auto arrowScanState = static_cast<ArrowNodeTableScanState*>(scanState);
        std::lock_guard<std::mutex> lock(mtx);

        while (currentBatchIdx < batchSizes.size()) {
            auto batchLength = batchSizes[currentBatchIdx];

            if (currentMorselStartOffset < batchLength) {
                arrowScanState->currentBatchIdx = currentBatchIdx;
                arrowScanState->currentMorselStartOffset = currentMorselStartOffset;
                arrowScanState->currentMorselEndOffset =
                    std::min(currentMorselStartOffset + morselSize, batchLength);
                this->currentMorselStartOffset = arrowScanState->currentMorselEndOffset;

                return true;
            }

            this->currentBatchIdx++;
            this->currentMorselStartOffset = 0;
        }

        return false;
    }
};

class ArrowNodeTable final : public ColumnarNodeTableBase {
public:
    ArrowNodeTable(const StorageManager* storageManager,
        const catalog::NodeTableCatalogEntry* nodeTableEntry, MemoryManager* memoryManager,
        ArrowSchemaWrapper schema, std::vector<ArrowArrayWrapper> arrays, std::string arrowId);

    ~ArrowNodeTable();

    void initializeScanCoordination(const transaction::Transaction* transaction) override;

    void initScanState(transaction::Transaction* transaction, TableScanState& scanState,
        bool resetCachedBoundNodeSelVec = true) const override;

    bool scanInternal(transaction::Transaction* transaction, TableScanState& scanState) override;

    const ArrowSchemaWrapper& getArrowSchema() const { return schema; }
    const std::vector<ArrowArrayWrapper>& getArrowArrays() const { return arrays; }

    common::node_group_idx_t getNumBatches(
        const transaction::Transaction* transaction) const override;

    size_t getNumScanMorsels(const transaction::Transaction* transaction) const;

    const catalog::NodeTableCatalogEntry* getCatalogEntry() const { return nodeTableCatalogEntry; }

protected:
    std::string getColumnarFormatName() const override { return "Arrow"; }
    common::row_idx_t getTotalRowCount(const transaction::Transaction* transaction) const override;

private:
    std::vector<size_t> getBatchSizes(
        [[maybe_unused]] const transaction::Transaction* transaction) const;

    void copyArrowMorselToOutputVectors(const ArrowArrayWrapper& batch,
        const size_t currentMorselStartOffset, const uint64_t numRowsToCopy,
        const std::vector<common::ValueVector*>& outputVectors,
        const std::vector<int64_t>& outputToArrowColumnIdx) const;

private:
    ArrowSchemaWrapper schema;
    std::vector<ArrowArrayWrapper> arrays;
    std::vector<size_t> batchStartOffsets;
    size_t totalRows;
    std::string arrowId;                           // ID in registry for cleanup
    constexpr static size_t scanMorselSize = 2048; // Default morsel size
};

} // namespace storage
} // namespace lbug