Support streaming write of files duckdb ext (#3227)

Author: joseph-isaacs

Cargo.lock

@@ -6,6 +6,40 @@
 
 #include <duckdb/common/unique_ptr.hpp>
 
+struct DType {
+	explicit DType(vx_dtype *dtype): dtype(dtype) {}
+
+	static duckdb::unique_ptr<DType> FromDuckDBTable(
+		const std::vector<duckdb_logical_type> &column_types,
+		const std::vector<unsigned char> &column_nullable,
+		const std::vector<const char *> &column_names
+	) {
+		D_ASSERT(column_names.size() == column_nullable.size());
+		D_ASSERT(column_names.size() == column_types.size());
+
+		vx_error *error = nullptr;
+		auto dtype = vx_duckdb_logical_type_to_dtype(
+			column_types.data(),
+			column_nullable.data(),
+			column_names.data(),
+			column_names.size(),
+			&error
+		);
+		HandleError(error);
+
+		return duckdb::make_uniq<DType>(dtype);
+	}
+
+
+	~DType() {
+		if (dtype != nullptr) {
+			vx_dtype_free(dtype);
+		}
+	}
+
+	vx_dtype *dtype;
+};
+
 struct VortexConversionCache {
 	explicit VortexConversionCache(const unsigned long cache_id) : cache(vx_conversion_cache_create(cache_id)) {
 	}
@@ -32,6 +66,10 @@ struct VortexFileReader {
 		return duckdb::make_uniq<VortexFileReader>(file);
 	}
 
+	struct DType DType() {
+		return ::DType(vx_file_dtype(file));
+	}
+
 	vx_file_reader *file;
 };
 
@@ -45,6 +83,13 @@ struct VortexArray {
 		}
 	}
 
