duckdb 1.1.2-dev6.0 → 1.1.4-dev11.0

package/src/duckdb/src/optimizer/sum_rewriter.cpp

@@ -0,0 +1,174 @@
+#include "duckdb/optimizer/sum_rewriter.hpp"
+
+#include "duckdb/planner/expression/bound_aggregate_expression.hpp"
+#include "duckdb/planner/expression/bound_cast_expression.hpp"
+#include "duckdb/planner/expression/bound_columnref_expression.hpp"
+#include "duckdb/planner/expression/bound_function_expression.hpp"
+#include "duckdb/planner/operator/logical_aggregate.hpp"
+#include "duckdb/planner/operator/logical_projection.hpp"
+#include "duckdb/optimizer/matcher/expression_matcher.hpp"
+#include "duckdb/function/aggregate/distributive_function_utils.hpp"
+#include "duckdb/function/function_binder.hpp"
+#include "duckdb/planner/binder.hpp"
+#include "duckdb/optimizer/optimizer.hpp"
+
+namespace duckdb {
+
+unique_ptr<SetTypesMatcher> GetSmallIntegerTypesMatcher() {
+	vector<LogicalType> types {LogicalTypeId::TINYINT,  LogicalTypeId::SMALLINT, LogicalTypeId::INTEGER,
+	                           LogicalTypeId::BIGINT,   LogicalTypeId::UTINYINT, LogicalTypeId::USMALLINT,
+	                           LogicalTypeId::UINTEGER, LogicalTypeId::UBIGINT};
+	return make_uniq<SetTypesMatcher>(std::move(types));
+}
+
+SumRewriterOptimizer::SumRewriterOptimizer(Optimizer &optimizer) : optimizer(optimizer) {
+	// set up an expression matcher that detects SUM(x + C) or SUM(C + x)
+	auto op = make_uniq<AggregateExpressionMatcher>();
+	op->function = make_uniq<SpecificFunctionMatcher>("sum");
+	op->policy = SetMatcher::Policy::UNORDERED;
+
+	auto arithmetic = make_uniq<FunctionExpressionMatcher>();
+	// handle X + C where
+	arithmetic->function = make_uniq<SpecificFunctionMatcher>("+");
+	// we match only on integral numeric types
+	arithmetic->type = make_uniq<IntegerTypeMatcher>();
+	auto child_constant_matcher = make_uniq<ConstantExpressionMatcher>();
+	auto child_expression_matcher = make_uniq<StableExpressionMatcher>();
+	child_constant_matcher->type = GetSmallIntegerTypesMatcher();
+	child_expression_matcher->type = GetSmallIntegerTypesMatcher();
+	arithmetic->matchers.push_back(std::move(child_constant_matcher));
+	arithmetic->matchers.push_back(std::move(child_expression_matcher));
+	arithmetic->policy = SetMatcher::Policy::SOME;
+	op->matchers.push_back(std::move(arithmetic));
+
+	sum_matcher = std::move(op);
+}
+
+SumRewriterOptimizer::~SumRewriterOptimizer() {
+}
+
+void SumRewriterOptimizer::Optimize(unique_ptr<LogicalOperator> &op) {
+	if (op->type == LogicalOperatorType::LOGICAL_AGGREGATE_AND_GROUP_BY) {
+		RewriteSums(op);
+	}
+	VisitOperator(*op);
+}
+
+void SumRewriterOptimizer::StandardVisitOperator(LogicalOperator &op) {
+	for (auto &child : op.children) {
+		Optimize(child);
+	}
+	if (!aggregate_map.empty()) {
+		VisitOperatorExpressions(op);
+	}
+}
+
+void SumRewriterOptimizer::VisitOperator(LogicalOperator &op) {
+	switch (op.type) {
+	case LogicalOperatorType::LOGICAL_UNION:
+	case LogicalOperatorType::LOGICAL_EXCEPT:
+	case LogicalOperatorType::LOGICAL_INTERSECT:
+	case LogicalOperatorType::LOGICAL_MATERIALIZED_CTE:
+	case LogicalOperatorType::LOGICAL_PROJECTION: {
+		SumRewriterOptimizer sum_rewriter(optimizer);
+		sum_rewriter.StandardVisitOperator(op);
+		return;
+	}
+	default:
+		break;
+	}
+
+	StandardVisitOperator(op);
+}
+
+unique_ptr<Expression> SumRewriterOptimizer::VisitReplace(BoundColumnRefExpression &expr,
+                                                          unique_ptr<Expression> *expr_ptr) {
+	// check if this column ref points to an aggregate that was remapped; if it does we remap it
+	auto entry = aggregate_map.find(expr.binding);
+	if (entry != aggregate_map.end()) {
+		expr.binding = entry->second;
+	}
+	return nullptr;
+}
+
+void SumRewriterOptimizer::RewriteSums(unique_ptr<LogicalOperator> &op) {
+	auto &aggr = op->Cast<LogicalAggregate>();
+	if (!aggr.groups.empty()) {
+		return;
+	}
+
+	unordered_set<idx_t> rewrote_map;
+	vector<unique_ptr<Expression>> constants;
+	idx_t aggr_count = aggr.expressions.size();
+	for (idx_t i = 0; i < aggr_count; ++i) {
+		auto &expr = aggr.expressions[i];
+		vector<reference<Expression>> bindings;
+		if (!sum_matcher->Match(*expr, bindings)) {
+			continue;
+		}
+		// found SUM(x + C)
+		auto &sum = bindings[0].get().Cast<BoundAggregateExpression>();
+		auto &addition = bindings[1].get().Cast<BoundFunctionExpression>();
+		idx_t const_idx = addition.children[0]->GetExpressionType() == ExpressionType::VALUE_CONSTANT ? 0 : 1;
+		auto const_expr = std::move(addition.children[const_idx]);
+		auto main_expr = std::move(addition.children[1 - const_idx]);
+
+		// turn this into SUM(x)
+		sum.children[0] = main_expr->Copy();
+
+		// push a new aggregate - COUNT(x)
+		FunctionBinder function_binder(optimizer.context);
+
+		auto count_fun = CountFunctionBase::GetFunction();
+		vector<unique_ptr<Expression>> children;
+		children.push_back(std::move(main_expr));
+		auto count_aggr =
+		    function_binder.BindAggregateFunction(count_fun, std::move(children), nullptr, AggregateType::NON_DISTINCT);
+		aggr.expressions.push_back(std::move(count_aggr));
+		constants.push_back(std::move(const_expr));
+		rewrote_map.insert(i);
+	}
+	if (rewrote_map.empty()) {
+		return;
+	}
+	vector<unique_ptr<Expression>> projection_expressions;
+	// we rewrote aggregates - we need to push a projection in which we re-compute the original result
+	idx_t rewritten_index = 0;
+	auto proj_index = optimizer.binder.GenerateTableIndex();
+	for (idx_t i = 0; i < aggr_count; i++) {
+		ColumnBinding aggregate_binding(aggr.aggregate_index, i);
+		aggregate_map[aggregate_binding] = ColumnBinding(proj_index, i);
+		auto &aggr_type = aggr.expressions[i]->return_type;
+		auto aggr_ref = make_uniq<BoundColumnRefExpression>(aggr_type, aggregate_binding);
+		if (rewrote_map.find(i) == rewrote_map.end()) {
+			// not rewritten - just push a reference
+			projection_expressions.push_back(std::move(aggr_ref));
+			continue;
+		}
+		// rewritten - need to compute the final result
+		idx_t count_idx = aggr_count + rewritten_index;
+		ColumnBinding count_binding(aggr.aggregate_index, count_idx);
+		auto count_ref = make_uniq<BoundColumnRefExpression>(aggr.expressions[count_idx]->return_type, count_binding);
+
+		// cast the count to the sum type
+		auto cast_count = BoundCastExpression::AddCastToType(optimizer.context, std::move(count_ref), aggr_type);
+		auto const_expr =
+		    BoundCastExpression::AddCastToType(optimizer.context, std::move(constants[rewritten_index]), aggr_type);
+
+		// bind the multiplication
+		auto multiply = optimizer.BindScalarFunction("*", std::move(cast_count), std::move(const_expr));
+
+		// add it to the sum
+		auto final_result = optimizer.BindScalarFunction("+", std::move(aggr_ref), std::move(multiply));
+		projection_expressions.push_back(std::move(final_result));
+
+		rewritten_index++;
+	}
+
+	// push the projection to replace the aggregate
+	auto proj = make_uniq<LogicalProjection>(proj_index, std::move(projection_expressions));
+	proj->children.push_back(std::move(op));
+	op = std::move(proj);
+}
+
+} // namespace duckdb

