duckdb 0.8.2-dev1764.0 → 0.8.2-dev1859.0

package/src/duckdb/src/include/duckdb/execution/window_executor.hpp

@@ -0,0 +1,313 @@
+//===----------------------------------------------------------------------===//
+//                         DuckDB
+//
+// duckdb/execution/window_executor.hpp
+//
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "duckdb/execution/expression_executor.hpp"
+#include "duckdb/execution/window_segment_tree.hpp"
+#include "duckdb/planner/expression/bound_window_expression.hpp"
+#include "duckdb/common/vector_operations/vector_operations.hpp"
+
+namespace duckdb {
+
+struct WindowInputExpression {
+	static void PrepareInputExpression(Expression &expr, ExpressionExecutor &executor, DataChunk &chunk) {
+		vector<LogicalType> types;
+		types.push_back(expr.return_type);
+		executor.AddExpression(expr);
+
+		auto &allocator = executor.GetAllocator();
+		chunk.Initialize(allocator, types);
+	}
+
+	WindowInputExpression(optional_ptr<Expression> expr_p, ClientContext &context)
+	    : expr(expr_p), ptype(PhysicalType::INVALID), scalar(true), executor(context) {
+		if (expr) {
+			PrepareInputExpression(*expr, executor, chunk);
+			ptype = expr->return_type.InternalType();
+			scalar = expr->IsScalar();
+		}
+	}
+
+	void Execute(DataChunk &input_chunk) {
+		if (expr) {
+			chunk.Reset();
+			executor.Execute(input_chunk, chunk);
+			chunk.Verify();
+		}
+	}
+
+	template <typename T>
+	inline T GetCell(idx_t i) const {
+		D_ASSERT(!chunk.data.empty());
+		const auto data = FlatVector::GetData<T>(chunk.data[0]);
+		return data[scalar ? 0 : i];
+	}
+
+	inline bool CellIsNull(idx_t i) const {
+		D_ASSERT(!chunk.data.empty());
+		if (chunk.data[0].GetVectorType() == VectorType::CONSTANT_VECTOR) {
+			return ConstantVector::IsNull(chunk.data[0]);
+		}
+		return FlatVector::IsNull(chunk.data[0], i);
+	}
+
+	inline void CopyCell(Vector &target, idx_t target_offset) const {
+		D_ASSERT(!chunk.data.empty());
+		auto &source = chunk.data[0];
+		auto source_offset = scalar ? 0 : target_offset;
+		VectorOperations::Copy(source, target, source_offset + 1, source_offset, target_offset);
+	}
+
+	optional_ptr<Expression> expr;
+	PhysicalType ptype;
+	bool scalar;
+	ExpressionExecutor executor;
+	DataChunk chunk;
+};
+
+struct WindowInputColumn {
+	WindowInputColumn(Expression *expr_p, ClientContext &context, idx_t capacity_p)
+	    : input_expr(expr_p, context), count(0), capacity(capacity_p) {
+		if (input_expr.expr) {
+			target = make_uniq<Vector>(input_expr.chunk.data[0].GetType(), capacity);
+		}
+	}
+
+	void Append(DataChunk &input_chunk) {
+		if (input_expr.expr) {
+			const auto source_count = input_chunk.size();
+			D_ASSERT(count + source_count <= capacity);
+			if (!input_expr.scalar || !count) {
+				input_expr.Execute(input_chunk);
+				auto &source = input_expr.chunk.data[0];
+				VectorOperations::Copy(source, *target, source_count, 0, count);
+			}
+			count += source_count;
+		}
+	}
+
+	inline bool CellIsNull(idx_t i) const {
+		D_ASSERT(target);
+		D_ASSERT(i < count);
+		return FlatVector::IsNull(*target, input_expr.scalar ? 0 : i);
+	}
+
+	template <typename T>
+	inline T GetCell(idx_t i) const {
+		D_ASSERT(target);
+		D_ASSERT(i < count);
+		const auto data = FlatVector::GetData<T>(*target);
+		return data[input_expr.scalar ? 0 : i];
+	}
+
+	WindowInputExpression input_expr;
+
+private:
+	unique_ptr<Vector> target;
+	idx_t count;
+	idx_t capacity;
+};
+
+//	Column indexes of the bounds chunk
+enum WindowBounds : uint8_t { PARTITION_BEGIN, PARTITION_END, PEER_BEGIN, PEER_END, WINDOW_BEGIN, WINDOW_END };
+
+class WindowExecutorState {
+public:
+	WindowExecutorState() {};
+	virtual ~WindowExecutorState() {
+	}
+
+	template <class TARGET>
+	TARGET &Cast() {
+		D_ASSERT(dynamic_cast<TARGET *>(this));
+		return reinterpret_cast<TARGET &>(*this);
+	}
+	template <class TARGET>
+	const TARGET &Cast() const {
+		D_ASSERT(dynamic_cast<const TARGET *>(this));
+		return reinterpret_cast<const TARGET &>(*this);
+	}
+};
+
+class WindowExecutor {
+public:
+	WindowExecutor(BoundWindowExpression &wexpr, ClientContext &context, const idx_t payload_count,
+	               const ValidityMask &partition_mask, const ValidityMask &order_mask);
+	virtual ~WindowExecutor() {
+	}
+
+	virtual void Sink(DataChunk &input_chunk, const idx_t input_idx, const idx_t total_count) {
+		range.Append(input_chunk);
+	}
+
+	virtual void Finalize() {
+	}
+
+	virtual unique_ptr<WindowExecutorState> GetExecutorState() const;
+
+	void Evaluate(idx_t row_idx, DataChunk &input_chunk, Vector &result, WindowExecutorState &lstate) const;
+
+protected:
+	// The function
+	BoundWindowExpression &wexpr;
+	ClientContext &context;
+	const idx_t payload_count;
+	const ValidityMask &partition_mask;
+	const ValidityMask &order_mask;
+
+	// Expression collections
+	DataChunk payload_collection;
+	ExpressionExecutor payload_executor;
+	DataChunk payload_chunk;
+
+	// evaluate RANGE expressions, if needed
+	WindowInputColumn range;
+
+	virtual void EvaluateInternal(WindowExecutorState &lstate, Vector &result, idx_t count, idx_t row_idx) const = 0;
+};
+
+class WindowAggregateExecutor : public WindowExecutor {
+public:
+	bool IsConstantAggregate();
+	bool IsCustomAggregate();
+
+	WindowAggregateExecutor(BoundWindowExpression &wexpr, ClientContext &context, const idx_t payload_count,
+	                        const ValidityMask &partition_mask, const ValidityMask &order_mask,
+	                        WindowAggregationMode mode);
+
+	void Sink(DataChunk &input_chunk, const idx_t input_idx, const idx_t total_count) override;
+	void Finalize() override;
+
+	unique_ptr<WindowExecutorState> GetExecutorState() const override;
+
+	const WindowAggregationMode mode;
+
+protected:
+	ExpressionExecutor filter_executor;
+	SelectionVector filter_sel;
+
+	// aggregate computation algorithm
+	unique_ptr<WindowAggregator> aggregator;
+
+	void EvaluateInternal(WindowExecutorState &lstate, Vector &result, idx_t count, idx_t row_idx) const override;
+};
+
+class WindowRowNumberExecutor : public WindowExecutor {
+public:
+	WindowRowNumberExecutor(BoundWindowExpression &wexpr, ClientContext &context, const idx_t payload_count,
+	                        const ValidityMask &partition_mask, const ValidityMask &order_mask);
+
+protected:
+	void EvaluateInternal(WindowExecutorState &lstate, Vector &result, idx_t count, idx_t row_idx) const override;
+};
+
+//	Base class for non-aggregate functions that use peer boundaries
+class WindowRankExecutor : public WindowExecutor {
+public:
+	WindowRankExecutor(BoundWindowExpression &wexpr, ClientContext &context, const idx_t payload_count,
+	                   const ValidityMask &partition_mask, const ValidityMask &order_mask);
+
+	unique_ptr<WindowExecutorState> GetExecutorState() const override;
+
+protected:
+	void EvaluateInternal(WindowExecutorState &lstate, Vector &result, idx_t count, idx_t row_idx) const override;
+};
+
+class WindowDenseRankExecutor : public WindowExecutor {
+public:
+	WindowDenseRankExecutor(BoundWindowExpression &wexpr, ClientContext &context, const idx_t payload_count,
+	                        const ValidityMask &partition_mask, const ValidityMask &order_mask);
+
+	unique_ptr<WindowExecutorState> GetExecutorState() const override;
+
+protected:
+	void EvaluateInternal(WindowExecutorState &lstate, Vector &result, idx_t count, idx_t row_idx) const override;
+};
+
+class WindowPercentRankExecutor : public WindowExecutor {
+public:
+	WindowPercentRankExecutor(BoundWindowExpression &wexpr, ClientContext &context, const idx_t payload_count,
+	                          const ValidityMask &partition_mask, const ValidityMask &order_mask);
+
+	unique_ptr<WindowExecutorState> GetExecutorState() const override;
+
+protected:
+	void EvaluateInternal(WindowExecutorState &lstate, Vector &result, idx_t count, idx_t row_idx) const override;
+};
+
+class WindowCumeDistExecutor : public WindowExecutor {
+public:
+	WindowCumeDistExecutor(BoundWindowExpression &wexpr, ClientContext &context, const idx_t payload_count,
+	                       const ValidityMask &partition_mask, const ValidityMask &order_mask);
+
+protected:
+	void EvaluateInternal(WindowExecutorState &lstate, Vector &result, idx_t count, idx_t row_idx) const override;
+};
+
+// Base class for non-aggregate functions that have a payload
+class WindowValueExecutor : public WindowExecutor {
+public:
+	WindowValueExecutor(BoundWindowExpression &wexpr, ClientContext &context, const idx_t payload_count,
+	                    const ValidityMask &partition_mask, const ValidityMask &order_mask);
+
+	void Sink(DataChunk &input_chunk, const idx_t input_idx, const idx_t total_count) override;
+
+protected:
+	// IGNORE NULLS
+	ValidityMask ignore_nulls;
+};
+
+//
+class WindowNtileExecutor : public WindowValueExecutor {
+public:
+	WindowNtileExecutor(BoundWindowExpression &wexpr, ClientContext &context, const idx_t payload_count,
+	                    const ValidityMask &partition_mask, const ValidityMask &order_mask);
+
+protected:
+	void EvaluateInternal(WindowExecutorState &lstate, Vector &result, idx_t count, idx_t row_idx) const override;
+};
+class WindowLeadLagExecutor : public WindowValueExecutor {
+public:
+	WindowLeadLagExecutor(BoundWindowExpression &wexpr, ClientContext &context, const idx_t payload_count,
+	                      const ValidityMask &partition_mask, const ValidityMask &order_mask);
+
+	unique_ptr<WindowExecutorState> GetExecutorState() const override;
+
+protected:
+	void EvaluateInternal(WindowExecutorState &lstate, Vector &result, idx_t count, idx_t row_idx) const override;
+};
+
+class WindowFirstValueExecutor : public WindowValueExecutor {
+public:
+	WindowFirstValueExecutor(BoundWindowExpression &wexpr, ClientContext &context, const idx_t payload_count,
+	                         const ValidityMask &partition_mask, const ValidityMask &order_mask);
+
+protected:
+	void EvaluateInternal(WindowExecutorState &lstate, Vector &result, idx_t count, idx_t row_idx) const override;
+};
+
+class WindowLastValueExecutor : public WindowValueExecutor {
+public:
+	WindowLastValueExecutor(BoundWindowExpression &wexpr, ClientContext &context, const idx_t payload_count,
+	                        const ValidityMask &partition_mask, const ValidityMask &order_mask);
+
+protected:
+	void EvaluateInternal(WindowExecutorState &lstate, Vector &result, idx_t count, idx_t row_idx) const override;
+};
+
+class WindowNthValueExecutor : public WindowValueExecutor {
+public:
+	WindowNthValueExecutor(BoundWindowExpression &wexpr, ClientContext &context, const idx_t payload_count,
+	                       const ValidityMask &partition_mask, const ValidityMask &order_mask);
+
+protected:
+	void EvaluateInternal(WindowExecutorState &lstate, Vector &result, idx_t count, idx_t row_idx) const override;
+};
+
+} // namespace duckdb