+	static duckdb::unique_ptr<VortexArray> FromDuckDBChunk(DType &dtype, duckdb::DataChunk &chunk) {
+		vx_error *error;
+		auto array = vx_duckdb_chunk_to_array(reinterpret_cast<duckdb_data_chunk>(&chunk), dtype.dtype, &error);
+		HandleError(error);
+		return duckdb::make_uniq<VortexArray>(array);
+	}
+
 	idx_t ToDuckDBVector(idx_t current_row, duckdb_data_chunk output, const VortexConversionCache *cache) const {
 		vx_error *error;
 		auto idx = vx_array_to_duckdb_chunk(array, current_row, output, cache->cache, &error);
@@ -75,3 +120,40 @@ struct VortexArrayStream {
 
 	vx_array_stream *array_stream;
 };
+
+struct ArrayStreamSink {
+	explicit ArrayStreamSink(vx_array_sink *sink, duckdb::unique_ptr<DType> dtype) : sink(sink), dtype(std::move(dtype)) {
+	}
+
+	static duckdb::unique_ptr<ArrayStreamSink> Create(std::string file_path, duckdb::unique_ptr<DType> &&dtype) {
+	    vx_error *error = nullptr;
+        auto sink = vx_array_sink_open_file(file_path.c_str(), dtype->dtype, &error);
+        HandleError(error);
+
+        return duckdb::make_uniq<ArrayStreamSink>(sink, std::move(dtype));
+	}
+
+	void PushChunk(duckdb::DataChunk &chunk) {
+		vx_error *error = nullptr;
+		auto array = VortexArray::FromDuckDBChunk(*dtype, chunk);
+		vx_array_sink_push(sink, array->array, &error);
+		HandleError(error);
+	}
+
+	void Close() {
+		vx_error *error;
+		vx_array_sink_close(sink, &error);
+		HandleError(error);
+
+		this->sink = nullptr;
+	}
+
+	~ArrayStreamSink() {
+		// "should dctor a sink, before closing it
+		D_ASSERT(sink == nullptr);
+	}
+
+
+	vx_array_sink *sink;
+	duckdb::unique_ptr<DType> dtype;
+};

duckdb-vortex/src/include/vortex_common.hpp

@@ -129,17 +129,17 @@ void CreateFilterExpression(google::protobuf::Arena &arena, vector<std::string>
 }
 
 /// Extracts schema information from a Vortex file's data type.
-static void ExtractVortexSchema(const vx_dtype *file_dtype, vector<LogicalType> &column_types,
+static void ExtractVortexSchema(DType &file_dtype, vector<LogicalType> &column_types,
                                 vector<string> &column_names) {
-	uint32_t field_count = vx_dtype_field_count(file_dtype);
+	uint32_t field_count = vx_dtype_field_count(file_dtype.dtype);
 	for (uint32_t idx = 0; idx < field_count; idx++) {
 		char name_buffer[512];
 		int name_len = 0;
 
-		vx_dtype_field_name(file_dtype, idx, name_buffer, &name_len);
+		vx_dtype_field_name(file_dtype.dtype, idx, name_buffer, &name_len);
 		std::string field_name(name_buffer, name_len);
 
-		vx_dtype *field_dtype = vx_dtype_field_dtype(file_dtype, idx);
+		vx_dtype *field_dtype = vx_dtype_field_dtype(file_dtype.dtype, idx);
 		vx_error *error = nullptr;
 		auto duckdb_type = vx_dtype_to_duckdb_logical_type(field_dtype, &error);
 		HandleError(error);
@@ -179,8 +179,8 @@ static unique_ptr<VortexFileReader> OpenFile(const std::string &filename, vector
 	}
 
 	// This Ptr is owned by the file
-	const vx_dtype *file_dtype = vx_file_dtype(file->file);
-	if (vx_dtype_get(file_dtype) != DTYPE_STRUCT) {
+	auto file_dtype = file->DType();
+	if (vx_dtype_get(file_dtype.dtype) != DTYPE_STRUCT) {
 		vx_file_reader_free(file->file);
 		throw FatalException("Vortex file does not contain a struct array as a top-level dtype");
 	}

duckdb-vortex/src/vortex_scan.cpp

@@ -10,6 +10,8 @@
 #include "duckdb/function/copy_function.hpp"
 #include "duckdb/parser/constraints/not_null_constraint.hpp"
 
+// TODO(joe): enable multi-threaded writes, see `VortexWriteSink`.
+
 namespace duckdb {
 
 struct VortexWriteBindData : public TableFunctionData {
@@ -21,9 +23,7 @@ struct VortexWriteBindData : public TableFunctionData {
 };
 
 struct VortexWriteGlobalData : public GlobalFunctionData {
-	std::string file_name;
-	std::unique_ptr<VortexFileReader> file;
-	unique_ptr<VortexArray> array;
+	unique_ptr<ArrayStreamSink> sink;
 };
 
 struct VortexWriteLocalData : public LocalFunctionData {};
@@ -33,16 +33,12 @@ void VortexWriteSink(ExecutionContext &context, FunctionData &bind_data, GlobalF
 	auto &global_state = gstate.Cast<VortexWriteGlobalData>();
 	auto bind = bind_data.Cast<VortexWriteBindData>();
 
-	auto chunk = DataChunk();
-	chunk.Initialize(Allocator::Get(context.client), bind.sql_types);
-
 	for (auto i = 0u; i < input.ColumnCount(); i++) {
 		input.data[i].Flatten(input.size());
 	}
-
-	auto new_array = vx_array_append_duckdb_chunk(
-	    global_state.array->array, reinterpret_cast<duckdb_data_chunk>(&input), bind.column_nullable.data());
-	global_state.array = make_uniq<VortexArray>(new_array);
+	// TODO(joe): go to a model of combining local chunked into arrays of a specific size
+	// before push each of these larger chunks into the global_state
+	global_state.sink->PushChunk(input);
 }
 
 std::vector<idx_t> TableNullability(ClientContext &context, const string &catalog_name, const string &schema,
@@ -92,7 +88,6 @@ void RegisterVortexWriteFunction(DatabaseInstance &instance) {
 	                                        const string &file_path) -> unique_ptr<GlobalFunctionData> {
 		auto &bind = bind_data.Cast<VortexWriteBindData>();
 		auto gstate = make_uniq<VortexWriteGlobalData>();
-		gstate->file_name = file_path;
 
 		auto column_names = std::vector<const char *>();
 		for (const auto &col_id : bind.column_names) {
@@ -104,12 +99,8 @@ void RegisterVortexWriteFunction(DatabaseInstance &instance) {
 			column_types.push_back(reinterpret_cast<duckdb_logical_type>(&col_type));
 		}
 
-		vx_error *error = nullptr;
-		auto array = vx_array_create_empty_from_duckdb_table(column_types.data(), bind.column_nullable.data(),
-		                                                     column_names.data(), column_names.size(), &error);
-		HandleError(error);
-
-		gstate->array = make_uniq<VortexArray>(array);
+		auto dtype = DType::FromDuckDBTable(column_types, bind.column_nullable, column_names);
+		gstate->sink = ArrayStreamSink::Create(file_path, std::move(dtype));
 		return std::move(gstate);
 	};
 	function.copy_to_initialize_local = [](ExecutionContext &context,
@@ -119,9 +110,7 @@ void RegisterVortexWriteFunction(DatabaseInstance &instance) {
 	function.copy_to_sink = VortexWriteSink;
 	function.copy_to_finalize = [](ClientContext &context, FunctionData &bind_data, GlobalFunctionData &gstate) {
 		auto &global_state = gstate.Cast<VortexWriteGlobalData>();
-		vx_error *error;
-		vx_file_write_array(global_state.file_name.c_str(), global_state.array->array, &error);
-		HandleError(error);
+		global_state.sink->Close();
 	};
 	function.execution_mode = [](bool preserve_insertion_order,
 	                             bool supports_batch_index) -> CopyFunctionExecutionMode {

duckdb-vortex/src/vortex_write.cpp

@@ -27,6 +27,7 @@ paste = { workspace = true }
 prost = { workspace = true }
 simplelog = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread"] }
+tokio-stream = { workspace = true }
 url = { workspace = true, features = [] }
 vortex = { workspace = true, features = ["object_store", "proto"] }
 vortex-duckdb = { workspace = true, optional = true }

vortex-ffi/Cargo.toml

@@ -94,6 +94,12 @@ typedef struct vx_dtype vx_dtype;
  */
 typedef struct vx_array vx_array;
 
+/**
+ * The `sink` interface is used to collect array chunks and place them into a resource
+ * (e.g. an array stream or file (`vx_array_sink_open_file`)).
+ */
+typedef struct vx_array_sink vx_array_sink;
+
 /**
  * FFI-exposed stream interface.
  */
@@ -284,6 +290,17 @@ duckdb_logical_type vx_dtype_to_duckdb_logical_type(struct vx_dtype *dtype,
                                                     struct vx_error **error);
 #endif
 
+#if defined(ENABLE_DUCKDB_FFI)
+/**
+ * Converts a DuckDB type into a vortex type
+ */
+struct vx_dtype *vx_duckdb_logical_type_to_dtype(const duckdb_logical_type *column_types,
+                                                 const unsigned char *column_nullable,
+                                                 const char *const *column_names,
+                                                 int column_count,
+                                                 struct vx_error **error);
+#endif
+
 #if defined(ENABLE_DUCKDB_FFI)
 /**
  * Back a single chunk of the array as a duckdb data chunk.
@@ -300,22 +317,11 @@ unsigned int vx_array_to_duckdb_chunk(struct vx_array *stream,
 
 #if defined(ENABLE_DUCKDB_FFI)
 /**
- * Returns an empty vortex array constructed from three arrays of len `len`, the (types, null, names).
- */
-struct vx_array *vx_array_create_empty_from_duckdb_table(const duckdb_logical_type *type_array,
-                                                         const unsigned char *nullable,
-                                                         const char *const *names,
-                                                         int len,
-                                                         struct vx_error **error);
-#endif
-
-#if defined(ENABLE_DUCKDB_FFI)
-/**
- * Requires a vortex array, a duckdb data chunk and a nullable array (equal to len(chunk.columns)).
+ * Pushed a single duckdb chunk into a file sink.
  */
-struct vx_array *vx_array_append_duckdb_chunk(struct vx_array *array,
-                                              duckdb_data_chunk chunk,
-                                              const unsigned char *nullable);
+struct vx_array *vx_duckdb_chunk_to_array(duckdb_data_chunk chunk,
+                                          struct vx_dtype *dtype,
+                                          struct vx_error **error);
 #endif
 
 #if defined(ENABLE_DUCKDB_FFI)
@@ -345,19 +351,14 @@ void vx_error_free(struct vx_error *error);
 struct vx_file_reader *vx_file_open_reader(const struct vx_file_open_options *options,
                                            struct vx_error **error);
 
-void vx_file_write_array(const char *path, struct vx_array *ffi_array, struct vx_error **error);
-
 struct vx_file_statistics *vx_file_extract_statistics(struct vx_file_reader *file);
 
 void vx_file_statistics_free(struct vx_file_statistics *stat);
 
 /**
- * Get a readonly pointer to the DType of the data inside of the file.
- *
- * The pointer's lifetime is tied to the lifetime of the underlying file, so it should not be
- * dereferenced after the file has been freed.
+ * Get the DType of the data inside of the file.
  */
-const struct vx_dtype *vx_file_dtype(const struct vx_file_reader *file);
+struct vx_dtype *vx_file_dtype(const struct vx_file_reader *file);
 
 /**
  * Build a new `vx_array_stream` that return a series of `vx_array`s scan over a `vx_file`.
@@ -382,6 +383,27 @@ void vx_file_reader_free(struct vx_file_reader *file);
  */
 void vx_init_logging(enum vx_log_level level);
 
+/**
+ * Opens a writable array stream, where sink is used to push values into the stream.
+ * To close the stream close the sink with `vx_array_sink_close`.
+ */
+struct vx_array_sink *vx_array_sink_open_file(const char *path,
+                                              const struct vx_dtype *dtype,
+                                              struct vx_error **error);
+
+/**
+ * Pushed a single array chunk into a file sink.
+ */
+void vx_array_sink_push(struct vx_array_sink *sink,
+                        const struct vx_array *array,
+                        struct vx_error **error);
+
+/**
+ * Closes an array sink, must be called to ensure all the values pushed to the sink are written
+ * to the external resource.
+ */
+void vx_array_sink_close(struct vx_array_sink *sink, struct vx_error **error);
+
 /**
  * Gets the dtype from an array `stream`, if the stream is finished the `DType` is null
  */