auto &column_ids = get.GetMutableColumnIds();

vector<idx_t> result;

for (idx_t r_idx = 0; r_idx < row_id_column_ids.size(); ++r_idx) {

// check if it is already projected

auto row_id_column_id = row_id_column_ids[r_idx];

auto &row_id_column = row_id_columns[r_idx];

optional_idx row_id_index;

for (idx_t i = 0; i < column_ids.size(); ++i) {

if (column_ids[i].GetPrimaryIndex() == row_id_column_id) {

// already projected - return the id

row_id_index = i;

break;

}

}

if (row_id_index.IsValid()) {

result.push_back(row_id_index.GetIndex());

continue;

}

// row id is not yet projected - push it and return the new index

column_ids.push_back(ColumnIndex(row_id_column_id));

if (!get.projection_ids.empty()) {

get.projection_ids.push_back(column_ids.size() - 1);

}

if (!get.types.empty()) {

get.types.push_back(row_id_column.type);

}

result.push_back(column_ids.size() - 1);

}

return result;

// we need to construct a new scan of the same table

auto table_index = optimizer.binder.GenerateTableIndex();

auto new_get = make_uniq<LogicalGet>(table_index, get.function, get.bind_data->Copy(), get.returned_types,

get.names, get.virtual_columns);

new_get->GetMutableColumnIds() = get.GetColumnIds();

new_get->projection_ids = get.projection_ids;

return new_get;

// traverse down until we reach the LogicalGet

vector<reference<LogicalOperator>> stack;

reference<LogicalOperator> child = *op->children[0];

while (child.get().type != LogicalOperatorType::LOGICAL_GET) {

stack.push_back(child);

D_ASSERT(child.get().children.size() == 1);

child = *child.get().children[0];

}

// we have reached the logical get - now we need to push the row-id column (if it is not yet projected out)

auto &get = child.get().Cast<LogicalGet>();

auto row_id_indexes = GetOrInsertRowIds(get);

idx_t column_count = get.projection_ids.empty() ? get.GetColumnIds().size() : get.projection_ids.size();

D_ASSERT(column_count == get.GetColumnBindings().size());

// the row id has been projected - now project it up the stack

vector<ColumnBinding> row_id_bindings;

for (auto &row_id_index : row_id_indexes) {

row_id_bindings.emplace_back(get.table_index, row_id_index);

}

for (idx_t i = stack.size(); i > 0; i--) {

auto &op = stack[i - 1].get();

switch (op.type) {

case LogicalOperatorType::LOGICAL_PROJECTION: {

auto &proj = op.Cast<LogicalProjection>();

// push projection of the row-id columns

for (idx_t r_idx = 0; r_idx < row_id_columns.size(); r_idx++) {

auto &r_col = row_id_columns[r_idx];

proj.expressions.push_back(

make_uniq<BoundColumnRefExpression>(r_col.name, r_col.type, row_id_bindings[r_idx]));

// modify the row-id-binding to the new projection

row_id_bindings[r_idx] = ColumnBinding(proj.table_index, proj.expressions.size() - 1);

}

column_count = proj.expressions.size();

break;

}

case LogicalOperatorType::LOGICAL_FILTER: {

auto &filter = op.Cast<LogicalFilter>();

// column bindings pass-through this operator as-is UNLESS the filter has a projection map

if (filter.HasProjectionMap()) {

// if the filter has a projection map, we need to project the new column

filter.projection_map.push_back(column_count - 1);

}

break;

}

default:

throw InternalException("Unsupported logical operator in LateMaterialization::ConstructRHS");

}

}

return row_id_bindings;

reference<LogicalOperator> current_op = root;

while (true) {

auto &op = current_op.get();

switch (op.type) {

case LogicalOperatorType::LOGICAL_PROJECTION: {

// reached a projection - modify the table index and return

auto &proj = op.Cast<LogicalProjection>();

proj.table_index = new_index;

return;

}

case LogicalOperatorType::LOGICAL_GET: {

// reached the root get - modify the table index and return

auto &get = op.Cast<LogicalGet>();

get.table_index = new_index;

return;

}

case LogicalOperatorType::LOGICAL_TOP_N: {

// visit the expressions of the operator and continue into the child node

auto &top_n = op.Cast<LogicalTopN>();

for (auto &order : top_n.orders) {

ReplaceTableReferences(*order.expression, new_index);

}

current_op = *op.children[0];

break;

}

case LogicalOperatorType::LOGICAL_FILTER:

case LogicalOperatorType::LOGICAL_SAMPLE:

case LogicalOperatorType::LOGICAL_LIMIT: {

// visit the expressions of the operator and continue into the child node

for (auto &expr : op.expressions) {

ReplaceTableReferences(*expr, new_index);

}

current_op = *op.children[0];

break;

}

default:

throw InternalException("Unsupported operator type in LateMaterialization::ReplaceTopLevelTableIndex");

}

}