package/src/duckdb/src/optimizer/topn_optimizer.cpp

@@ -1,9 +1,16 @@
 #include "duckdb/optimizer/topn_optimizer.hpp"
 
 #include "duckdb/common/limits.hpp"
+#include "duckdb/planner/operator/logical_get.hpp"
 #include "duckdb/planner/operator/logical_limit.hpp"
 #include "duckdb/planner/operator/logical_order.hpp"
 #include "duckdb/planner/operator/logical_top_n.hpp"
+#include "duckdb/planner/filter/constant_filter.hpp"
+#include "duckdb/planner/filter/dynamic_filter.hpp"
+#include "duckdb/planner/filter/optional_filter.hpp"
+#include "duckdb/execution/operator/join/join_filter_pushdown.hpp"
+#include "duckdb/optimizer/join_filter_pushdown_optimizer.hpp"
+#include "duckdb/planner/expression/bound_columnref_expression.hpp"
 
 namespace duckdb {
 
@@ -32,6 +39,69 @@ bool TopN::CanOptimize(LogicalOperator &op) {
 	return false;
 }
 
+void TopN::PushdownDynamicFilters(LogicalTopN &op) {
+	// pushdown dynamic filters through the Top-N operator
+	if (op.orders[0].null_order == OrderByNullType::NULLS_FIRST) {
+		// FIXME: not supported for NULLS FIRST quite yet
+		// we can support NULLS FIRST by doing (x IS NULL) OR [boundary value]
+		return;
+	}
+	auto &type = op.orders[0].expression->return_type;
+	if (!TypeIsIntegral(type.InternalType()) && type.id() != LogicalTypeId::VARCHAR) {
+		// only supported for integral types currently
+		return;
+	}
+	if (op.orders[0].expression->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) {
+		// we can only pushdown on ORDER BY [col] currently
+		return;
+	}
+	auto &colref = op.orders[0].expression->Cast<BoundColumnRefExpression>();
+	vector<JoinFilterPushdownColumn> columns;
+	JoinFilterPushdownColumn column;
+	column.probe_column_index = colref.binding;
+	columns.emplace_back(column);
+	vector<PushdownFilterTarget> pushdown_targets;
+	JoinFilterPushdownOptimizer::GetPushdownFilterTargets(*op.children[0], std::move(columns), pushdown_targets);
+	if (pushdown_targets.empty()) {
+		// no pushdown targets
+		return;
+	}
+	// found pushdown targets! generate dynamic filters
+	ExpressionType comparison_type;
+	if (op.orders[0].type == OrderType::ASCENDING) {
+		// for ascending order, we want the lowest N elements, so we filter on C <= [boundary]
+		// if we only have a single order clause, we can filter on C < boundary
+		comparison_type =
+		    op.orders.size() == 1 ? ExpressionType::COMPARE_LESSTHAN : ExpressionType::COMPARE_LESSTHANOREQUALTO;
+	} else {
+		// for descending order, we want the highest N elements, so we filter on C >= [boundary]
+		// if we only have a single order clause, we can filter on C > boundary
+		comparison_type =
+		    op.orders.size() == 1 ? ExpressionType::COMPARE_GREATERTHAN : ExpressionType::COMPARE_GREATERTHANOREQUALTO;
+	}
+	Value minimum_value = type.InternalType() == PhysicalType::VARCHAR ? Value("") : Value::MinimumValue(type);
+	auto base_filter = make_uniq<ConstantFilter>(comparison_type, std::move(minimum_value));
+	auto filter_data = make_shared_ptr<DynamicFilterData>();
+	filter_data->filter = std::move(base_filter);
+
+	// put the filter into the Top-N clause
+	op.dynamic_filter = filter_data;
+
+	for (auto &target : pushdown_targets) {
+		auto &get = target.get;
+		D_ASSERT(target.columns.size() == 1);
+		auto col_idx = target.columns[0].probe_column_index.column_index;
+
+		// create the actual dynamic filter
+		auto dynamic_filter = make_uniq<DynamicFilter>(filter_data);
+		auto optional_filter = make_uniq<OptionalFilter>(std::move(dynamic_filter));
+
+		// push the filter into the table scan
+		auto &column_index = get.GetColumnIds()[col_idx];
+		get.table_filters.PushFilter(column_index, std::move(optional_filter));
+	}
+}
+
 unique_ptr<LogicalOperator> TopN::Optimize(unique_ptr<LogicalOperator> op) {
 	if (CanOptimize(*op)) {
 
@@ -66,6 +136,7 @@ unique_ptr<LogicalOperator> TopN::Optimize(unique_ptr<LogicalOperator> op) {
 		if (topn->children[0]->has_estimated_cardinality && topn->children[0]->estimated_cardinality < limit_val) {
 			cardinality = topn->children[0]->estimated_cardinality;
 		}
+		PushdownDynamicFilters(*topn);
 		topn->SetEstimatedCardinality(cardinality);
 		op = std::move(topn);
 

package/src/duckdb/src/optimizer/unnest_rewriter.cpp

@@ -19,7 +19,7 @@ void UnnestRewriterPlanUpdater::VisitOperator(LogicalOperator &op) {
 void UnnestRewriterPlanUpdater::VisitExpression(unique_ptr<Expression> *expression) {
 	auto &expr = *expression;
 
-	if (expr->expression_class == ExpressionClass::BOUND_COLUMN_REF) {
+	if (expr->GetExpressionClass() == ExpressionClass::BOUND_COLUMN_REF) {
 		auto &bound_column_ref = expr->Cast<BoundColumnRefExpression>();
 		for (idx_t i = 0; i < replace_bindings.size(); i++) {
 			if (bound_column_ref.binding == replace_bindings[i].old_binding) {
@@ -122,7 +122,7 @@ bool UnnestRewriter::RewriteCandidate(unique_ptr<LogicalOperator> &candidate) {
 	D_ASSERT(delim_join.type == LogicalOperatorType::LOGICAL_DELIM_JOIN);
 	GetDelimColumns(delim_join);
 
-	// LHS of the LOGICAL_DELIM_JOIN is a LOGICAL_WINDOW that contains a LOGICAL_PROJECTION
+	// LHS of the LOGICAL_DELIM_JOIN is a LOGICAL_WINDOW that contains a LOGICAL_PROJECTION/LOGICAL_CROSS_JOIN
 	// this lhs_proj later becomes the child of the UNNEST
 
 	idx_t delim_idx = delim_join.delim_flipped ? 1 : 0;
@@ -279,7 +279,7 @@ void UnnestRewriter::UpdateBoundUnnestBindings(UnnestRewriterPlanUpdater &update
 
 			if (delim_binding.table_index == unnest_binding.table_index) {
 				unnest_binding.table_index = overwritten_tbl_idx;
-				unnest_binding.column_index++;
+				unnest_binding.column_index = i;
 				updater.replace_bindings.emplace_back(unnest_binding, delim_binding);
 				unnest_cols.erase(unnest_it);
 				break;
@@ -301,7 +301,7 @@ void UnnestRewriter::GetDelimColumns(LogicalOperator &op) {
 	auto &delim_join = op.Cast<LogicalComparisonJoin>();
 	for (idx_t i = 0; i < delim_join.duplicate_eliminated_columns.size(); i++) {
 		auto &expr = *delim_join.duplicate_eliminated_columns[i];
-		D_ASSERT(expr.type == ExpressionType::BOUND_COLUMN_REF);
+		D_ASSERT(expr.GetExpressionType() == ExpressionType::BOUND_COLUMN_REF);
 		auto &bound_colref_expr = expr.Cast<BoundColumnRefExpression>();
 		delim_columns.push_back(bound_colref_expr.binding);
 	}
@@ -326,7 +326,7 @@ void UnnestRewriter::GetLHSExpressions(LogicalOperator &op) {
 		lhs_bindings.emplace_back(col_bindings[i], op.types[i]);
 		if (set_alias) {
 			auto &proj = op.Cast<LogicalProjection>();
-			lhs_bindings.back().alias = proj.expressions[i]->alias;
+			lhs_bindings.back().alias = proj.expressions[i]->GetAlias();
 		}
 	}
 }

package/src/duckdb/src/parallel/event.cpp

@@ -63,6 +63,10 @@ void Event::FinishTask() {
 	}
 }
 
+ClientContext &Event::GetClientContext() {
+	return executor.context;
+}
+
 void Event::InsertEvent(shared_ptr<Event> replacement_event) {
 	replacement_event->parents = std::move(parents);
 #ifdef DEBUG

package/src/duckdb/src/parallel/executor.cpp

@@ -27,7 +27,7 @@ Executor::Executor(ClientContext &context) : context(context), executor_tasks(0)
 }
 
 Executor::~Executor() {
-	D_ASSERT(executor_tasks == 0);
+	D_ASSERT(Exception::UncaughtException() || executor_tasks == 0);
 }
 
 Executor &Executor::Get(ClientContext &context) {
@@ -684,39 +684,21 @@ void Executor::Flush(ThreadContext &thread_context) {
 	}
 }
 
-bool Executor::GetPipelinesProgress(double &current_progress, uint64_t &current_cardinality,
-                                    uint64_t &total_cardinality) { // LCOV_EXCL_START
+idx_t Executor::GetPipelinesProgress(ProgressData &progress) { // LCOV_EXCL_START
 	lock_guard<mutex> elock(executor_lock);
 
-	vector<double> progress;
-	vector<idx_t> cardinality;
-	total_cardinality = 0;
-	current_cardinality = 0;
+	progress.done = 0;
+	progress.total = 0;
+	idx_t count_invalid = 0;
 	for (auto &pipeline : pipelines) {
-		double child_percentage;
-		idx_t child_cardinality;
-
-		if (!pipeline->GetProgress(child_percentage, child_cardinality)) {
-			return false;
+		ProgressData p;
+		if (!pipeline->GetProgress(p)) {
+			count_invalid++;
+		} else {
+			progress.Add(p);
 		}
-		progress.push_back(child_percentage);
-		cardinality.push_back(child_cardinality);
-		total_cardinality += child_cardinality;
-	}
-	if (total_cardinality == 0) {
-		return true;
 	}
-	current_progress = 0;
-
-	for (size_t i = 0; i < progress.size(); i++) {
-		progress[i] = MaxValue(0.0, MinValue(100.0, progress[i]));
-		current_cardinality = LossyNumericCast<idx_t>(static_cast<double>(
-		    static_cast<double>(current_cardinality) +
-		    static_cast<double>(progress[i]) * static_cast<double>(cardinality[i]) / static_cast<double>(100)));
-		current_progress += progress[i] * double(cardinality[i]) / double(total_cardinality);
-		D_ASSERT(current_cardinality <= total_cardinality);
-	}
-	return true;
+	return count_invalid;
 } // LCOV_EXCL_STOP
 
 bool Executor::HasResultCollector() {

package/src/duckdb/src/parallel/executor_task.cpp

@@ -36,9 +36,14 @@ void ExecutorTask::Reschedule() {
 TaskExecutionResult ExecutorTask::Execute(TaskExecutionMode mode) {
 	try {
 		if (thread_context) {
-			thread_context->profiler.StartOperator(op);
-			auto result = ExecuteTask(mode);
-			thread_context->profiler.EndOperator(nullptr);
+			TaskExecutionResult result;
+			do {
+				thread_context->profiler.StartOperator(op);
+				// to allow continuous profiling, always execute in small steps
+				result = ExecuteTask(TaskExecutionMode::PROCESS_PARTIAL);
+				thread_context->profiler.EndOperator(nullptr);
+				executor.Flush(*thread_context);
+			} while (mode == TaskExecutionMode::PROCESS_ALL && result == TaskExecutionResult::TASK_NOT_FINISHED);
 			return result;
 		} else {
 			return ExecuteTask(mode);

package/src/duckdb/src/parallel/pipeline.cpp

@@ -72,17 +72,23 @@ ClientContext &Pipeline::GetClientContext() {
 	return executor.context;
 }
 
-bool Pipeline::GetProgress(double &current_percentage, idx_t &source_cardinality) {
+bool Pipeline::GetProgress(ProgressData &progress) {
 	D_ASSERT(source);
-	source_cardinality = MinValue<idx_t>(source->estimated_cardinality, 1ULL << 48ULL);
+	idx_t source_cardinality = MinValue<idx_t>(source->estimated_cardinality, 1ULL << 48ULL);
+	if (source_cardinality < 1) {
+		source_cardinality = 1;
+	}
 	if (!initialized) {
-		current_percentage = 0;
+		progress.done = 0;
+		progress.total = double(source_cardinality);
 		return true;
 	}
 	auto &client = executor.context;
-	current_percentage = source->GetProgress(client, *source_state);
-	current_percentage = sink->GetSinkProgress(client, *sink->sink_state, current_percentage);
-	return current_percentage >= 0;
+
+	progress = source->GetProgress(client, *source_state);
+	progress.Normalize(double(source_cardinality));
+	progress = sink->GetSinkProgress(client, *sink->sink_state, progress);
+	return progress.IsValid();
 }
 
 void Pipeline::ScheduleSequentialTask(shared_ptr<Event> &event) {
@@ -105,8 +111,9 @@ bool Pipeline::ScheduleParallel(shared_ptr<Event> &event) {
 			return false;
 		}
 	}
-	if (sink->RequiresBatchIndex()) {
-		if (!source->SupportsBatchIndex()) {
+	auto partition_info = sink->RequiredPartitionInfo();
+	if (partition_info.batch_index) {
+		if (!source->SupportsPartitioning(OperatorPartitionInfo::BatchIndex())) {
 			throw InternalException(
 			    "Attempting to schedule a pipeline where the sink requires batch index but source does not support it");
 		}