duckdb 0.8.2-dev33.0 → 0.8.2-dev3334.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,47 +16,79 @@
 
 namespace duckdb {
 
-class WindowAggregateState {
+class WindowAggregatorState {
 public:
-	WindowAggregateState(AggregateObject aggr, const LogicalType &result_type_p);
-	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);
 
-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;
 
-	//! 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 statev;
-	//! A vector of pointers to "state", used for intermediate window segment aggregation
-	Vector statep;
-	//! Input data chunk, used for intermediate window segment aggregation
+	//! The cardinality of the partition
+	const idx_t partition_count;
+	//! The size of a single aggregate state
+	const idx_t state_size;
+	//! Partition data chunk
 	DataChunk inputs;
+
+	//! The filtered rows in inputs.
+	vector<validity_t> filter_bits;
+	ValidityMask filter_mask;
+	idx_t filter_pos;
+	//! The state used by the aggregator to build.
+	unique_ptr<WindowAggregatorState> gstate;
 };
 
-class WindowConstantAggregate : public WindowAggregateState {
+class WindowConstantAggregator : public WindowAggregator {
 public:
-	static bool IsConstantAggregate(const BoundWindowExpression &wexpr);
-
-	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 Compute(Vector &result, idx_t rid, idx_t start, idx_t end) 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
@@ -65,53 +97,43 @@ 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;
+	//! Reused result state container for the window functions
+	Vector statef;
 };
 
-class WindowSegmentTree {
+class WindowCustomAggregator : public WindowAggregator {
 public:
-	using FrameBounds = std::pair<idx_t, idx_t>;
+	WindowCustomAggregator(AggregateObject aggr, const LogicalType &result_type_p, idx_t partition_count);
+	~WindowCustomAggregator() override;
 
-	WindowSegmentTree(AggregateObject aggr, const LogicalType &result_type, DataChunk *input,
-	                  const ValidityMask &filter_mask, WindowAggregationMode mode);
-	~WindowSegmentTree();
+	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;
+};
 
-	//! First row contains the result.
-	void Compute(Vector &result, idx_t rid, idx_t start, idx_t end);
+class WindowSegmentTree : public WindowAggregator {
+public:
+	WindowSegmentTree(AggregateObject aggr, const LogicalType &result_type, idx_t count, WindowAggregationMode mode_p);
+	~WindowSegmentTree() override;
 
-private:
+	void Finalize() 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;
+
+public:
 	void ConstructTree();
-	void ExtractFrame(idx_t begin, idx_t end);
-	void WindowSegmentValue(idx_t l_idx, idx_t begin, idx_t end);
-	void AggregateInit();
-	void AggegateFinal(Vector &result, idx_t rid);
 
-	//! Use the window API, if available
-	inline bool UseWindowAPI() const {
-		return mode < WindowAggregationMode::COMBINE;
-	}
 	//! Use the combine API, if available
 	inline bool UseCombineAPI() const {
 		return mode < WindowAggregationMode::SEPARATE;
 	}
 
-	//! The aggregate that the window function is computed over
-	AggregateObject aggr;
-	//! The result type of the window function
-	LogicalType result_type;
-
-	//! Data pointer that contains a single state, used for intermediate window segment aggregation
-	vector<data_t> state;
-	//! Input data chunk, used for intermediate window segment aggregation
-	DataChunk inputs;
-	//! 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 result state container for the window functions
-	Vector statev;
-
 	//! 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
@@ -120,17 +142,11 @@ private:
 	//! The total number of internal nodes of the tree, stored in levels_flat_native
 	idx_t internal_nodes;
 
-	//! The (sorted) input chunk collection on which the tree is built
-	DataChunk *input_ref;
-
-	//! The filtered rows in input_ref.
-	const ValidityMask &filter_mask;
-
-	//! Use the window API, if available
+	//! Use the combine API, if available
 	WindowAggregationMode mode;
 
 	// TREE_FANOUT needs to cleanly divide STANDARD_VECTOR_SIZE
-	static constexpr idx_t TREE_FANOUT = 64;
+	static constexpr idx_t TREE_FANOUT = 16;
 };
 
 } // namespace duckdb

package/src/duckdb/src/include/duckdb/function/aggregate_function.hpp

@@ -8,10 +8,10 @@
 
 #pragma once
 
+#include "duckdb/common/vector_operations/aggregate_executor.hpp"
 #include "duckdb/function/aggregate_state.hpp"
 #include "duckdb/planner/bound_result_modifier.hpp"
 #include "duckdb/planner/expression.hpp"
-#include "duckdb/common/vector_operations/aggregate_executor.hpp"
 
 namespace duckdb {
 
@@ -41,7 +41,6 @@ typedef void (*aggregate_simple_update_t)(Vector inputs[], AggregateInputData &a
                                           data_ptr_t state, idx_t count);
 
 //! The type used for updating complex windowed aggregate functions (optional)
-typedef std::pair<idx_t, idx_t> FrameBounds;
 typedef void (*aggregate_window_t)(Vector inputs[], const ValidityMask &filter_mask,
                                    AggregateInputData &aggr_input_data, idx_t input_count, data_ptr_t state,
                                    const FrameBounds &frame, const FrameBounds &prev, Vector &result, idx_t rid,
@@ -52,6 +51,11 @@ typedef void (*aggregate_serialize_t)(FieldWriter &writer, const FunctionData *b
 typedef unique_ptr<FunctionData> (*aggregate_deserialize_t)(PlanDeserializationState &context, FieldReader &reader,
                                                             AggregateFunction &function);
 
+typedef void (*aggregate_format_serialize_t)(FormatSerializer &serializer, const optional_ptr<FunctionData> bind_data,
+                                             const AggregateFunction &function);
+typedef unique_ptr<FunctionData> (*aggregate_format_deserialize_t)(FormatDeserializer &deserializer,
+                                                                   AggregateFunction &function);
+
 class AggregateFunction : public BaseScalarFunction {
 public:
 	AggregateFunction(const string &name, const vector<LogicalType> &arguments, const LogicalType &return_type,
@@ -66,7 +70,8 @@ public:
 	                         LogicalType(LogicalTypeId::INVALID), null_handling),
 	      state_size(state_size), initialize(initialize), update(update), combine(combine), finalize(finalize),
 	      simple_update(simple_update), window(window), bind(bind), destructor(destructor), statistics(statistics),
-	      serialize(serialize), deserialize(deserialize), order_dependent(AggregateOrderDependent::ORDER_DEPENDENT) {
+	      serialize(serialize), deserialize(deserialize), format_serialize(nullptr), format_deserialize(nullptr),
+	      order_dependent(AggregateOrderDependent::ORDER_DEPENDENT) {
 	}
 
 	AggregateFunction(const string &name, const vector<LogicalType> &arguments, const LogicalType &return_type,
@@ -80,7 +85,8 @@ public:
 	                         LogicalType(LogicalTypeId::INVALID)),
 	      state_size(state_size), initialize(initialize), update(update), combine(combine), finalize(finalize),
 	      simple_update(simple_update), window(window), bind(bind), destructor(destructor), statistics(statistics),
-	      serialize(serialize), deserialize(deserialize), order_dependent(AggregateOrderDependent::ORDER_DEPENDENT) {
+	      serialize(serialize), deserialize(deserialize), format_serialize(nullptr), format_deserialize(nullptr),
+	      order_dependent(AggregateOrderDependent::ORDER_DEPENDENT) {
 	}
 
 	AggregateFunction(const vector<LogicalType> &arguments, const LogicalType &return_type, aggregate_size_t state_size,
@@ -131,6 +137,8 @@ public:
 
 	aggregate_serialize_t serialize;
 	aggregate_deserialize_t deserialize;
+	aggregate_format_serialize_t format_serialize;
+	aggregate_format_deserialize_t format_deserialize;
 	//! Whether or not the aggregate is order dependent
 	AggregateOrderDependent order_dependent;
 

package/src/duckdb/src/include/duckdb/function/aggregate_state.hpp

@@ -21,11 +21,11 @@ enum class AggregateOrderDependent : uint8_t { ORDER_DEPENDENT = 1, NOT_ORDER_DE
 class BoundAggregateExpression;
 
 struct AggregateInputData {
-	AggregateInputData(optional_ptr<FunctionData> bind_data_p, Allocator &allocator_p)
+	AggregateInputData(optional_ptr<FunctionData> bind_data_p, ArenaAllocator &allocator_p)
 	    : bind_data(bind_data_p), allocator(allocator_p) {
 	}
 	optional_ptr<FunctionData> bind_data;
-	Allocator &allocator;
+	ArenaAllocator &allocator;
 };
 
 struct AggregateUnaryInput {

package/src/duckdb/src/include/duckdb/function/built_in_functions.hpp

@@ -57,6 +57,7 @@ private:
 	void RegisterDistributiveAggregates();
 
 	// scalar functions
+	void RegisterCompressedMaterializationFunctions();
 	void RegisterGenericFunctions();
 	void RegisterOperators();
 	void RegisterStringFunctions();

package/src/duckdb/src/include/duckdb/function/copy_function.hpp

@@ -99,13 +99,16 @@ typedef void (*copy_flush_batch_t)(ClientContext &context, FunctionData &bind_da
                                    PreparedBatchData &batch);
 typedef idx_t (*copy_desired_batch_size_t)(ClientContext &context, FunctionData &bind_data);
 
+typedef bool (*copy_supports_type_t)(const LogicalType &type);
+
 class CopyFunction : public Function {
 public:
 	explicit CopyFunction(string name)
 	    : Function(name), plan(nullptr), copy_to_bind(nullptr), copy_to_initialize_local(nullptr),
 	      copy_to_initialize_global(nullptr), copy_to_sink(nullptr), copy_to_combine(nullptr),
 	      copy_to_finalize(nullptr), execution_mode(nullptr), prepare_batch(nullptr), flush_batch(nullptr),
-	      desired_batch_size(nullptr), serialize(nullptr), deserialize(nullptr), copy_from_bind(nullptr) {
+	      desired_batch_size(nullptr), serialize(nullptr), deserialize(nullptr), supports_type(nullptr),
+	      copy_from_bind(nullptr) {
 	}
 
 	//! Plan rewrite copy function
@@ -126,6 +129,8 @@ public:
 	copy_to_serialize_t serialize;
 	copy_to_deserialize_t deserialize;
 
+	copy_supports_type_t supports_type;
+
 	copy_from_bind_t copy_from_bind;
 	TableFunction copy_from_function;