duckdb 1.0.1-dev21.0 → 1.0.1-dev27.0

package/src/duckdb/src/include/duckdb/function/table/arrow/arrow_duck_schema.hpp

@@ -12,85 +12,39 @@
 #include "duckdb/common/unordered_map.hpp"
 #include "duckdb/common/vector.hpp"
 #include "duckdb/common/unique_ptr.hpp"
+#include "duckdb/function/table/arrow/arrow_type_info.hpp"
 
 namespace duckdb {
-//===--------------------------------------------------------------------===//
-// Arrow Variable Size Types
-//===--------------------------------------------------------------------===//
-enum class ArrowVariableSizeType : uint8_t { FIXED_SIZE = 0, NORMAL = 1, SUPER_SIZE = 2 };
-
-//===--------------------------------------------------------------------===//
-// Arrow Time/Date Types
-//===--------------------------------------------------------------------===//
-enum class ArrowDateTimeType : uint8_t {
-	MILLISECONDS = 0,
-	MICROSECONDS = 1,
-	NANOSECONDS = 2,
-	SECONDS = 3,
-	DAYS = 4,
-	MONTHS = 5,
-	MONTH_DAY_NANO = 6
-};
 
 class ArrowType {
 public:
 	//! From a DuckDB type
-	ArrowType(LogicalType type_p) // NOLINT: allow implicit conversion
-	    : type(std::move(type_p)), size_type(ArrowVariableSizeType::NORMAL),
-	      date_time_precision(ArrowDateTimeType::DAYS) {};
-
-	//! From a DuckDB type + fixed_size
-	ArrowType(LogicalType type_p, idx_t fixed_size_p) // NOLINT: work-around bug in clang-tidy
-	    : type(std::move(type_p)), size_type(ArrowVariableSizeType::FIXED_SIZE),
-	      date_time_precision(ArrowDateTimeType::DAYS), fixed_size(fixed_size_p) {};
-
-	//! From a DuckDB type + variable size type
-	ArrowType(LogicalType type_p, ArrowVariableSizeType size_type_p) // NOLINT: work-around bug in clang-tidy
-	    : type(std::move(type_p)), size_type(size_type_p), date_time_precision(ArrowDateTimeType::DAYS) {};
-
-	//! From a DuckDB type + datetime type
-	ArrowType(LogicalType type_p, ArrowDateTimeType date_time_precision_p) // NOLINT: work-around bug in clang-tidy
-	    : type(std::move(type_p)), size_type(ArrowVariableSizeType::NORMAL),
-	      date_time_precision(date_time_precision_p) {};
-
-	void AddChild(unique_ptr<ArrowType> child);
-
-	void AssignChildren(vector<unique_ptr<ArrowType>> children);
+	explicit ArrowType(LogicalType type_p, unique_ptr<ArrowTypeInfo> type_info = nullptr)
+	    : type(std::move(type_p)), type_info(std::move(type_info)) {
+	}
 
+public:
 	LogicalType GetDuckType(bool use_dictionary = false) const;
 
-	ArrowVariableSizeType GetSizeType() const;
-
-	idx_t FixedSize() const;
-
 	void SetDictionary(unique_ptr<ArrowType> dictionary);
-
 	bool HasDictionary() const;
-
-	ArrowDateTimeType GetDateTimeType() const;
-
-	void SetRunEndEncoded();
-
 	const ArrowType &GetDictionary() const;
 
 	bool RunEndEncoded() const;
+	void SetRunEndEncoded();
 
-	const ArrowType &operator[](idx_t index) const;
+	template <class T>
+	const T &GetTypeInfo() const {
+		return type_info->Cast<T>();
+	}
 
 private:
 	LogicalType type;
-	//! If we have a nested type, their children's type.
-	vector<unique_ptr<ArrowType>> children;
-	//! If its a variable size type (e.g., strings, blobs, lists) holds which type it is
-	ArrowVariableSizeType size_type;
-	//! If this is a date/time holds its precision
-	ArrowDateTimeType date_time_precision;
-	//! Only for size types with fixed size
-	idx_t fixed_size = 0;
 	//! Hold the optional type if the array is a dictionary
 	unique_ptr<ArrowType> dictionary_type;
 	//! Is run-end-encoded
 	bool run_end_encoded = false;
+	unique_ptr<ArrowTypeInfo> type_info;
 };
 
 using arrow_column_map_t = unordered_map<idx_t, unique_ptr<ArrowType>>;

package/src/duckdb/src/include/duckdb/function/table/arrow/arrow_type_info.hpp

@@ -0,0 +1,142 @@
+//===----------------------------------------------------------------------===//
+//                         DuckDB
+//
+// duckdb/function/table/arrow/arrow_type_info.hpp
+//
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "duckdb/common/types.hpp"
+#include "duckdb/common/unordered_map.hpp"
+#include "duckdb/common/vector.hpp"
+#include "duckdb/common/unique_ptr.hpp"
+#include "duckdb/function/table/arrow/enum/arrow_type_info_type.hpp"
+#include "duckdb/function/table/arrow/enum/arrow_datetime_type.hpp"
+#include "duckdb/function/table/arrow/enum/arrow_variable_size_type.hpp"
+#include "duckdb/common/enum_util.hpp"
+
+namespace duckdb {
+
+class ArrowType;
+
+struct ArrowTypeInfo {
+public:
+	explicit ArrowTypeInfo(ArrowTypeInfoType type);
+	virtual ~ArrowTypeInfo();
+
+	ArrowTypeInfoType type;
+
+public:
+	template <class TARGET>
+	TARGET &Cast() {
+		D_ASSERT(dynamic_cast<TARGET *>(this));
+		if (type != TARGET::TYPE) {
+			throw InternalException("Failed to cast ArrowTypeInfo, type mismatch (expected: %s, got: %s)",
+			                        EnumUtil::ToString(TARGET::TYPE), EnumUtil::ToString(type));
+		}
+		return reinterpret_cast<TARGET &>(*this);
+	}
+	template <class TARGET>
+	const TARGET &Cast() const {
+		D_ASSERT(dynamic_cast<const TARGET *>(this));
+		if (type != TARGET::TYPE) {
+			throw InternalException("Failed to cast ArrowTypeInfo, type mismatch (expected: %s, got: %s)",
+			                        EnumUtil::ToString(TARGET::TYPE), EnumUtil::ToString(type));
+		}
+		return reinterpret_cast<const TARGET &>(*this);
+	}
+};
+
+struct ArrowStructInfo : public ArrowTypeInfo {
+public:
+	static constexpr const ArrowTypeInfoType TYPE = ArrowTypeInfoType::STRUCT;
+
+public:
+	explicit ArrowStructInfo(vector<unique_ptr<ArrowType>> children);
+	~ArrowStructInfo() override;
+
+public:
+	idx_t ChildCount() const;
+	const ArrowType &GetChild(idx_t index) const;
+	const vector<unique_ptr<ArrowType>> &GetChildren() const;
+
+private:
+	vector<unique_ptr<ArrowType>> children;
+};
+
+struct ArrowDateTimeInfo : public ArrowTypeInfo {
+public:
+	static constexpr const ArrowTypeInfoType TYPE = ArrowTypeInfoType::DATE_TIME;
+
+public:
+	explicit ArrowDateTimeInfo(ArrowDateTimeType size);
+	~ArrowDateTimeInfo() override;
+
+public:
+	ArrowDateTimeType GetDateTimeType() const;
+
+private:
+	ArrowDateTimeType size_type;
+};
+
+struct ArrowStringInfo : public ArrowTypeInfo {
+public:
+	static constexpr const ArrowTypeInfoType TYPE = ArrowTypeInfoType::STRING;
+
+public:
+	explicit ArrowStringInfo(ArrowVariableSizeType size);
+	explicit ArrowStringInfo(idx_t fixed_size);
+	~ArrowStringInfo() override;
+
+public:
+	ArrowVariableSizeType GetSizeType() const;
+	idx_t FixedSize() const;
+
+private:
+	ArrowVariableSizeType size_type;
+	idx_t fixed_size;
+};
+
+struct ArrowListInfo : public ArrowTypeInfo {
+public:
+	static constexpr const ArrowTypeInfoType TYPE = ArrowTypeInfoType::LIST;
+
+public:
+	static unique_ptr<ArrowListInfo> ListView(unique_ptr<ArrowType> child, ArrowVariableSizeType size);
+	static unique_ptr<ArrowListInfo> List(unique_ptr<ArrowType> child, ArrowVariableSizeType size);
+	~ArrowListInfo() override;
+
+public:
+	ArrowVariableSizeType GetSizeType() const;
+	bool IsView() const;
+	ArrowType &GetChild() const;
+
+private:
+	explicit ArrowListInfo(unique_ptr<ArrowType> child, ArrowVariableSizeType size);
+
+private:
+	ArrowVariableSizeType size_type;
+	bool is_view = false;
+	unique_ptr<ArrowType> child;
+};
+
+struct ArrowArrayInfo : public ArrowTypeInfo {
+public:
+	static constexpr const ArrowTypeInfoType TYPE = ArrowTypeInfoType::ARRAY;
+
+public:
+	explicit ArrowArrayInfo(unique_ptr<ArrowType> child, idx_t fixed_size);
+	~ArrowArrayInfo() override;
+
+public:
+	idx_t FixedSize() const;
+	ArrowType &GetChild() const;
+
+private:
+	unique_ptr<ArrowType> child;
+	idx_t fixed_size;
+};
+
+} // namespace duckdb

package/src/duckdb/src/include/duckdb/function/table/arrow/enum/arrow_datetime_type.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+namespace duckdb {
+
+//===--------------------------------------------------------------------===//
+// Arrow Time/Date Types
+//===--------------------------------------------------------------------===//
+enum class ArrowDateTimeType : uint8_t {
+	MILLISECONDS,
+	MICROSECONDS,
+	NANOSECONDS,
+	SECONDS,
+	DAYS,
+	MONTHS,
+	MONTH_DAY_NANO
+};
+
+} // namespace duckdb

package/src/duckdb/src/include/duckdb/function/table/arrow/enum/arrow_type_info_type.hpp

@@ -0,0 +1,7 @@
+#pragma once
+
+namespace duckdb {
+
+enum class ArrowTypeInfoType : uint8_t { LIST, STRUCT, DATE_TIME, STRING, ARRAY };
+
+} // namespace duckdb

package/src/duckdb/src/include/duckdb/function/table/arrow/enum/arrow_variable_size_type.hpp

@@ -0,0 +1,10 @@
+#pragma once
+
+namespace duckdb {
+
+//===--------------------------------------------------------------------===//
+// Arrow Variable Size Types
+//===--------------------------------------------------------------------===//
+enum class ArrowVariableSizeType : uint8_t { NORMAL, FIXED_SIZE, SUPER_SIZE, VIEW };
+
+} // namespace duckdb

package/src/duckdb/src/include/duckdb/function/table/arrow.hpp

@@ -217,6 +217,8 @@ protected:
 	static idx_t ArrowGetBatchIndex(ClientContext &context, const FunctionData *bind_data_p,
 	                                LocalTableFunctionState *local_state, GlobalTableFunctionState *global_state);
 
+	//! Specify if a given type can be pushed-down by the arrow engine
+	static bool ArrowPushdownType(const LogicalType &type);
 	//! -----Utility Functions:-----
 	//! Gets Arrow Table's Cardinality
 	static unique_ptr<NodeStatistics> ArrowScanCardinality(ClientContext &context, const FunctionData *bind_data);

package/src/duckdb/src/include/duckdb/function/table/range.hpp

@@ -49,4 +49,8 @@ struct ReadTextFunction {
 	static void RegisterFunction(BuiltinFunctions &set);
 };
 
+struct QueryTableFunction {
+	static void RegisterFunction(BuiltinFunctions &set);
+};
+
 } // namespace duckdb