if (expr.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) {

auto &bound_column_ref = expr.Cast<BoundColumnRefExpression>();

bound_column_ref.binding.table_index = new_table_index;

}

ExpressionIterator::EnumerateChildren(expr,

[&](Expression &child) { ReplaceTableReferences(child, new_table_index); });

switch (op.type) {

case LogicalOperatorType::LOGICAL_GET: {

auto &get = op.Cast<LogicalGet>();

auto &column_id = get.GetColumnIds()[column_index];

auto column_name = get.GetColumnName(column_id);

auto &column_type = get.GetColumnType(column_id);

auto expr =

make_uniq<BoundColumnRefExpression>(column_name, column_type, ColumnBinding(get.table_index, column_index));

return std::move(expr);

}

case LogicalOperatorType::LOGICAL_PROJECTION:

return op.expressions[column_index]->Copy();

default:

throw InternalException("Unsupported operator type for LateMaterialization::GetExpression");

}

if (expr->GetExpressionType() == ExpressionType::BOUND_COLUMN_REF) {

auto &bound_column_ref = expr->Cast<BoundColumnRefExpression>();

expr = GetExpression(next_op, bound_column_ref.binding.column_index);

return;

}

ExpressionIterator::EnumerateChildren(

*expr, [&](unique_ptr<Expression> &child) { ReplaceExpressionReferences(next_op, child); });

// check if we can benefit from late materialization

// we need to see how many columns we require in the pipeline versus how many columns we emit in the scan

// for example, in a query like SELECT * FROM tbl ORDER BY ts LIMIT 5, the top-n only needs the "ts" column

// the other columns can be fetched later on using late materialization

// we can only push late materialization through a subset of operators

// and we can only do it for scans that support the row-id pushdown (currently only DuckDB table scans)

// visit the expressions for each operator in the chain

vector<reference<LogicalOperator>> source_operators;

VisitOperatorExpressions(*op);

reference<LogicalOperator> child = *op->children[0];

while (child.get().type != LogicalOperatorType::LOGICAL_GET) {

switch (child.get().type) {

case LogicalOperatorType::LOGICAL_PROJECTION: {

// recurse into the child node - but ONLY visit expressions that are referenced

auto &proj = child.get().Cast<LogicalProjection>();

source_operators.push_back(child);

for (auto &expr : proj.expressions) {

if (expr->IsVolatile()) {

// we cannot do this optimization if any of the columns are volatile

return false;

}

}

// figure out which projection expressions we are currently referencing

set<idx_t> referenced_columns;

for (auto &entry : column_references) {

auto &column_binding = entry.first;

if (column_binding.table_index == proj.table_index) {

referenced_columns.insert(column_binding.column_index);

}

}

// clear the list of referenced expressions and visit those columns

column_references.clear();

for (auto &col_idx : referenced_columns) {

VisitExpression(&proj.expressions[col_idx]);

}

// continue into child

child = *child.get().children[0];

break;

}

case LogicalOperatorType::LOGICAL_FILTER: {

// visit filter expressions - we need these columns

VisitOperatorExpressions(child.get());

// continue into child

child = *child.get().children[0];

break;

}

default:

// unsupported operator for late materialization

return false;

}

}

auto &get = child.get().Cast<LogicalGet>();

if (column_references.size() >= get.GetColumnIds().size()) {

// we do not benefit from late materialization

// we need all of the columns to compute the root node anyway (Top-N/Limit/etc)

return false;

}

if (!get.function.late_materialization) {

// this function does not support late materialization

return false;

}

if (!get.function.get_row_id_columns) {

throw InternalException("Function supports late materialization but not get_row_id_columns");

}

row_id_column_ids = get.function.get_row_id_columns(optimizer.context, get.bind_data.get());

if (row_id_column_ids.empty()) {

throw InternalException("Row Id Columns must not be empty");

}

row_id_columns.clear();

