mask.h

Go to the documentation of this file.

#pragma once
 
#include "types.h"
 
namespace kuzu {
namespace common {
 
// Note that this class is NOT thread-safe.
class SemiMask {
public:
    explicit SemiMask(offset_t maxOffset) : maxOffset{maxOffset}, enabled{false} {}
 
    virtual ~SemiMask() = default;
 
    virtual void mask(offset_t nodeOffset) = 0;
    virtual void maskRange(offset_t startNodeOffset, offset_t endNodeOffset) = 0;
 
    virtual bool isMasked(offset_t startNodeOffset) = 0;
 
    // include&exclude
    virtual offset_vec_t range(uint32_t start, uint32_t end) = 0;
 
    virtual uint64_t getNumMaskedNodes() const = 0;
 
    virtual offset_vec_t collectMaskedNodes(uint64_t size) const = 0;
 
    offset_t getMaxOffset() const { return maxOffset; }
 
    bool isEnabled() const { return enabled; }
    void enable() { enabled = true; }
 
private:
    offset_t maxOffset;
    bool enabled;
};
 
struct SemiMaskUtil {
    KUZU_API static std::unique_ptr<SemiMask> createMask(offset_t maxOffset);
};
 
class NodeOffsetMaskMap {
public:
    NodeOffsetMaskMap() = default;
 
    offset_t getNumMaskedNode() const;
 
    void addMask(table_id_t tableID, std::unique_ptr<SemiMask> mask) {
        KU_ASSERT(!maskMap.contains(tableID));
        maskMap.insert({tableID, std::move(mask)});
    }
 
    table_id_map_t<SemiMask*> getMasks() const {
        table_id_map_t<SemiMask*> result;
        for (auto& [tableID, mask] : maskMap) {
            result.emplace(tableID, mask.get());
        }
        return result;
    }
 
    bool containsTableID(table_id_t tableID) const { return maskMap.contains(tableID); }
    SemiMask* getOffsetMask(table_id_t tableID) const {
        KU_ASSERT(containsTableID(tableID));
        return maskMap.at(tableID).get();
    }
 
    void pin(table_id_t tableID) {
        if (maskMap.contains(tableID)) {
            pinnedMask = maskMap.at(tableID).get();
        } else {
            pinnedMask = nullptr;
        }
    }
    bool hasPinnedMask() const { return pinnedMask != nullptr; }
    SemiMask* getPinnedMask() const { return pinnedMask; }
 
    bool valid(offset_t offset) const {
        KU_ASSERT(pinnedMask != nullptr);
        return pinnedMask->isMasked(offset);
    }
    bool valid(nodeID_t nodeID) const {
        KU_ASSERT(maskMap.contains(nodeID.tableID));
        return maskMap.at(nodeID.tableID)->isMasked(nodeID.offset);
    }
 
private:
    table_id_map_t<std::unique_ptr<SemiMask>> maskMap;
    SemiMask* pinnedMask = nullptr;
};
 
} // namespace common
} // namespace kuzu