#include "binder/expression/rel_expression.h"
#include "catalog/catalog.h"
#include "catalog/catalog_entry/node_table_catalog_entry.h"
#include "common/partition_routing.h"
#include "common/partition_routing_hook.h"
#include "main/client_context.h"
#include "planner/operator/persistent/logical_insert.h"
#include "processor/expression_mapper.h"
#include "processor/operator/persistent/insert.h"
#include "processor/plan_mapper.h"
#include "storage/partition_storage_registry.h"
#include "storage/storage_manager.h"
#include "transaction/transaction.h"
using namespace lbug::evaluator;
using namespace lbug::planner;
using namespace lbug::storage;
using namespace lbug::catalog;
using namespace lbug::common;
using namespace lbug::binder;
namespace lbug {
namespace processor {
static std::vector<DataPos> populateReturnColumnsPos(const LogicalInsertInfo& info,
const Schema& schema) {
std::vector<DataPos> result;
for (auto i = 0u; i < info.columnDataExprs.size(); ++i) {
if (info.isReturnColumnExprs[i]) {
result.emplace_back(schema.getExpressionPos(*info.columnExprs[i]));
} else {
result.push_back(DataPos::getInvalidPos());
}
}
return result;
}
NodeInsertExecutor PlanMapper::getNodeInsertExecutor(const LogicalInsertInfo* boundInfo,
const Schema& inSchema, const Schema& outSchema) const {
auto& node = boundInfo->pattern->constCast<NodeExpression>();
auto nodeIDPos = getDataPos(*node.getInternalID(), outSchema);
auto columnsPos = populateReturnColumnsPos(*boundInfo, outSchema);
auto info = NodeInsertInfo(nodeIDPos, columnsPos, boundInfo->conflictAction);
evaluator_vector_t evaluators;
auto exprMapper = ExpressionMapper(&inSchema);
for (auto& expr : boundInfo->columnDataExprs) {
evaluators.push_back(exprMapper.getEvaluator(expr));
}
const auto* firstEntry = node.getEntry(0)->ptrCast<NodeTableCatalogEntry>();
bool isListParent = false;
if (firstEntry->isPartitionChild()) {
auto transaction = transaction::Transaction::Get(*clientContext);
const auto* parentEntry =
Catalog::Get(*clientContext)
->getTableCatalogEntry(transaction, firstEntry->getParentTableID())
->ptrCast<NodeTableCatalogEntry>();
isListParent = parentEntry->isPartitioned() &&
static_cast<common::PartitionMethod>(*parentEntry->getPartitionMethod()) ==
common::PartitionMethod::LIST;
}
const bool parentPattern =
firstEntry->isPartitionChild() &&
(node.getNumEntries() > 1 || firstEntry->getScanFunction().has_value() || isListParent);
if (parentPattern) {
const auto parentID = firstEntry->getParentTableID();
DASSERT(parentID != INVALID_TABLE_ID);
auto transaction = transaction::Transaction::Get(*clientContext);
const auto* parent = Catalog::Get(*clientContext)
->getTableCatalogEntry(transaction, parentID)
->ptrCast<NodeTableCatalogEntry>();
DASSERT(parent->isPartitioned());
const auto childTableIDs = parent->getChildTableIDs();
DASSERT(!childTableIDs.empty());
const auto* hooks = common::getPartitionRoutingHooks();
common::PartitionHandle handle = nullptr;
const bool firstClaimed =
hooks != nullptr && hooks->locate != nullptr &&
hooks->locate(hooks->context, common::PartitionRef{parentID, 0}, &handle);
auto firstTable = firstClaimed ? nullptr :
storage::PartitionStorageRegistry::resolveNodeTableByID(
clientContext, childTableIDs[0]);
auto tableInfo = NodeTableInsertInfo(firstTable, std::move(evaluators));
tableInfo.partitionKeyColumnID = parent->getPartitionColumnID();
tableInfo.partitionMethod =
static_cast<common::PartitionMethod>(*parent->getPartitionMethod());
tableInfo.parentTableID = parentID;
if (tableInfo.partitionMethod != common::PartitionMethod::LIST) {
tableInfo.partitionTables.reserve(childTableIDs.size());
for (auto i = 0u; i < childTableIDs.size(); ++i) {
const auto ref = common::PartitionRef{parentID, i};
common::PartitionHandle partHandle = nullptr;
const bool claimed = hooks != nullptr && hooks->locate != nullptr &&
hooks->locate(hooks->context, ref, &partHandle);
tableInfo.partitionTables.push_back(
claimed ? nullptr :
storage::PartitionStorageRegistry::resolveNodeTableByID(clientContext,
childTableIDs[i]));
tableInfo.partitionChildIDs.push_back(childTableIDs[i]);
tableInfo.partitionRefs.push_back(ref);
tableInfo.partitionHandles.push_back(claimed ? partHandle : nullptr);
}
}
return NodeInsertExecutor(std::move(info), std::move(tableInfo));
}
auto table =
storage::PartitionStorageRegistry::resolveNodeTable(clientContext, *node.getEntry(0));
auto tableInfo = NodeTableInsertInfo(table, std::move(evaluators));
return NodeInsertExecutor(std::move(info), std::move(tableInfo));
}
RelInsertExecutor PlanMapper::getRelInsertExecutor(const LogicalInsertInfo* boundInfo,
const Schema& inSchema, const Schema& outSchema) const {
auto& rel = boundInfo->pattern->constCast<RelExpression>();
auto srcNode = rel.getSrcNode();
auto dstNode = rel.getDstNode();
auto srcNodeIDPos = getDataPos(*srcNode->getInternalID(), inSchema);
auto dstNodeIDPos = getDataPos(*dstNode->getInternalID(), inSchema);
auto columnsPos = populateReturnColumnsPos(*boundInfo, outSchema);
auto info = RelInsertInfo(srcNodeIDPos, dstNodeIDPos, std::move(columnsPos));
auto storageManager = StorageManager::Get(*clientContext);
DASSERT(srcNode->getNumEntries() == 1 && dstNode->getNumEntries() == 1);
auto srcTableID = srcNode->getEntry(0)->getTableID();
auto dstTableID = dstNode->getEntry(0)->getTableID();
DASSERT(rel.getNumEntries() == 1);
auto& relGroupEntry = rel.getEntry(0)->constCast<RelGroupCatalogEntry>();
auto relEntryInfo = relGroupEntry.getRelEntryInfo(srcTableID, dstTableID);
auto table = storageManager->getTable(relEntryInfo->oid)->ptrCast<RelTable>();
evaluator_vector_t evaluators;
auto exprMapper = ExpressionMapper(&outSchema);
for (auto& expr : boundInfo->columnDataExprs) {
evaluators.push_back(exprMapper.getEvaluator(expr));
}
auto tableInfo = RelTableInsertInfo(table, std::move(evaluators));
return RelInsertExecutor(std::move(info), std::move(tableInfo));
}
std::unique_ptr<PhysicalOperator> PlanMapper::mapInsert(const LogicalOperator* logicalOperator) {
auto& logicalInsert = logicalOperator->constCast<LogicalInsert>();
auto inSchema = logicalInsert.getChild(0)->getSchema();
auto outSchema = logicalInsert.getSchema();
auto prevOperator = mapOperator(logicalOperator->getChild(0).get());
std::vector<NodeInsertExecutor> nodeExecutors;
std::vector<RelInsertExecutor> relExecutors;
for (auto& info : logicalInsert.getInfos()) {
switch (info.tableType) {
case TableType::NODE: {
nodeExecutors.push_back(getNodeInsertExecutor(&info, *inSchema, *outSchema));
} break;
case TableType::REL: {
relExecutors.push_back(getRelInsertExecutor(&info, *inSchema, *outSchema));
} break;
default:
UNREACHABLE_CODE;
}
}
expression_vector expressions;
for (auto& info : logicalInsert.getInfos()) {
for (auto& expr : info.columnExprs) {
expressions.push_back(expr);
}
}
auto printInfo =
std::make_unique<InsertPrintInfo>(expressions, logicalInsert.getInfos()[0].conflictAction);
return std::make_unique<Insert>(std::move(nodeExecutors), std::move(relExecutors),
std::move(prevOperator), getOperatorID(), std::move(printInfo));
}
} }