for (auto &col_id : row_id_column_ids) {

auto entry = get.virtual_columns.find(col_id);

if (entry == get.virtual_columns.end()) {

throw InternalException("Row id column id not found in virtual column list");

}

row_id_columns.push_back(entry->second);

}

// we benefit from late materialization

// we need to transform this plan into a semi-join with the row-id

// we need to ensure the operator returns exactly the same column bindings as before

// construct the LHS from the LogicalGet

auto lhs = ConstructLHS(get);

// insert the row-id column on the left hand side

auto &lhs_get = *lhs;

auto lhs_index = lhs_get.table_index;

auto lhs_columns = lhs_get.GetColumnIds().size();

auto lhs_row_indexes = GetOrInsertRowIds(lhs_get);

vector<ColumnBinding> lhs_bindings;

for (auto &lhs_row_index : lhs_row_indexes) {

lhs_bindings.emplace_back(lhs_index, lhs_row_index);

}

// after constructing the LHS but before constructing the RHS we construct the final projections/orders

// - we do this before constructing the RHS because that alter the original plan

vector<unique_ptr<Expression>> final_proj_list;

// construct the final projection list from either (1) the root projection, or (2) the logical get

if (!source_operators.empty()) {

// construct the columns from the root projection

auto &root_proj = source_operators[0].get();

for (auto &expr : root_proj.expressions) {

final_proj_list.push_back(expr->Copy());

}

// now we need to "flatten" the projection list by traversing the set of projections and inlining them

for (idx_t i = 0; i < source_operators.size(); i++) {

auto &next_operator = i + 1 < source_operators.size() ? source_operators[i + 1].get() : lhs_get;

for (auto &expr : final_proj_list) {

ReplaceExpressionReferences(next_operator, expr);

}

}

} else {

// if we have no projection directly construct the columns from the root get

for (idx_t i = 0; i < lhs_columns; i++) {

final_proj_list.push_back(GetExpression(lhs_get, i));

}

}

// we need to re-order again at the end

vector<BoundOrderByNode> final_orders;

auto root_type = op->type;

if (root_type == LogicalOperatorType::LOGICAL_TOP_N) {

// for top-n we need to re-order by the top-n conditions

auto &top_n = op->Cast<LogicalTopN>();

for (auto &order : top_n.orders) {

auto expr = order.expression->Copy();

final_orders.emplace_back(order.type, order.null_order, std::move(expr));

}

} else {

// for limit/sample we order by row-id

for (idx_t r_idx = 0; r_idx < row_id_columns.size(); r_idx++) {

auto &row_id_col = row_id_columns[r_idx];

auto row_id_expr =

make_uniq<BoundColumnRefExpression>(row_id_col.name, row_id_col.type, lhs_bindings[r_idx]);

final_orders.emplace_back(OrderType::ASCENDING, OrderByNullType::NULLS_LAST, std::move(row_id_expr));

}

}

// construct the RHS for the join

// this is essentially the old pipeline, but with the `rowid` column added

// note that the purpose of this optimization is to remove columns from the RHS

// we don't do that here yet though - we do this in a later step using the RemoveUnusedColumns optimizer

auto rhs_bindings = ConstructRHS(op);

// the final table index emitted must be the table index of the original operator

// this ensures any upstream operators that refer to the original get will keep on referring to the correct columns

auto final_index = rhs_bindings[0].table_index;

// we need to replace any references to "rhs_binding.table_index" in the rhs to a new table index

auto rhs_table_index = optimizer.binder.GenerateTableIndex();

for (auto &rhs_binding : rhs_bindings) {

rhs_binding.table_index = rhs_table_index;

}

ReplaceTopLevelTableIndex(*op, rhs_table_index);

if (rhs_bindings.size() != lhs_bindings.size()) {

throw InternalException("Mismatch in late materialization binding sizes");

}

// construct a semi join between the lhs and rhs

auto join = make_uniq<LogicalComparisonJoin>(JoinType::SEMI);

join->children.push_back(std::move(lhs));

join->children.push_back(std::move(op));

for (idx_t r_idx = 0; r_idx < row_id_columns.size(); r_idx++) {

auto &row_id_col = row_id_columns[r_idx];

JoinCondition condition;

condition.comparison = ExpressionType::COMPARE_EQUAL;

condition.left = make_uniq<BoundColumnRefExpression>(row_id_col.name, row_id_col.type, lhs_bindings[r_idx]);

condition.right = make_uniq<BoundColumnRefExpression>(row_id_col.name, row_id_col.type, rhs_bindings[r_idx]);

join->conditions.push_back(std::move(condition));

}