package/src/duckdb/src/include/duckdb/function/table/read_csv.hpp

@@ -92,7 +92,7 @@ struct ReadCSVData : public BaseCSVData {
 	unique_ptr<CSVFileScan> initial_reader;
 	//! The union readers are created (when csv union_by_name option is on) during binding
 	//! Those readers can be re-used during ReadCSVFunction
-	vector<unique_ptr<CSVFileScan>> union_readers;
+	vector<unique_ptr<CSVUnionData>> union_readers;
 	//! Reader bind data
 	MultiFileReaderBindData reader_bind;
 
@@ -101,6 +101,9 @@ struct ReadCSVData : public BaseCSVData {
 	void Initialize(unique_ptr<CSVFileScan> &reader) {
 		this->initial_reader = std::move(reader);
 	}
+	void Initialize(ClientContext &, unique_ptr<CSVUnionData> &data) {
+		this->initial_reader = std::move(data->reader);
+	}
 	void FinalizeRead(ClientContext &context);
 
 	void Serialize(Serializer &serializer) const;

package/src/duckdb/src/include/duckdb/function/table/system_functions.hpp

@@ -113,6 +113,10 @@ struct DuckDBTypesFun {
 	static void RegisterFunction(BuiltinFunctions &set);
 };
 
+struct DuckDBVariablesFun {
+	static void RegisterFunction(BuiltinFunctions &set);
+};
+
 struct DuckDBViewsFun {
 	static void RegisterFunction(BuiltinFunctions &set);
 };

package/src/duckdb/src/include/duckdb/function/table/table_scan.hpp

@@ -23,17 +23,17 @@ struct TableScanBindData : public TableFunctionData {
 	//! The table to scan
 	DuckTableEntry &table;
 
-	//! Whether or not the table scan is an index scan
+	//! Whether or not the table scan is an index scan.
 	bool is_index_scan;
-	//! Whether or not the table scan is for index creation
+	//! Whether or not the table scan is for index creation.
 	bool is_create_index;
-	//! The row ids to fetch (in case of an index scan)
-	vector<row_t> result_ids;
+	//! The row ids to fetch in case of an index scan.
+	unsafe_vector<row_t> row_ids;
 
 public:
 	bool Equals(const FunctionData &other_p) const override {
 		auto &other = other_p.Cast<TableScanBindData>();
-		return &other.table == &table && result_ids == other.result_ids;
+		return &other.table == &table && row_ids == other.row_ids;
 	}
 };
 

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

@@ -85,7 +85,7 @@ struct TableFunctionBindInput {
 	TableFunctionBindInput(vector<Value> &inputs, named_parameter_map_t &named_parameters,
 	                       vector<LogicalType> &input_table_types, vector<string> &input_table_names,
 	                       optional_ptr<TableFunctionInfo> info, optional_ptr<Binder> binder,
-	                       const TableFunction &table_function, const TableFunctionRef &ref)
+	                       TableFunction &table_function, const TableFunctionRef &ref)
 	    : inputs(inputs), named_parameters(named_parameters), input_table_types(input_table_types),
 	      input_table_names(input_table_names), info(info), binder(binder), table_function(table_function), ref(ref) {
 	}
@@ -96,7 +96,7 @@ struct TableFunctionBindInput {
 	vector<string> &input_table_names;
 	optional_ptr<TableFunctionInfo> info;
 	optional_ptr<Binder> binder;
-	const TableFunction &table_function;
+	TableFunction &table_function;
 	const TableFunctionRef &ref;
 };
 
@@ -204,6 +204,8 @@ typedef BindInfo (*table_function_get_bind_info_t)(const optional_ptr<FunctionDa
 
 typedef unique_ptr<MultiFileReader> (*table_function_get_multi_file_reader_t)();
 
+typedef bool (*table_function_supports_pushdown_type_t)(const LogicalType &type);
+
 typedef double (*table_function_progress_t)(ClientContext &context, const FunctionData *bind_data,
                                             const GlobalTableFunctionState *global_state);
 typedef void (*table_function_dependency_t)(LogicalDependencyList &dependencies, const FunctionData *bind_data);
@@ -221,6 +223,9 @@ typedef unique_ptr<FunctionData> (*table_function_deserialize_t)(Deserializer &d
 typedef void (*table_function_type_pushdown_t)(ClientContext &context, optional_ptr<FunctionData> bind_data,
                                                const unordered_map<idx_t, LogicalType> &new_column_types);
 
+//! When to call init_global to initialize the table function
+enum class TableFunctionInitialization { INITIALIZE_ON_EXECUTE, INITIALIZE_ON_SCHEDULE };
+
 class TableFunction : public SimpleNamedParameterFunction { // NOLINT: work-around bug in clang-tidy
 public:
 	DUCKDB_API
@@ -280,6 +285,8 @@ public:
 	table_function_type_pushdown_t type_pushdown;
 	//! (Optional) allows injecting a custom MultiFileReader implementation
 	table_function_get_multi_file_reader_t get_multi_file_reader;
+	//! (Optional) If this scanner supports filter pushdown, but not to all data types
+	table_function_supports_pushdown_type_t supports_pushdown_type;
 
 	table_function_serialize_t serialize;
 	table_function_deserialize_t deserialize;
@@ -297,6 +304,11 @@ public:
 	//! Additional function info, passed to the bind
 	shared_ptr<TableFunctionInfo> function_info;
 
+	//! When to call init_global
+	//! By default init_global is called when the pipeline is ready for execution
+	//! If this is set to `INITIALIZE_ON_SCHEDULE` the table function is initialized when the query is scheduled
+	TableFunctionInitialization global_initialization = TableFunctionInitialization::INITIALIZE_ON_EXECUTE;
+
 	DUCKDB_API bool Equal(const TableFunction &rhs) const;
 };
 

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

@@ -32,7 +32,7 @@ public:
 public:
 	unique_ptr<MacroFunction> Copy() const override;
 
-	string ToSQL(const string &schema, const string &name) const override;
+	string ToSQL() const override;
 
 	void Serialize(Serializer &serializer) const override;
 	static unique_ptr<MacroFunction> Deserialize(Deserializer &deserializer);

package/src/duckdb/src/include/duckdb/main/appender.hpp

@@ -27,10 +27,11 @@ enum class AppenderType : uint8_t {
 
 //! The Appender class can be used to append elements to a table.
 class BaseAppender {
-protected:
+public:
 	//! The amount of tuples that will be gathered in the column data collection before flushing
-	static constexpr const idx_t FLUSH_COUNT = STANDARD_VECTOR_SIZE * 100ULL;
+	static constexpr const idx_t DEFAULT_FLUSH_COUNT = STANDARD_VECTOR_SIZE * 100ULL;
 
+protected:
 	Allocator &allocator;
 	//! The append types
 	vector<LogicalType> types;
@@ -42,10 +43,13 @@ protected:
 	idx_t column = 0;
 	//! The type of the appender
 	AppenderType appender_type;
+	//! The amount of rows after which we flush the appender automatically
+	idx_t flush_count = DEFAULT_FLUSH_COUNT;
 
 protected:
 	DUCKDB_API BaseAppender(Allocator &allocator, AppenderType type);
-	DUCKDB_API BaseAppender(Allocator &allocator, vector<LogicalType> types, AppenderType type);
+	DUCKDB_API BaseAppender(Allocator &allocator, vector<LogicalType> types, AppenderType type,
+	                        idx_t flush_count = DEFAULT_FLUSH_COUNT);
 
 public:
 	DUCKDB_API virtual ~BaseAppender();
@@ -114,12 +118,17 @@ class Appender : public BaseAppender {
 	shared_ptr<ClientContext> context;
 	//! The table description (including column names)
 	unique_ptr<TableDescription> description;
+	//! The default expressions
+	unordered_map<idx_t, Value> default_values;
 
 public:
 	DUCKDB_API Appender(Connection &con, const string &schema_name, const string &table_name);
 	DUCKDB_API Appender(Connection &con, const string &table_name);
 	DUCKDB_API ~Appender() override;
 
+public:
+	void AppendDefault();
+
 protected:
 	void FlushInternal(ColumnDataCollection &collection) override;
 };
@@ -131,7 +140,8 @@ class InternalAppender : public BaseAppender {
 	TableCatalogEntry &table;
 
 public:
-	DUCKDB_API InternalAppender(ClientContext &context, TableCatalogEntry &table);
+	DUCKDB_API InternalAppender(ClientContext &context, TableCatalogEntry &table,
+	                            idx_t flush_count = DEFAULT_FLUSH_COUNT);
 	DUCKDB_API ~InternalAppender() override;
 
 protected:

package/src/duckdb/src/include/duckdb/main/attached_database.hpp

@@ -31,19 +31,37 @@ enum class AttachedDatabaseType {
 	TEMP_DATABASE,
 };
 
-//! The AttachedDatabase represents an attached database instance
+//! AttachOptions holds information about a database we plan to attach. These options are generalized, i.e.,
+//! they have to apply to any database file type (duckdb, sqlite, etc.).
+struct AttachOptions {
+	//! Constructor for databases we attach outside of the ATTACH DATABASE statement.
+	explicit AttachOptions(const DBConfigOptions &options);
+	//! Constructor for databases we attach when using ATTACH DATABASE.
+	AttachOptions(const unique_ptr<AttachInfo> &info, const AccessMode default_access_mode);
+
+	//! Defaults to the access mode configured in the DBConfig, unless specified otherwise.
+	AccessMode access_mode;
+	//! The file format type. The default type is a duckdb database file, but other file formats are possible.
+	string db_type;
+	//! We only set this, if we detect any unrecognized option.
+	string unrecognized_option;
+};
+
+//! The AttachedDatabase represents an attached database instance.
 class AttachedDatabase : public CatalogEntry {
 public:
-	//! Create the built-in system attached database (without storage)
+	//! Create the built-in system database (without storage).
 	explicit AttachedDatabase(DatabaseInstance &db, AttachedDatabaseType type = AttachedDatabaseType::SYSTEM_DATABASE);
-	//! Create an attached database instance with the specified name and storage
-	AttachedDatabase(DatabaseInstance &db, Catalog &catalog, string name, string file_path, AccessMode access_mode);
-	//! Create an attached database instance with the specified storage extension
+	//! Create an attached database instance with the specified name and storage.
+	AttachedDatabase(DatabaseInstance &db, Catalog &catalog, string name, string file_path,
+	                 const AttachOptions &options);
+	//! Create an attached database instance with the specified storage extension.
 	AttachedDatabase(DatabaseInstance &db, Catalog &catalog, StorageExtension &ext, ClientContext &context, string name,
-	                 const AttachInfo &info, AccessMode access_mode);
+	                 const AttachInfo &info, const AttachOptions &options);
 	~AttachedDatabase() override;
 
-	void Initialize();
+	//! Initializes the catalog and storage of the attached database.
+	void Initialize(const optional_idx block_alloc_size = optional_idx());
 	void Close();
 
 	Catalog &ParentCatalog() override;

package/src/duckdb/src/include/duckdb/main/buffered_data/batched_buffered_data.hpp

@@ -0,0 +1,79 @@
+//===----------------------------------------------------------------------===//
+//                         DuckDB
+//
+// duckdb/main/batched_buffered_data.hpp
+//
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "duckdb/parallel/interrupt.hpp"
+#include "duckdb/common/deque.hpp"
+#include "duckdb/common/vector_size.hpp"
+#include "duckdb/common/types/data_chunk.hpp"
+#include "duckdb/main/buffered_data/simple_buffered_data.hpp"
+#include "duckdb/common/map.hpp"
+
+namespace duckdb {
+
+class StreamQueryResult;
+
+class InProgressBatch {
+public:
+	//! The chunks that make up the batch
+	deque<unique_ptr<DataChunk>> chunks;
+	//! Whether the batch is completed (NextBatch has been called)
+	bool completed = false;
+};
+
+class BatchedBufferedData : public BufferedData {
+public:
+	static constexpr const BufferedData::Type TYPE = BufferedData::Type::BATCHED;
+
+public:
+	explicit BatchedBufferedData(weak_ptr<ClientContext> context);
+
+public:
+	void Append(const DataChunk &chunk, idx_t batch);
+	void BlockSink(const InterruptState &blocked_sink, idx_t batch);
+
+	bool ShouldBlockBatch(idx_t batch);
+	StreamExecutionResult ExecuteTaskInternal(StreamQueryResult &result, ClientContextLock &context_lock) override;
+	unique_ptr<DataChunk> Scan() override;
+	void UpdateMinBatchIndex(idx_t min_batch_index);
+	bool IsMinimumBatchIndex(lock_guard<mutex> &lock, idx_t batch);
+	void CompleteBatch(idx_t batch);
+	bool BufferIsEmpty();
+	void UnblockSinks() override;
+
+	inline idx_t ReadQueueCapacity() const {
+		return read_queue_capacity;
+	}
+	inline idx_t BufferCapacity() const {
+		return buffer_capacity;
+	}
+
+private:
+	void ResetReplenishState();
+	void MoveCompletedBatches(lock_guard<mutex> &lock);
+
+private:
+	//! The buffer where chunks are written before they are ready to be read.
+	map<idx_t, InProgressBatch> buffer;
+	idx_t buffer_capacity;
+	atomic<idx_t> buffer_byte_count;
+
+	//! The queue containing the chunks that can be read.
+	deque<unique_ptr<DataChunk>> read_queue;
+	idx_t read_queue_capacity;
+	atomic<idx_t> read_queue_byte_count;
+
+	map<idx_t, InterruptState> blocked_sinks;
+
+	idx_t min_batch;
+	//! Debug variable to verify that order is preserved correctly.
+	idx_t lowest_moved_batch = 0;
+};
+
+} // namespace duckdb

package/src/duckdb/src/include/duckdb/main/buffered_data/buffered_data.hpp

@@ -15,6 +15,7 @@
 #include "duckdb/common/optional_idx.hpp"
 #include "duckdb/execution/physical_operator_states.hpp"
 #include "duckdb/common/enums/pending_execution_result.hpp"
+#include "duckdb/common/enums/stream_execution_result.hpp"
 #include "duckdb/common/shared_ptr.hpp"
 
 namespace duckdb {
@@ -22,38 +23,25 @@ namespace duckdb {
 class StreamQueryResult;
 class ClientContextLock;
 
-struct BlockedSink {
-public:
-	BlockedSink(InterruptState state, idx_t chunk_size) : state(std::move(state)), chunk_size(chunk_size) {
-	}
-
-public:
-	//! The handle to reschedule the blocked sink
-	InterruptState state;
-	//! The amount of tuples this sink would add
-	idx_t chunk_size;
-};
-
 class BufferedData {
 protected:
-	enum class Type { SIMPLE };
+	enum class Type { SIMPLE, BATCHED };
 
 public:
-	BufferedData(Type type, weak_ptr<ClientContext> context) : type(type), context(std::move(context)) {
-	}
-	virtual ~BufferedData() {
-	}
+	BufferedData(Type type, weak_ptr<ClientContext> context_p);
+	virtual ~BufferedData();
 
 public:
-	virtual bool BufferIsFull() = 0;
-	virtual PendingExecutionResult ReplenishBuffer(StreamQueryResult &result, ClientContextLock &context_lock) = 0;
+	StreamExecutionResult ReplenishBuffer(StreamQueryResult &result, ClientContextLock &context_lock);
+	virtual StreamExecutionResult ExecuteTaskInternal(StreamQueryResult &result, ClientContextLock &context_lock) = 0;
 	virtual unique_ptr<DataChunk> Scan() = 0;
+	virtual void UnblockSinks() = 0;
 	shared_ptr<ClientContext> GetContext() {
 		return context.lock();
 	}
 	bool Closed() const {
 		if (context.expired()) {
-			return false;
+			return true;
 		}
 		auto c = context.lock();
 		return c == nullptr;
@@ -83,6 +71,8 @@ protected:
 	Type type;
 	//! This is weak to avoid a cyclical reference
 	weak_ptr<ClientContext> context;
+	//! The maximum amount of memory we should keep buffered
+	idx_t total_buffer_size;
 	//! Protect against populate/fetch race condition
 	mutex glock;
 };