query_graph.cpp

#include "binder/query/query_graph.h"

#include "binder/expression_visitor.h"

using namespace kuzu::common;

namespace kuzu {
namespace binder {

std::size_t SubqueryGraphHasher::operator()(const SubqueryGraph& key) const {
    if (0 == key.queryRelsSelector.count()) {
        return std::hash<std::bitset<MAX_NUM_QUERY_VARIABLES>>{}(key.queryNodesSelector);
    }
    return std::hash<std::bitset<MAX_NUM_QUERY_VARIABLES>>{}(key.queryRelsSelector);
}

std::unordered_map<common::idx_t, std::vector<common::idx_t>>
SubqueryGraph::getWCOJRelCandidates() const {
    std::unordered_map<common::idx_t, std::vector<common::idx_t>> candidates;
    for (auto relPos : getRelNbrPositions()) {
        auto rel = queryGraph.getQueryRel(relPos);
        // TODO(Xiyang): is the following check relevant?
        if (!queryGraph.containsQueryNode(rel->getSrcNodeName()) ||
            !queryGraph.containsQueryNode(rel->getDstNodeName())) {
            continue;
        }
        auto srcNodePos = queryGraph.getQueryNodePos(rel->getSrcNodeName());
        auto dstNodePos = queryGraph.getQueryNodePos(rel->getDstNodeName());
        auto isSrcConnected = queryNodesSelector[srcNodePos];
        auto isDstConnected = queryNodesSelector[dstNodePos];
        // Closing rel should be handled with inner join.
        if (isSrcConnected && isDstConnected) {
            continue;
        }
        auto intersectNodePos = isSrcConnected ? dstNodePos : srcNodePos;
        if (!candidates.contains(intersectNodePos)) {
            candidates.insert({intersectNodePos, std::vector<common::idx_t>{}});
        }
        candidates.at(intersectNodePos).push_back(relPos);
    }
    return candidates;
}

bool SubqueryGraph::containAllVariables(std::unordered_set<std::string>& variables) const {
    for (auto& var : variables) {
        if (queryGraph.containsQueryNode(var) &&
            !queryNodesSelector[queryGraph.getQueryNodePos(var)]) {
            return false;
        }
        if (queryGraph.containsQueryRel(var) &&
            !queryRelsSelector[queryGraph.getQueryRelPos(var)]) {
            return false;
        }
    }
    return true;
}

std::unordered_set<uint32_t> SubqueryGraph::getNodeNbrPositions() const {
    std::unordered_set<uint32_t> result;
    for (auto relPos = 0u; relPos < queryGraph.getNumQueryRels(); ++relPos) {
        if (!queryRelsSelector[relPos]) { // rel not in subgraph, no need to check
            continue;
        }
        auto rel = queryGraph.getQueryRel(relPos);
        auto srcNodePos = queryGraph.getQueryNodePos(*rel->getSrcNode());
        if (!queryNodesSelector[srcNodePos]) {
            result.insert(srcNodePos);
        }
        auto dstNodePos = queryGraph.getQueryNodePos(*rel->getDstNode());
        if (!queryNodesSelector[dstNodePos]) {
            result.insert(dstNodePos);
        }
    }
    return result;
}

std::unordered_set<uint32_t> SubqueryGraph::getRelNbrPositions() const {
    std::unordered_set<uint32_t> result;
    for (auto relPos = 0u; relPos < queryGraph.getNumQueryRels(); ++relPos) {
        if (queryRelsSelector[relPos]) { // rel already in subgraph, cannot be rel neighbour
            continue;
        }
        auto rel = queryGraph.getQueryRel(relPos);
        auto srcNodePos = queryGraph.getQueryNodePos(*rel->getSrcNode());
        auto dstNodePos = queryGraph.getQueryNodePos(*rel->getDstNode());
        if (queryNodesSelector[srcNodePos] || queryNodesSelector[dstNodePos]) {
            result.insert(relPos);
        }
    }
    return result;
}

subquery_graph_set_t SubqueryGraph::getNbrSubgraphs(uint32_t size) const {
    auto result = getBaseNbrSubgraph();
    for (auto i = 1u; i < size; ++i) {
        std::unordered_set<SubqueryGraph, SubqueryGraphHasher> tmp;
        for (auto& prevNbr : result) {
            for (auto& subgraph : getNextNbrSubgraphs(prevNbr)) {
                tmp.insert(subgraph);
            }
        }
        result = std::move(tmp);
    }
    return result;
}

std::vector<uint32_t> SubqueryGraph::getConnectedNodePos(const SubqueryGraph& nbr) const {
    std::vector<uint32_t> result;
    for (auto& nodePos : getNodeNbrPositions()) {
        if (nbr.queryNodesSelector[nodePos]) {
            result.push_back(nodePos);
        }
    }
    for (auto& nodePos : nbr.getNodeNbrPositions()) {
        if (queryNodesSelector[nodePos]) {
            result.push_back(nodePos);
        }
    }
    return result;
}

std::unordered_set<uint32_t> SubqueryGraph::getNodePositionsIgnoringNodeSelector() const {
    std::unordered_set<uint32_t> result;
    for (auto nodePos = 0u; nodePos < queryGraph.getNumQueryNodes(); ++nodePos) {
        if (queryNodesSelector[nodePos]) {
            result.insert(nodePos);
        }
    }
    for (auto relPos = 0u; relPos < queryGraph.getNumQueryRels(); ++relPos) {
        auto rel = queryGraph.getQueryRel(relPos);
        if (queryRelsSelector[relPos]) {
            result.insert(queryGraph.getQueryNodePos(rel->getSrcNodeName()));
            result.insert(queryGraph.getQueryNodePos(rel->getDstNodeName()));
        }
    }
    return result;
}

subquery_graph_set_t SubqueryGraph::getBaseNbrSubgraph() const {
    subquery_graph_set_t result;
    for (auto& nodePos : getNodeNbrPositions()) {
        auto nbr = SubqueryGraph(queryGraph);
        nbr.addQueryNode(nodePos);
        result.insert(std::move(nbr));
    }
    for (auto& relPos : getRelNbrPositions()) {
        auto nbr = SubqueryGraph(queryGraph);
        nbr.addQueryRel(relPos);
        result.insert(std::move(nbr));
    }
    return result;
}

subquery_graph_set_t SubqueryGraph::getNextNbrSubgraphs(const SubqueryGraph& prevNbr) const {
    subquery_graph_set_t result;
    for (auto& nodePos : prevNbr.getNodeNbrPositions()) {
        if (queryNodesSelector[nodePos]) {
            continue;
        }
        auto nbr = prevNbr;
        nbr.addQueryNode(nodePos);
        result.insert(std::move(nbr));
    }
    for (auto& relPos : prevNbr.getRelNbrPositions()) {
        if (queryRelsSelector[relPos]) {
            continue;
        }
        auto nbr = prevNbr;
        nbr.addQueryRel(relPos);
        result.insert(std::move(nbr));
    }
    return result;
}

bool QueryGraph::isEmpty() const {
    for (auto& n : queryNodes) {
        if (n->isEmpty()) {
            return true;
        }
    }
    for (auto& r : queryRels) {
        if (r->isEmpty()) {
            return true;
        }
    }
    return false;
}

std::vector<std::shared_ptr<NodeOrRelExpression>> QueryGraph::getAllPatterns() const {
    std::vector<std::shared_ptr<NodeOrRelExpression>> patterns;
    for (auto& p : queryNodes) {
        patterns.push_back(p);
    }
    for (auto& p : queryRels) {
        patterns.push_back(p);
    }
    return patterns;
}

std::vector<std::shared_ptr<NodeExpression>> QueryGraph::getQueryNodes(
    const std::vector<idx_t>& indices) const {
    std::vector<std::shared_ptr<NodeExpression>> result;
    result.reserve(indices.size());
    for (auto idx : indices) {
        result.push_back(queryNodes[idx]);
    }
    return result;
}

std::vector<std::shared_ptr<RelExpression>> QueryGraph::getQueryRels(
    const std::vector<idx_t>& indices) const {
    std::vector<std::shared_ptr<RelExpression>> result;
    result.reserve(indices.size());
    for (auto idx : indices) {
        result.push_back(queryRels[idx]);
    }
    return result;
}

void QueryGraph::addQueryNode(std::shared_ptr<NodeExpression> queryNode) {
    // Note that a node may be added multiple times. We should only keep one of it.
    // E.g. MATCH (a:person)-[:knows]->(b:person), (a)-[:knows]->(c:person)
    // a will be added twice during binding
    if (containsQueryNode(queryNode->getUniqueName())) {
        return;
    }
    queryNodeNameToPosMap.insert({queryNode->getUniqueName(), queryNodes.size()});
    queryNodes.push_back(std::move(queryNode));
}

void QueryGraph::addQueryRel(std::shared_ptr<RelExpression> queryRel) {
    if (containsQueryRel(queryRel->getUniqueName())) {
        return;
    }
    queryRelNameToPosMap.insert({queryRel->getUniqueName(), queryRels.size()});
    queryRels.push_back(std::move(queryRel));
}

void QueryGraph::merge(const QueryGraph& other) {
    for (auto& otherNode : other.queryNodes) {
        addQueryNode(otherNode);
    }
    for (auto& otherRel : other.queryRels) {
        addQueryRel(otherRel);
    }
}

bool QueryGraph::canProjectExpression(const std::shared_ptr<Expression>& expression) const {
    auto expressionCollector = std::make_unique<ExpressionCollector>();
    for (auto& variable : expressionCollector->getDependentVariableNames(expression)) {
        if (!containsQueryNode(variable) && !containsQueryRel(variable)) {
            return false;
        }
    }
    return true;
}

bool QueryGraph::isConnected(const QueryGraph& other) {
    for (auto& queryNode : queryNodes) {
        if (other.containsQueryNode(queryNode->getUniqueName())) {
            return true;
        }
    }
    return false;
}

void QueryGraphCollection::addAndMergeQueryGraphIfConnected(QueryGraph queryGraphToAdd) {
    bool isMerged = false;
    for (auto& queryGraph : queryGraphs) {
        if (queryGraph.isConnected(queryGraphToAdd)) {
            queryGraph.merge(queryGraphToAdd);
            isMerged = true;
        }
    }
    if (!isMerged) {
        queryGraphs.push_back(std::move(queryGraphToAdd));
    }
}

void QueryGraphCollection::finalize() {
    common::idx_t baseGraphIdx = 0;
    while (true) {
        auto prevNumGraphs = queryGraphs.size();
        queryGraphs = mergeGraphs(baseGraphIdx++);
        if (queryGraphs.size() == prevNumGraphs || baseGraphIdx == queryGraphs.size()) {
            return;
        }
    }
}

std::vector<QueryGraph> QueryGraphCollection::mergeGraphs(common::idx_t baseGraphIdx) {
    KU_ASSERT(baseGraphIdx < queryGraphs.size());
    auto baseGraph = std::move(queryGraphs[baseGraphIdx]);
    std::unordered_set<common::idx_t> mergedGraphIndices;
    mergedGraphIndices.insert(baseGraphIdx);
    while (true) {
        // find graph to merge
        common::idx_t graphToMergeIdx = common::INVALID_IDX;
        for (auto i = 0u; i < queryGraphs.size(); ++i) {
            if (mergedGraphIndices.contains(i)) { // graph has been merged.
                continue;
            }
            if (baseGraph.isConnected(queryGraphs[i])) { // find graph to merge.
                graphToMergeIdx = i;
                break;
            }
        }
        if (graphToMergeIdx == common::INVALID_IDX) { // No graph can be merged. Terminate.
            break;
        }
        // Perform merge
        baseGraph.merge(queryGraphs[graphToMergeIdx]);
        mergedGraphIndices.insert(graphToMergeIdx);
    }
    std::vector<QueryGraph> finalGraphs;
    for (auto i = 0u; i < queryGraphs.size(); ++i) {
        if (i == baseGraphIdx) {
            finalGraphs.push_back(baseGraph);
            continue;
        }
        if (mergedGraphIndices.contains(i)) {
            continue;
        }
        finalGraphs.push_back(std::move(queryGraphs[i]));
    }
    return finalGraphs;
}

std::vector<std::shared_ptr<NodeExpression>> QueryGraphCollection::getQueryNodes() const {
    std::vector<std::shared_ptr<NodeExpression>> result;
    for (auto& queryGraph : queryGraphs) {
        for (auto& node : queryGraph.getQueryNodes()) {
            result.push_back(node);
        }
    }
    return result;
}

std::vector<std::shared_ptr<RelExpression>> QueryGraphCollection::getQueryRels() const {
    std::vector<std::shared_ptr<RelExpression>> result;
    for (auto& queryGraph : queryGraphs) {
        for (auto& rel : queryGraph.getQueryRels()) {
            result.push_back(rel);
        }
    }
    return result;
}

} // namespace binder
} // namespace kuzu