package/src/duckdb/src/include/duckdb/execution/window_segment_tree.hpp

@@ -16,20 +16,42 @@
 
 namespace duckdb {
 
-class WindowAggregateState {
+class WindowAggregatorState {
 public:
-	WindowAggregateState(AggregateObject aggr, const LogicalType &result_type_p, idx_t partition_count);
-	virtual ~WindowAggregateState();
+	WindowAggregatorState();
+	virtual ~WindowAggregatorState() {
+	}
+
+	template <class TARGET>
+	TARGET &Cast() {
+		D_ASSERT(dynamic_cast<TARGET *>(this));
+		return reinterpret_cast<TARGET &>(*this);
+	}
+	template <class TARGET>
+	const TARGET &Cast() const {
+		D_ASSERT(dynamic_cast<const TARGET *>(this));
+		return reinterpret_cast<const TARGET &>(*this);
+	}
 
+	//! Allocator for aggregates
+	ArenaAllocator allocator;
+};
+
+class WindowAggregator {
+public:
+	WindowAggregator(AggregateObject aggr, const LogicalType &result_type_p, idx_t partition_count);
+	virtual ~WindowAggregator();
+
+	//	Build
 	virtual void Sink(DataChunk &payload_chunk, SelectionVector *filter_sel, idx_t filtered);
 	virtual void Finalize();
-	virtual void Compute(Vector &result, idx_t rid, idx_t start, idx_t end);
-	virtual void Evaluate(const idx_t *begins, const idx_t *ends, Vector &result, idx_t count);
 
-protected:
-	void AggregateInit();
-	void AggegateFinal(Vector &result, idx_t rid);
+	//	Probe
+	virtual unique_ptr<WindowAggregatorState> GetLocalState() const = 0;
+	virtual void Evaluate(WindowAggregatorState &lstate, const idx_t *begins, const idx_t *ends, Vector &result,
+	                      idx_t count) const = 0;
 
+protected:
 	AggregateObject aggr;
 	//! The result type of the window function
 	LogicalType result_type;
@@ -38,10 +60,6 @@ protected:
 	const idx_t partition_count;
 	//! The size of a single aggregate state
 	const idx_t state_size;
-	//! Data pointer that contains a single state, used for intermediate window segment aggregation
-	vector<data_t> state;
-	//! Reused result state container for the window functions
-	Vector statef;
 	//! Partition data chunk
 	DataChunk inputs;
 
@@ -49,23 +67,28 @@ protected:
 	vector<validity_t> filter_bits;
 	ValidityMask filter_mask;
 	idx_t filter_pos;
-
-	//! Allocator for aggregates
-	ArenaAllocator allocator;
+	//! The state used by the aggregator to build.
+	unique_ptr<WindowAggregatorState> gstate;
 };
 
-class WindowConstantAggregate : public WindowAggregateState {
+class WindowConstantAggregator : public WindowAggregator {
 public:
-	WindowConstantAggregate(AggregateObject aggr, const LogicalType &result_type_p, const ValidityMask &partition_mask,
-	                        const idx_t count);
-	~WindowConstantAggregate() override {
+	WindowConstantAggregator(AggregateObject aggr, const LogicalType &result_type_p, const ValidityMask &partition_mask,
+	                         const idx_t count);
+	~WindowConstantAggregator() override {
 	}
 
 	void Sink(DataChunk &payload_chunk, SelectionVector *filter_sel, idx_t filtered) override;
 	void Finalize() override;
-	void Evaluate(const idx_t *begins, const idx_t *ends, Vector &result, idx_t count) override;
+
+	unique_ptr<WindowAggregatorState> GetLocalState() const override;
+	void Evaluate(WindowAggregatorState &lstate, const idx_t *begins, const idx_t *ends, Vector &result,
+	              idx_t count) const override;
 
 private:
+	void AggregateInit();
+	void AggegateFinal(Vector &result, idx_t rid);
+
 	//! Partition starts
 	vector<idx_t> partition_offsets;
 	//! Aggregate results
@@ -74,25 +97,25 @@ private:
 	idx_t partition;
 	//! The current input row being built/read
 	idx_t row;
+	//! Data pointer that contains a single state, used for intermediate window segment aggregation
+	vector<data_t> state;
 	//! A vector of pointers to "state", used for intermediate window segment aggregation
 	Vector statep;
-	//! Shared SV for evaluation
-	SelectionVector matches;
+	//! Reused result state container for the window functions
+	Vector statef;
 };
 
-class WindowCustomAggregate : public WindowAggregateState {
+class WindowCustomAggregator : public WindowAggregator {
 public:
-	WindowCustomAggregate(AggregateObject aggr, const LogicalType &result_type_p, idx_t partition_count);
-	~WindowCustomAggregate() override;
+	WindowCustomAggregator(AggregateObject aggr, const LogicalType &result_type_p, idx_t partition_count);
+	~WindowCustomAggregator() override;
 
-	void Compute(Vector &result, idx_t rid, idx_t start, idx_t end) override;
-
-private:
-	//! The frame boundaries, used for the window functions
-	FrameBounds frame;
+	unique_ptr<WindowAggregatorState> GetLocalState() const override;
+	void Evaluate(WindowAggregatorState &lstate, const idx_t *begins, const idx_t *ends, Vector &result,
+	              idx_t count) const override;
 };
 
-class WindowSegmentTree : public WindowAggregateState {
+class WindowSegmentTree : public WindowAggregator {
 public:
 	using FrameBounds = std::pair<idx_t, idx_t>;
 
@@ -100,32 +123,19 @@ public:
 	~WindowSegmentTree() override;
 
 	void Finalize() override;
-	void Evaluate(const idx_t *begins, const idx_t *ends, Vector &result, idx_t count) override;
 
-private:
+	unique_ptr<WindowAggregatorState> GetLocalState() const override;
+	void Evaluate(WindowAggregatorState &lstate, const idx_t *begins, const idx_t *ends, Vector &result,
+	              idx_t count) const override;
+
+public:
 	void ConstructTree();
-	void ExtractFrame(idx_t begin, idx_t end, data_ptr_t current_state);
-	void FlushStates(bool combining);
-	void WindowSegmentValue(idx_t l_idx, idx_t begin, idx_t end, data_ptr_t current_state);
 
 	//! Use the combine API, if available
 	inline bool UseCombineAPI() const {
 		return mode < WindowAggregationMode::SEPARATE;
 	}
 
-	//! Input data chunk, used for leaf segment aggregation
-	DataChunk leaves;
-	//! The filtered rows in inputs.
-	SelectionVector filter_sel;
-	//! A vector of pointers to "state", used for intermediate window segment aggregation
-	Vector statep;
-	//! The frame boundaries, used for the window functions
-	FrameBounds frame;
-	//! Reused state pointers for combining segment tree levels
-	Vector statel;
-	//! Count of buffered values
-	idx_t flush_count;
-
 	//! The actual window segment tree: an array of aggregate states that represent all the intermediate nodes
 	unsafe_unique_array<data_t> levels_flat_native;
 	//! For each level, the starting location in the levels_flat_native array
@@ -134,12 +144,9 @@ private:
 	//! The total number of internal nodes of the tree, stored in levels_flat_native
 	idx_t internal_nodes;
 
-	//! Use the window API, if available
+	//! Use the combine API, if available
 	WindowAggregationMode mode;
 
-	//! Aggregate allocator
-	ArenaAllocator allocator;
-
 	// TREE_FANOUT needs to cleanly divide STANDARD_VECTOR_SIZE
 	static constexpr idx_t TREE_FANOUT = 16;
 };

package/src/duckdb/src/include/duckdb/storage/arena_allocator.hpp

@@ -46,6 +46,7 @@ public:
 	DUCKDB_API ArenaChunk *GetTail();
 
 	DUCKDB_API bool IsEmpty() const;
+	DUCKDB_API idx_t SizeInBytes() const;
 
 	//! Returns an "Allocator" wrapper for this arena allocator
 	Allocator &GetAllocator() {

package/src/duckdb/src/parser/parser.cpp

@@ -193,48 +193,53 @@ void Parser::ParseQuery(const string &query) {
 			auto query_statements = SplitQueryStringIntoStatements(query);
 			auto stmt_loc = 0;
 			for (auto const &query_statement : query_statements) {
-				PostgresParser another_parser;
-				another_parser.Parse(query_statement);
-				// LCOV_EXCL_START
-				// first see if DuckDB can parse this individual query statement
-				if (another_parser.success) {
-					if (!another_parser.parse_tree) {
-						// empty statement
-						continue;
-					}
-					transformer.TransformParseTree(another_parser.parse_tree, statements);
-					// important to set in the case of a mixture of DDB and parser ext statements
-					statements.back()->stmt_length = query_statement.size() - 1;
-					statements.back()->stmt_location = stmt_loc;
-					stmt_loc += query_statement.size();
-				} else {
-					// let extensions parse the statement which DuckDB failed to parse
-					bool parsed_single_statement = false;
-					for (auto &ext : *options.extensions) {
-						D_ASSERT(!parsed_single_statement);
-						D_ASSERT(ext.parse_function);
-						auto result = ext.parse_function(ext.parser_info.get(), query_statement);
-						if (result.type == ParserExtensionResultType::PARSE_SUCCESSFUL) {
-							auto statement = make_uniq<ExtensionStatement>(ext, std::move(result.parse_data));
-							statement->stmt_length = query_statement.size() - 1;
-							statement->stmt_location = stmt_loc;
-							stmt_loc += query_statement.size();
-							statements.push_back(std::move(statement));
-							parsed_single_statement = true;
-							break;
-						} else if (result.type == ParserExtensionResultType::DISPLAY_EXTENSION_ERROR) {
-							throw ParserException(result.error);
-						} else {
-							// We move to the next one!
+				string another_parser_error;
+				// Creating a new scope to allow extensions to use PostgresParser, which is not reentrant
+				{
+					PostgresParser another_parser;
+					another_parser.Parse(query_statement);
+					// LCOV_EXCL_START
+					// first see if DuckDB can parse this individual query statement
+					if (another_parser.success) {
+						if (!another_parser.parse_tree) {
+							// empty statement
+							continue;
 						}
+						transformer.TransformParseTree(another_parser.parse_tree, statements);
+						// important to set in the case of a mixture of DDB and parser ext statements
+						statements.back()->stmt_length = query_statement.size() - 1;
+						statements.back()->stmt_location = stmt_loc;
+						stmt_loc += query_statement.size();
+						continue;
+					} else {
+						another_parser_error = QueryErrorContext::Format(query, another_parser.error_message,
+						                                                 another_parser.error_location - 1);
 					}
-					if (!parsed_single_statement) {
-						parser_error = QueryErrorContext::Format(query, another_parser.error_message,
-						                                         another_parser.error_location - 1);
-						throw ParserException(parser_error);
+				} // LCOV_EXCL_STOP
+				// LCOV_EXCL_START
+				// let extensions parse the statement which DuckDB failed to parse
+				bool parsed_single_statement = false;
+				for (auto &ext : *options.extensions) {
+					D_ASSERT(!parsed_single_statement);
+					D_ASSERT(ext.parse_function);
+					auto result = ext.parse_function(ext.parser_info.get(), query_statement);
+					if (result.type == ParserExtensionResultType::PARSE_SUCCESSFUL) {
+						auto statement = make_uniq<ExtensionStatement>(ext, std::move(result.parse_data));
+						statement->stmt_length = query_statement.size() - 1;
+						statement->stmt_location = stmt_loc;
+						stmt_loc += query_statement.size();
+						statements.push_back(std::move(statement));
+						parsed_single_statement = true;
+						break;
+					} else if (result.type == ParserExtensionResultType::DISPLAY_EXTENSION_ERROR) {
+						throw ParserException(result.error);
+					} else {
+						// We move to the next one!
 					}
 				}
-				// LCOV_EXCL_STOP
+				if (!parsed_single_statement) {
+					throw ParserException(parser_error);
+				} // LCOV_EXCL_STOP
 			}
 		}
 	}

package/src/duckdb/src/storage/arena_allocator.cpp

@@ -151,4 +151,16 @@ bool ArenaAllocator::IsEmpty() const {
 	return head == nullptr;
 }
 
+idx_t ArenaAllocator::SizeInBytes() const {
+	idx_t total_size = 0;
+	if (!IsEmpty()) {
+		auto current = head.get();
+		while (current != nullptr) {
+			total_size += current->current_position;
+			current = current->next.get();
+		}
+	}
+	return total_size;
+}
+
 } // namespace duckdb