duckdb 0.8.2-dev33.0 → 0.8.2-dev3300.0

package/src/duckdb/src/execution/join_hashtable.cpp

@@ -902,16 +902,16 @@ bool JoinHashTable::RequiresPartitioning(ClientConfig &config, vector<unique_ptr
 		const auto partition_count = partition_counts[max_partition_idx];
 		const auto partition_size = partition_sizes[max_partition_idx];
 
-		const auto max_added_bits = 8 - radix_bits;
-		idx_t added_bits;
-		for (added_bits = 1; added_bits < max_added_bits; added_bits++) {
+		const auto max_added_bits = RadixPartitioning::MAX_RADIX_BITS - radix_bits;
+		idx_t added_bits = config.force_external ? 2 : 1;
+		for (; added_bits < max_added_bits; added_bits++) {
 			double partition_multiplier = RadixPartitioning::NumberOfPartitions(added_bits);
 
 			auto new_estimated_count = double(partition_count) / partition_multiplier;
 			auto new_estimated_size = double(partition_size) / partition_multiplier;
 			auto new_estimated_ht_size = new_estimated_size + PointerTableSize(new_estimated_count);
 
-			if (new_estimated_ht_size <= double(max_ht_size) / 4) {
+			if (config.force_external || new_estimated_ht_size <= double(max_ht_size) / 4) {
 				// Aim for an estimated partition size of max_ht_size / 4
 				break;
 			}

package/src/duckdb/src/execution/operator/aggregate/aggregate_object.cpp

@@ -1,4 +1,5 @@
 #include "duckdb/execution/operator/aggregate/aggregate_object.hpp"
+
 #include "duckdb/planner/expression/bound_aggregate_expression.hpp"
 #include "duckdb/planner/expression/bound_window_expression.hpp"
 

package/src/duckdb/src/execution/operator/aggregate/physical_hash_aggregate.cpp

@@ -396,10 +396,10 @@ SinkResultType PhysicalHashAggregate::Sink(ExecutionContext &context, DataChunk
 	return SinkResultType::NEED_MORE_INPUT;
 }
 
-void PhysicalHashAggregate::CombineDistinct(ExecutionContext &context, GlobalSinkState &state,
-                                            LocalSinkState &lstate) const {
-	auto &global_sink = state.Cast<HashAggregateGlobalState>();
-	auto &sink = lstate.Cast<HashAggregateLocalState>();
+void PhysicalHashAggregate::CombineDistinct(ExecutionContext &context, OperatorSinkCombineInput &input) const {
+
+	auto &global_sink = input.global_state.Cast<HashAggregateGlobalState>();
+	auto &sink = input.local_state.Cast<HashAggregateLocalState>();
 
 	if (!distinct_collection_info) {
 		return;
@@ -425,14 +425,15 @@ void PhysicalHashAggregate::CombineDistinct(ExecutionContext &context, GlobalSin
 	}
 }
 
-void PhysicalHashAggregate::Combine(ExecutionContext &context, GlobalSinkState &state, LocalSinkState &lstate) const {
-	auto &gstate = state.Cast<HashAggregateGlobalState>();
-	auto &llstate = lstate.Cast<HashAggregateLocalState>();
+SinkCombineResultType PhysicalHashAggregate::Combine(ExecutionContext &context, OperatorSinkCombineInput &input) const {
+	auto &gstate = input.global_state.Cast<HashAggregateGlobalState>();
+	auto &llstate = input.local_state.Cast<HashAggregateLocalState>();
 
-	CombineDistinct(context, state, lstate);
+	OperatorSinkCombineInput combine_distinct_input {gstate, llstate, input.interrupt_state};
+	CombineDistinct(context, combine_distinct_input);
 
 	if (CanSkipRegularSink()) {
-		return;
+		return SinkCombineResultType::FINISHED;
 	}
 	for (idx_t i = 0; i < groupings.size(); i++) {
 		auto &grouping_gstate = gstate.grouping_states[i];
@@ -442,6 +443,8 @@ void PhysicalHashAggregate::Combine(ExecutionContext &context, GlobalSinkState &
 		auto &table = grouping.table_data;
 		table.Combine(context, *grouping_gstate.table_state, *grouping_lstate.table_state);
 	}
+
+	return SinkCombineResultType::FINISHED;
 }
 
 //! REGULAR FINALIZE EVENT
@@ -818,8 +821,8 @@ SinkFinalizeType PhysicalHashAggregate::FinalizeInternal(Pipeline &pipeline, Eve
 }
 
 SinkFinalizeType PhysicalHashAggregate::Finalize(Pipeline &pipeline, Event &event, ClientContext &context,
-                                                 GlobalSinkState &gstate_p) const {
-	return FinalizeInternal(pipeline, event, context, gstate_p, true);
+                                                 OperatorSinkFinalizeInput &input) const {
+	return FinalizeInternal(pipeline, event, context, input.global_state, true);
 }
 
 //===--------------------------------------------------------------------===//

package/src/duckdb/src/execution/operator/aggregate/physical_perfecthash_aggregate.cpp

@@ -159,13 +159,15 @@ SinkResultType PhysicalPerfectHashAggregate::Sink(ExecutionContext &context, Dat
 //===--------------------------------------------------------------------===//
 // Combine
 //===--------------------------------------------------------------------===//
-void PhysicalPerfectHashAggregate::Combine(ExecutionContext &context, GlobalSinkState &gstate_p,
-                                           LocalSinkState &lstate_p) const {
-	auto &lstate = lstate_p.Cast<PerfectHashAggregateLocalState>();
-	auto &gstate = gstate_p.Cast<PerfectHashAggregateGlobalState>();
+SinkCombineResultType PhysicalPerfectHashAggregate::Combine(ExecutionContext &context,
+                                                            OperatorSinkCombineInput &input) const {
+	auto &lstate = input.local_state.Cast<PerfectHashAggregateLocalState>();
+	auto &gstate = input.global_state.Cast<PerfectHashAggregateGlobalState>();
 
 	lock_guard<mutex> l(gstate.lock);
 	gstate.ht->Combine(*lstate.ht);
+
+	return SinkCombineResultType::FINISHED;
 }
 
 //===--------------------------------------------------------------------===//

package/src/duckdb/src/execution/operator/aggregate/physical_streaming_window.cpp

@@ -26,14 +26,16 @@ class StreamingWindowState : public OperatorState {
 public:
 	using StateBuffer = vector<data_t>;
 
-	StreamingWindowState() : initialized(false), statev(LogicalType::POINTER, data_ptr_cast(&state_ptr)) {
+	StreamingWindowState()
+	    : initialized(false), allocator(Allocator::DefaultAllocator()),
+	      statev(LogicalType::POINTER, data_ptr_cast(&state_ptr)) {
 	}
 
 	~StreamingWindowState() override {
 		for (size_t i = 0; i < aggregate_dtors.size(); ++i) {
 			auto dtor = aggregate_dtors[i];
 			if (dtor) {
-				AggregateInputData aggr_input_data(aggregate_bind_data[i], Allocator::DefaultAllocator());
+				AggregateInputData aggr_input_data(aggregate_bind_data[i], allocator);
 				state_ptr = aggregate_states[i].data();
 				dtor(statev, aggr_input_data, 1);
 			}
@@ -89,6 +91,7 @@ public:
 public:
 	bool initialized;
 	vector<unique_ptr<Vector>> const_vectors;
+	ArenaAllocator allocator;
 
 	// Aggregation
 	vector<StateBuffer> aggregate_states;
@@ -110,6 +113,8 @@ OperatorResultType PhysicalStreamingWindow::Execute(ExecutionContext &context, D
                                                     GlobalOperatorState &gstate_p, OperatorState &state_p) const {
 	auto &gstate = gstate_p.Cast<StreamingWindowGlobalState>();
 	auto &state = state_p.Cast<StreamingWindowState>();
+	state.allocator.Reset();
+
 	if (!state.initialized) {
 		state.Initialize(context.client, input, select_list);
 	}
@@ -130,7 +135,7 @@ OperatorResultType PhysicalStreamingWindow::Execute(ExecutionContext &context, D
 			auto &aggregate = *wexpr.aggregate;
 			auto &statev = state.statev;
 			state.state_ptr = state.aggregate_states[expr_idx].data();
-			AggregateInputData aggr_input_data(wexpr.bind_info.get(), Allocator::DefaultAllocator());
+			AggregateInputData aggr_input_data(wexpr.bind_info.get(), state.allocator);
 
 			// Check for COUNT(*)
 			if (wexpr.children.empty()) {

package/src/duckdb/src/execution/operator/aggregate/physical_ungrouped_aggregate.cpp

@@ -1,20 +1,21 @@
 #include "duckdb/execution/operator/aggregate/physical_ungrouped_aggregate.hpp"
 
 #include "duckdb/catalog/catalog_entry/aggregate_function_catalog_entry.hpp"
+#include "duckdb/common/algorithm.hpp"
+#include "duckdb/common/unordered_set.hpp"
 #include "duckdb/common/vector_operations/vector_operations.hpp"
 #include "duckdb/execution/expression_executor.hpp"
 #include "duckdb/execution/operator/aggregate/aggregate_object.hpp"
+#include "duckdb/execution/operator/aggregate/distinct_aggregate_data.hpp"
+#include "duckdb/execution/radix_partitioned_hashtable.hpp"
 #include "duckdb/main/client_context.hpp"
+#include "duckdb/parallel/base_pipeline_event.hpp"
+#include "duckdb/parallel/interrupt.hpp"
 #include "duckdb/parallel/thread_context.hpp"
 #include "duckdb/planner/expression/bound_aggregate_expression.hpp"
 #include "duckdb/planner/expression/bound_reference_expression.hpp"
-#include "duckdb/execution/radix_partitioned_hashtable.hpp"
-#include "duckdb/parallel/base_pipeline_event.hpp"
-#include "duckdb/common/unordered_set.hpp"
-#include "duckdb/common/algorithm.hpp"
-#include "duckdb/parallel/interrupt.hpp"
+
 #include <functional>
-#include "duckdb/execution/operator/aggregate/distinct_aggregate_data.hpp"
 
 namespace duckdb {
 
@@ -58,7 +59,8 @@ struct AggregateState {
 			Vector state_vector(Value::POINTER(CastPointerToValue(aggregates[i].get())));
 			state_vector.SetVectorType(VectorType::FLAT_VECTOR);
 
-			AggregateInputData aggr_input_data(bind_data[i], Allocator::DefaultAllocator());
+			ArenaAllocator allocator(Allocator::DefaultAllocator());
+			AggregateInputData aggr_input_data(bind_data[i], allocator);
 			destructors[i](state_vector, aggr_input_data, 1);
 		}
 	}
@@ -81,7 +83,7 @@ struct AggregateState {
 class UngroupedAggregateGlobalState : public GlobalSinkState {
 public:
 	UngroupedAggregateGlobalState(const PhysicalUngroupedAggregate &op, ClientContext &client)
-	    : state(op.aggregates), finished(false) {
+	    : state(op.aggregates), finished(false), allocator(BufferAllocator::Get(client)) {
 		if (op.distinct_data) {
 			distinct_state = make_uniq<DistinctAggregateState>(*op.distinct_data, client);
 		}
@@ -95,16 +97,19 @@ public:
 	bool finished;
 	//! The data related to the distinct aggregates (if there are any)
 	unique_ptr<DistinctAggregateState> distinct_state;
+	//! Global arena allocator
+	ArenaAllocator allocator;
 };
 
 class UngroupedAggregateLocalState : public LocalSinkState {
 public:
 	UngroupedAggregateLocalState(const PhysicalUngroupedAggregate &op, const vector<LogicalType> &child_types,
 	                             GlobalSinkState &gstate_p, ExecutionContext &context)
-	    : state(op.aggregates), child_executor(context.client), aggregate_input_chunk(), filter_set() {
+	    : allocator(BufferAllocator::Get(context.client)), state(op.aggregates), child_executor(context.client),
+	      aggregate_input_chunk(), filter_set() {
 		auto &gstate = gstate_p.Cast<UngroupedAggregateGlobalState>();
 
-		auto &allocator = Allocator::Get(context.client);
+		auto &allocator = BufferAllocator::Get(context.client);
 		InitializeDistinctAggregates(op, gstate, context);
 
 		vector<LogicalType> payload_types;
@@ -125,6 +130,8 @@ public:
 		filter_set.Initialize(context.client, aggregate_objects, child_types);
 	}
 
+	//! Local arena allocator
+	ArenaAllocator allocator;
 	//! The local aggregate state
 	AggregateState state;
 	//! The executor
@@ -281,7 +288,7 @@ SinkResultType PhysicalUngroupedAggregate::Sink(ExecutionContext &context, DataC
 		}
 
 		auto start_of_input = payload_cnt == 0 ? nullptr : &payload_chunk.data[payload_idx];
-		AggregateInputData aggr_input_data(aggregate.bind_info.get(), Allocator::DefaultAllocator());
+		AggregateInputData aggr_input_data(aggregate.bind_info.get(), sink.allocator);
 		aggregate.function.simple_update(start_of_input, aggr_input_data, payload_cnt,
 		                                 sink.state.aggregates[aggr_idx].get(), payload_chunk.size());
 	}
@@ -292,10 +299,9 @@ SinkResultType PhysicalUngroupedAggregate::Sink(ExecutionContext &context, DataC
 // Finalize
 //===--------------------------------------------------------------------===//
 
-void PhysicalUngroupedAggregate::CombineDistinct(ExecutionContext &context, GlobalSinkState &state,
-                                                 LocalSinkState &lstate) const {
-	auto &global_sink = state.Cast<UngroupedAggregateGlobalState>();
-	auto &source = lstate.Cast<UngroupedAggregateLocalState>();
+void PhysicalUngroupedAggregate::CombineDistinct(ExecutionContext &context, OperatorSinkCombineInput &input) const {
+	auto &global_sink = input.global_state.Cast<UngroupedAggregateGlobalState>();
+	auto &source = input.local_state.Cast<UngroupedAggregateLocalState>();
 
 	if (!distinct_data) {
 		return;
@@ -312,10 +318,10 @@ void PhysicalUngroupedAggregate::CombineDistinct(ExecutionContext &context, Glob
 	}
 }
 
-void PhysicalUngroupedAggregate::Combine(ExecutionContext &context, GlobalSinkState &state,
-                                         LocalSinkState &lstate) const {
-	auto &gstate = state.Cast<UngroupedAggregateGlobalState>();
-	auto &source = lstate.Cast<UngroupedAggregateLocalState>();
+SinkCombineResultType PhysicalUngroupedAggregate::Combine(ExecutionContext &context,
+                                                          OperatorSinkCombineInput &input) const {
+	auto &gstate = input.global_state.Cast<UngroupedAggregateGlobalState>();
+	auto &source = input.local_state.Cast<UngroupedAggregateLocalState>();
 	D_ASSERT(!gstate.finished);
 
 	// finalize: combine the local state into the global state
@@ -323,7 +329,8 @@ void PhysicalUngroupedAggregate::Combine(ExecutionContext &context, GlobalSinkSt
 	// use the combine method to combine the partial aggregates
 	lock_guard<mutex> glock(gstate.lock);
 
-	CombineDistinct(context, state, lstate);
+	OperatorSinkCombineInput distinct_input {gstate, source, input.interrupt_state};
+	CombineDistinct(context, distinct_input);
 
 	for (idx_t aggr_idx = 0; aggr_idx < aggregates.size(); aggr_idx++) {
 		auto &aggregate = aggregates[aggr_idx]->Cast<BoundAggregateExpression>();
@@ -335,16 +342,19 @@ void PhysicalUngroupedAggregate::Combine(ExecutionContext &context, GlobalSinkSt
 		Vector source_state(Value::POINTER(CastPointerToValue(source.state.aggregates[aggr_idx].get())));
 		Vector dest_state(Value::POINTER(CastPointerToValue(gstate.state.aggregates[aggr_idx].get())));
 
-		AggregateInputData aggr_input_data(aggregate.bind_info.get(), Allocator::DefaultAllocator());
+		AggregateInputData aggr_input_data(aggregate.bind_info.get(), gstate.allocator);
 		aggregate.function.combine(source_state, dest_state, aggr_input_data, 1);
 #ifdef DEBUG
 		gstate.state.counts[aggr_idx] += source.state.counts[aggr_idx];
 #endif
 	}
+	source.allocator.Destroy();
 
 	auto &client_profiler = QueryProfiler::Get(context.client);
 	context.thread.profiler.Flush(*this, source.child_executor, "child_executor", 0);
 	client_profiler.Flush(context.thread.profiler);
+
+	return SinkCombineResultType::FINISHED;
 }
 
 class UngroupedDistinctAggregateFinalizeTask : public ExecutorTask {
@@ -352,7 +362,8 @@ public:
 	UngroupedDistinctAggregateFinalizeTask(Executor &executor, shared_ptr<Event> event_p,
 	                                       UngroupedAggregateGlobalState &state_p, ClientContext &context,
 	                                       const PhysicalUngroupedAggregate &op)
-	    : ExecutorTask(executor), event(std::move(event_p)), gstate(state_p), context(context), op(op) {
+	    : ExecutorTask(executor), event(std::move(event_p)), gstate(state_p), context(context), op(op),
+	      allocator(BufferAllocator::Get(context)) {
 	}
 
 	void AggregateDistinct() {
@@ -367,22 +378,22 @@ public:
 		idx_t payload_idx = 0;
 		idx_t next_payload_idx = 0;
 
-		for (idx_t i = 0; i < aggregates.size(); i++) {
-			auto &aggregate = aggregates[i]->Cast<BoundAggregateExpression>();
+		for (idx_t agg_idx = 0; agg_idx < aggregates.size(); agg_idx++) {
+			auto &aggregate = aggregates[agg_idx]->Cast<BoundAggregateExpression>();
 
 			// Forward the payload idx
 			payload_idx = next_payload_idx;
 			next_payload_idx = payload_idx + aggregate.children.size();
 
 			// If aggregate is not distinct, skip it
-			if (!distinct_data.IsDistinct(i)) {
+			if (!distinct_data.IsDistinct(agg_idx)) {
 				continue;
 			}
 
 			DataChunk payload_chunk;
 
-			D_ASSERT(distinct_data.info.table_map.count(i));
-			auto table_idx = distinct_data.info.table_map.at(i);
+			D_ASSERT(distinct_data.info.table_map.count(agg_idx));
+			auto table_idx = distinct_data.info.table_map.at(agg_idx);
 			auto &radix_table_p = distinct_data.radix_tables[table_idx];
 			auto &output_chunk = *distinct_state.distinct_output_chunks[table_idx];
 			auto &grouped_aggregate_data = *distinct_data.grouped_aggregate_data[table_idx];
@@ -417,14 +428,14 @@ public:
 				}
 				payload_chunk.SetCardinality(output_chunk);
 #ifdef DEBUG
-				gstate.state.counts[i] += payload_chunk.size();
+				gstate.state.counts[agg_idx] += payload_chunk.size();
 #endif
 
 				auto start_of_input = payload_cnt ? &payload_chunk.data[0] : nullptr;
 				//! Update the aggregate state
-				AggregateInputData aggr_input_data(aggregate.bind_info.get(), Allocator::DefaultAllocator());
+				AggregateInputData aggr_input_data(aggregate.bind_info.get(), allocator);
 				aggregate.function.simple_update(start_of_input, aggr_input_data, payload_cnt,
-				                                 gstate.state.aggregates[i].get(), payload_chunk.size());
+				                                 gstate.state.aggregates[agg_idx].get(), payload_chunk.size());
 			}
 		}
 		D_ASSERT(!gstate.finished);
@@ -442,6 +453,7 @@ private:
 	UngroupedAggregateGlobalState &gstate;
 	ClientContext &context;
 	const PhysicalUngroupedAggregate &op;
+	ArenaAllocator allocator;
 };
 
 // TODO: Create tasks and run these in parallel instead of doing this all in Schedule, single threaded
@@ -526,11 +538,11 @@ SinkFinalizeType PhysicalUngroupedAggregate::FinalizeDistinct(Pipeline &pipeline
 }
 
 SinkFinalizeType PhysicalUngroupedAggregate::Finalize(Pipeline &pipeline, Event &event, ClientContext &context,
-                                                      GlobalSinkState &gstate_p) const {
-	auto &gstate = gstate_p.Cast<UngroupedAggregateGlobalState>();
+                                                      OperatorSinkFinalizeInput &input) const {
+	auto &gstate = input.global_state.Cast<UngroupedAggregateGlobalState>();
 
 	if (distinct_data) {
-		return FinalizeDistinct(pipeline, event, context, gstate_p);
+		return FinalizeDistinct(pipeline, event, context, input.global_state);
 	}
 
 	D_ASSERT(!gstate.finished);
@@ -566,7 +578,7 @@ SourceResultType PhysicalUngroupedAggregate::GetData(ExecutionContext &context,
 		auto &aggregate = aggregates[aggr_idx]->Cast<BoundAggregateExpression>();
 
 		Vector state_vector(Value::POINTER(CastPointerToValue(gstate.state.aggregates[aggr_idx].get())));
-		AggregateInputData aggr_input_data(aggregate.bind_info.get(), Allocator::DefaultAllocator());
+		AggregateInputData aggr_input_data(aggregate.bind_info.get(), gstate.allocator);
 		aggregate.function.finalize(state_vector, aggr_input_data, chunk.data[aggr_idx], 1, 0);
 	}
 	VerifyNullHandling(chunk, gstate.state, aggregates);