// push a projection that removes the row id again from the lhs

// this is the final projection - so it should have the final table index

auto proj_index = final_index;

if (root_type == LogicalOperatorType::LOGICAL_TOP_N) {

// for top-n we need to order on expressions, so we need to order AFTER the final projection

auto proj = make_uniq<LogicalProjection>(proj_index, std::move(final_proj_list));

if (join->has_estimated_cardinality) {

proj->SetEstimatedCardinality(join->estimated_cardinality);

}

proj->children.push_back(std::move(join));

for (auto &order : final_orders) {

ReplaceTableReferences(*order.expression, proj_index);

}

auto order = make_uniq<LogicalOrder>(std::move(final_orders));

if (proj->has_estimated_cardinality) {

order->SetEstimatedCardinality(proj->estimated_cardinality);

}

order->children.push_back(std::move(proj));

op = std::move(order);

} else {

// for limit/sample we order on row-id, so we need to order BEFORE the final projection

// because the final projection removes row-ids

auto order = make_uniq<LogicalOrder>(std::move(final_orders));

if (join->has_estimated_cardinality) {

order->SetEstimatedCardinality(join->estimated_cardinality);

}

order->children.push_back(std::move(join));

auto proj = make_uniq<LogicalProjection>(proj_index, std::move(final_proj_list));

if (order->has_estimated_cardinality) {

proj->SetEstimatedCardinality(order->estimated_cardinality);

}

proj->children.push_back(std::move(order));

op = std::move(proj);

}

// run the RemoveUnusedColumns optimizer to prune the (now) unused columns the plan

RemoveUnusedColumns unused_optimizer(optimizer.binder, optimizer.context, true);

unused_optimizer.VisitOperator(*op);

return true;

auto &config = DBConfig::GetConfig(optimizer.context);

if (!has_offset && !config.options.preserve_insertion_order) {

// we avoid optimizing large limits if preserve insertion order is false

// since the limit is executed in parallel anyway

return false;

}

// we only perform this optimization until a certain amount of maximum values to reduce memory constraints

// since we still materialize the set of row-ids in the hash table this optimization can increase memory pressure

// FIXME: make this configurable as well

static constexpr const idx_t LIMIT_MAX_VAL = 1000000;

if (limit_val > LIMIT_MAX_VAL) {

return false;

}

// we only support large limits if they are directly below the source

reference<LogicalOperator> current_op = *limit.children[0];

while (current_op.get().type != LogicalOperatorType::LOGICAL_GET) {

if (current_op.get().type != LogicalOperatorType::LOGICAL_PROJECTION) {

return false;

}

current_op = *current_op.get().children[0];

}

return true;

switch (op->type) {

case LogicalOperatorType::LOGICAL_LIMIT: {

auto &limit = op->Cast<LogicalLimit>();

if (limit.limit_val.Type() != LimitNodeType::CONSTANT_VALUE) {

break;

}

auto limit_val = limit.limit_val.GetConstantValue();

bool has_offset = limit.offset_val.Type() != LimitNodeType::UNSET;

if (limit_val > max_row_count) {

// for large limits - we may still want to do this optimization if the limit is consecutive

// this is the case if there are only projections/get below the limit

// if the row-ids are not consecutive doing the join can worsen performance

if (!OptimizeLargeLimit(limit, limit_val, has_offset)) {

break;

}

} else {

// optimizing small limits really only makes sense if we have an offset

if (!has_offset) {

break;

}

}

if (TryLateMaterialization(op)) {

return op;

}

break;

}

case LogicalOperatorType::LOGICAL_TOP_N: {

auto &top_n = op->Cast<LogicalTopN>();

if (top_n.limit > max_row_count) {

break;

}

// for the top-n we need to visit the order elements

if (TryLateMaterialization(op)) {

return op;

}

break;

}

case LogicalOperatorType::LOGICAL_SAMPLE: {

auto &sample = op->Cast<LogicalSample>();

if (sample.sample_options->is_percentage) {

break;

}

if (sample.sample_options->sample_size.GetValue<uint64_t>() > max_row_count) {

break;

}

if (TryLateMaterialization(op)) {

return op;

}

break;

}

default:

break;

}

for (auto &child : op->children) {

child = Optimize(std::move(child));

}

return op;