duckdb 0.8.2-dev161.0 → 0.8.2-dev1724.0

package/src/duckdb/src/common/adbc/adbc.cpp

@@ -44,10 +44,23 @@ duckdb_adbc::AdbcStatusCode duckdb_adbc_init(size_t count, struct duckdb_adbc::A
 	driver->ConnectionRollback = duckdb_adbc::ConnectionRollback;
 	driver->ConnectionReadPartition = duckdb_adbc::ConnectionReadPartition;
 	driver->StatementExecutePartitions = duckdb_adbc::StatementExecutePartitions;
+	driver->ConnectionGetTableSchema = duckdb_adbc::ConnectionGetTableSchema;
+	driver->StatementSetSubstraitPlan = duckdb_adbc::StatementSetSubstraitPlan;
+
 	return ADBC_STATUS_OK;
 }
 
 namespace duckdb_adbc {
+
+struct DuckDBAdbcStatementWrapper {
+	::duckdb_connection connection;
+	::duckdb_arrow result;
+	::duckdb_prepared_statement statement;
+	char *ingestion_table_name;
+	ArrowArrayStream *ingestion_stream;
+};
+static AdbcStatusCode QueryInternal(struct AdbcConnection *connection, struct ArrowArrayStream *out, const char *query,
+                                    struct AdbcError *error);
 AdbcStatusCode SetErrorMaybe(const void *result, AdbcError *error, const std::string &error_message) {
 	if (!error) {
 		return ADBC_STATUS_INVALID_ARGUMENT;
@@ -109,6 +122,28 @@ AdbcStatusCode DatabaseNew(struct AdbcDatabase *database, struct AdbcError *erro
 	return CheckResult(res, error, "Failed to allocate");
 }
 
+AdbcStatusCode StatementSetSubstraitPlan(struct AdbcStatement *statement, const uint8_t *plan, size_t length,
+                                         struct AdbcError *error) {
+	if (!statement) {
+		SetError(error, "Statement is not set");
+		return ADBC_STATUS_INVALID_ARGUMENT;
+	}
+	if (!plan) {
+		SetError(error, "Substrait Plan is not set");
+		return ADBC_STATUS_INVALID_ARGUMENT;
+	}
+	if (length == 0) {
+		SetError(error, "Can't execute plan with size = 0");
+		return ADBC_STATUS_INVALID_ARGUMENT;
+	}
+	auto wrapper = reinterpret_cast<DuckDBAdbcStatementWrapper *>(statement->private_data);
+	auto plan_str = std::string(reinterpret_cast<const char *>(plan), length);
+	auto query = "CALL from_substrait('" + plan_str + "'::BLOB)";
+	auto res = duckdb_prepare(wrapper->connection, query.c_str(), &wrapper->statement);
+	auto error_msg = duckdb_prepare_error(wrapper->statement);
+	return CheckResult(res, error, error_msg);
+}
+
 AdbcStatusCode DatabaseSetOption(struct AdbcDatabase *database, const char *key, const char *value,
                                  struct AdbcError *error) {
 	auto status = SetErrorMaybe(database, error, "Missing database object");
@@ -159,6 +194,44 @@ AdbcStatusCode DatabaseRelease(struct AdbcDatabase *database, struct AdbcError *
 	return ADBC_STATUS_OK;
 }
 
+AdbcStatusCode ConnectionGetTableSchema(struct AdbcConnection *connection, const char *catalog, const char *db_schema,
+                                        const char *table_name, struct ArrowSchema *schema, struct AdbcError *error) {
+	if (!connection) {
+		SetError(error, "Connection is not set");
+		return ADBC_STATUS_INVALID_ARGUMENT;
+	}
+	if (catalog != nullptr && strlen(catalog) > 0) {
+		// In DuckDB this is the name of the database, not sure what's the expected functionality here, so for now,
+		// scream.
+		SetError(error, "Catalog Name is not used in DuckDB. It must be set to nullptr or an empty string");
+		return ADBC_STATUS_NOT_IMPLEMENTED;
+	} else if (db_schema == nullptr) {
+		SetError(error, "AdbcConnectionGetTableSchema: must provide db_schema");
+		return ADBC_STATUS_INVALID_ARGUMENT;
+	} else if (table_name == nullptr) {
+		SetError(error, "AdbcConnectionGetTableSchema: must provide table_name");
+		return ADBC_STATUS_INVALID_ARGUMENT;
+	} else if (strlen(table_name) == 0) {
+		SetError(error, "AdbcConnectionGetTableSchema: must provide table_name");
+		return ADBC_STATUS_INVALID_ARGUMENT;
+	}
+	ArrowArrayStream arrow_stream;
+
+	std::string query = "SELECT * FROM ";
+	if (strlen(db_schema) > 0) {
+		query += std::string(db_schema) + ".";
+	}
+	query += std::string(table_name) + " LIMIT 0;";
+
+	auto success = QueryInternal(connection, &arrow_stream, query.c_str(), error);
+	if (success != ADBC_STATUS_OK) {
+		return success;
+	}
+	arrow_stream.get_schema(&arrow_stream, schema);
+	arrow_stream.release(&arrow_stream);
+	return ADBC_STATUS_OK;
+}
+
 AdbcStatusCode ConnectionNew(struct AdbcConnection *connection, struct AdbcError *error) {
 	auto status = SetErrorMaybe(connection, error, "Missing connection object");
 	if (status != ADBC_STATUS_OK) {
@@ -400,14 +473,6 @@ AdbcStatusCode Ingest(duckdb_connection connection, const char *table_name, stru
 	return ADBC_STATUS_OK;
 }
 
-struct DuckDBAdbcStatementWrapper {
-	::duckdb_connection connection;
-	::duckdb_arrow result;
-	::duckdb_prepared_statement statement;
-	char *ingestion_table_name;
-	ArrowArrayStream *ingestion_stream;
-};
-
 AdbcStatusCode StatementNew(struct AdbcConnection *connection, struct AdbcStatement *statement,
                             struct AdbcError *error) {
 
@@ -581,8 +646,8 @@ AdbcStatusCode StatementSetOption(struct AdbcStatement *statement, const char *k
 	return ADBC_STATUS_INVALID_ARGUMENT;
 }
 
-static AdbcStatusCode QueryInternal(struct AdbcConnection *connection, struct ArrowArrayStream *out, const char *query,
-                                    struct AdbcError *error) {
+AdbcStatusCode QueryInternal(struct AdbcConnection *connection, struct ArrowArrayStream *out, const char *query,
+                             struct AdbcError *error) {
 	AdbcStatement statement;
 
 	auto status = StatementNew(connection, &statement, error);

package/src/duckdb/src/common/adbc/driver_manager.cpp

@@ -109,8 +109,9 @@ struct ManagerDriverState {
 static AdbcStatusCode ReleaseDriver(struct AdbcDriver *driver, struct AdbcError *error) {
 	AdbcStatusCode status = ADBC_STATUS_OK;
 
-	if (!driver->private_manager)
+	if (!driver->private_manager) {
 		return status;
+	}
 	ManagerDriverState *state = reinterpret_cast<ManagerDriverState *>(driver->private_manager);
 
 	if (state->driver_release) {
@@ -138,11 +139,6 @@ AdbcStatusCode ConnectionGetInfo(struct AdbcConnection *connection, uint32_t *in
 	return ADBC_STATUS_NOT_IMPLEMENTED;
 }
 
-AdbcStatusCode ConnectionGetTableSchema(struct AdbcConnection *, const char *, const char *, const char *,
-                                        struct ArrowSchema *, struct AdbcError *error) {
-	return ADBC_STATUS_NOT_IMPLEMENTED;
-}
-
 AdbcStatusCode StatementBind(struct AdbcStatement *, struct ArrowArray *, struct ArrowSchema *,
                              struct AdbcError *error) {
 	return ADBC_STATUS_NOT_IMPLEMENTED;
@@ -152,9 +148,6 @@ AdbcStatusCode StatementGetParameterSchema(struct AdbcStatement *statement, stru
                                            struct AdbcError *error) {
 	return ADBC_STATUS_NOT_IMPLEMENTED;
 }
-AdbcStatusCode StatementSetSubstraitPlan(struct AdbcStatement *, const uint8_t *, size_t, struct AdbcError *error) {
-	return ADBC_STATUS_NOT_IMPLEMENTED;
-}
 
 /// Temporary state while the database is being configured.
 struct TempDatabase {
@@ -351,15 +344,17 @@ AdbcStatusCode AdbcConnectionInit(struct AdbcConnection *connection, struct Adbc
 	delete args;
 
 	auto status = database->private_driver->ConnectionNew(connection, error);
-	if (status != ADBC_STATUS_OK)
+	if (status != ADBC_STATUS_OK) {
 		return status;
+	}
 	connection->private_driver = database->private_driver;
 
 	for (const auto &option : options) {
 		status = database->private_driver->ConnectionSetOption(connection, option.first.c_str(), option.second.c_str(),
 		                                                       error);
-		if (status != ADBC_STATUS_OK)
+		if (status != ADBC_STATUS_OK) {
 			return status;
+		}
 	}
 	return connection->private_driver->ConnectionInit(connection, database, error);
 }

package/src/duckdb/src/common/allocator.cpp

@@ -14,7 +14,13 @@
 #include <execinfo.h>
 #endif
 
-#if defined(BUILD_JEMALLOC_EXTENSION) && !defined(WIN32)
+#ifndef USE_JEMALLOC
+#if defined(DUCKDB_EXTENSION_JEMALLOC_LINKED) && DUCKDB_EXTENSION_JEMALLOC_LINKED && !defined(WIN32)
+#define USE_JEMALLOC
+#endif
+#endif
+
+#ifdef USE_JEMALLOC
 #include "jemalloc_extension.hpp"
 #endif
 
@@ -89,7 +95,7 @@ PrivateAllocatorData::~PrivateAllocatorData() {
 //===--------------------------------------------------------------------===//
 // Allocator
 //===--------------------------------------------------------------------===//
-#if defined(BUILD_JEMALLOC_EXTENSION) && !defined(WIN32)
+#ifdef USE_JEMALLOC
 Allocator::Allocator()
     : Allocator(JemallocExtension::Allocate, JemallocExtension::Free, JemallocExtension::Reallocate, nullptr) {
 }
@@ -177,6 +183,12 @@ Allocator &Allocator::DefaultAllocator() {
 	return *DefaultAllocatorReference();
 }
 
+void Allocator::ThreadFlush(idx_t threshold) {
+#ifdef USE_JEMALLOC
+	JemallocExtension::ThreadFlush(threshold);
+#endif
+}
+
 //===--------------------------------------------------------------------===//
 // Debug Info (extended)
 //===--------------------------------------------------------------------===//

package/src/duckdb/src/common/arrow/arrow_appender.cpp

@@ -153,7 +153,9 @@ struct ArrowIntervalConverter {
 template <class TGT, class SRC = TGT, class OP = ArrowScalarConverter>
 struct ArrowScalarBaseData {
 	static void Append(ArrowAppendData &append_data, Vector &input, idx_t from, idx_t to, idx_t input_size) {
+		D_ASSERT(to >= from);
 		idx_t size = to - from;
+		D_ASSERT(size <= input_size);
 		UnifiedVectorFormat format;
 		input.ToUnifiedFormat(input_size, format);
 
@@ -392,6 +394,76 @@ struct ArrowVarcharData {
 	}
 };
 
+//===--------------------------------------------------------------------===//
+// Unions
+//===--------------------------------------------------------------------===//
+/**
+ * Based on https://arrow.apache.org/docs/format/Columnar.html#union-layout &
+ * https://arrow.apache.org/docs/format/CDataInterface.html
+ */
+struct ArrowUnionData {
+	static void Initialize(ArrowAppendData &result, const LogicalType &type, idx_t capacity) {
+		result.main_buffer.reserve(capacity * sizeof(int8_t));
+
+		for (auto &child : UnionType::CopyMemberTypes(type)) {
+			auto child_buffer = InitializeArrowChild(child.second, capacity, result.options);
+			result.child_data.push_back(std::move(child_buffer));
+		}
+	}
+
+	static void Append(ArrowAppendData &append_data, Vector &input, idx_t from, idx_t to, idx_t input_size) {
+		UnifiedVectorFormat format;
+		input.ToUnifiedFormat(input_size, format);
+		idx_t size = to - from;
+
+		auto &types_buffer = append_data.main_buffer;
+
+		duckdb::vector<Vector> child_vectors;
+		for (const auto &child : UnionType::CopyMemberTypes(input.GetType())) {
+			child_vectors.emplace_back(child.second);
+		}
+
+		for (idx_t input_idx = from; input_idx < to; input_idx++) {
+			const auto &val = input.GetValue(input_idx);
+
+			idx_t tag = 0;
+			Value resolved_value(nullptr);
+			if (!val.IsNull()) {
+				tag = UnionValue::GetTag(val);
+
+				resolved_value = UnionValue::GetValue(val);
+			}
+
+			for (idx_t child_idx = 0; child_idx < child_vectors.size(); child_idx++) {
+				child_vectors[child_idx].SetValue(input_idx, child_idx == tag ? resolved_value : Value(nullptr));
+			}
+
+			types_buffer.data()[input_idx] = tag;
+		}
+
+		for (idx_t child_idx = 0; child_idx < child_vectors.size(); child_idx++) {
+			auto &child_buffer = append_data.child_data[child_idx];
+			auto &child = child_vectors[child_idx];
+			child_buffer->append_vector(*child_buffer, child, from, to, size);
+		}
+		append_data.row_count += size;
+	}
+
+	static void Finalize(ArrowAppendData &append_data, const LogicalType &type, ArrowArray *result) {
+		result->n_buffers = 2;
+		result->buffers[1] = append_data.main_buffer.data();
+
+		auto &child_types = UnionType::CopyMemberTypes(type);
+		append_data.child_pointers.resize(child_types.size());
+		result->children = append_data.child_pointers.data();
+		result->n_children = child_types.size();
+		for (idx_t i = 0; i < child_types.size(); i++) {
+			auto &child_type = child_types[i].second;
+			append_data.child_pointers[i] = FinalizeArrowChild(child_type, *append_data.child_data[i]);
+		}
+	}
+};
+
 //===--------------------------------------------------------------------===//
 // Structs
 //===--------------------------------------------------------------------===//
@@ -489,16 +561,9 @@ struct ArrowListData {
 		SelectionVector child_sel(child_indices.data());
 		auto &child = ListVector::GetEntry(input);
 		auto child_size = child_indices.size();
-		if (size != input_size) {
-			// Let's avoid doing this
-			Vector child_copy(child.GetType());
-			child_copy.Slice(child, child_sel, child_size);
-			append_data.child_data[0]->append_vector(*append_data.child_data[0], child_copy, 0, child_size, child_size);
-		} else {
-			// We don't care about the vector, slice it
-			child.Slice(child_sel, child_size);
-			append_data.child_data[0]->append_vector(*append_data.child_data[0], child, 0, child_size, child_size);
-		}
+		Vector child_copy(child.GetType());
+		child_copy.Slice(child, child_sel, child_size);
+		append_data.child_data[0]->append_vector(*append_data.child_data[0], child_copy, 0, child_size, child_size);
 		append_data.row_count += size;
 	}
 
@@ -721,6 +786,9 @@ static void InitializeFunctionPointers(ArrowAppendData &append_data, const Logic
 	case LogicalTypeId::INTERVAL:
 		InitializeFunctionPointers<ArrowScalarData<ArrowInterval, interval_t, ArrowIntervalConverter>>(append_data);
 		break;
+	case LogicalTypeId::UNION:
+		InitializeFunctionPointers<ArrowUnionData>(append_data);
+		break;
 	case LogicalTypeId::STRUCT:
 		InitializeFunctionPointers<ArrowStructData>(append_data);
 		break;
@@ -731,7 +799,7 @@ static void InitializeFunctionPointers(ArrowAppendData &append_data, const Logic
 		InitializeFunctionPointers<ArrowMapData>(append_data);
 		break;
 	default:
-		throw InternalException("Unsupported type in DuckDB -> Arrow Conversion: %s\n", type.ToString());
+		throw NotImplementedException("Unsupported type in DuckDB -> Arrow Conversion: %s\n", type.ToString());
 	}
 }
 
@@ -791,7 +859,6 @@ ArrowArray ArrowAppender::Finalize() {
 
 	// Configure root array
 	result.length = row_count;
-	result.n_children = types.size();
 	result.n_buffers = 1;
 	result.buffers = root_holder->buffers.data(); // there is no actual buffer there since we don't have NULLs
 	result.offset = 0;

package/src/duckdb/src/common/arrow/arrow_converter.cpp

@@ -76,6 +76,15 @@ void SetArrowMapFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child,
 	SetArrowFormat(root_holder, **child.children, ListType::GetChildType(type), options);
 }
 
+unsafe_unique_array<char> AddName(const string &name) {
+	auto name_ptr = make_unsafe_uniq_array<char>(name.size() + 1);
+	for (size_t i = 0; i < name.size(); i++) {
+		name_ptr[i] = name[i];
+	}
+	name_ptr[name.size()] = '\0';
+	return name_ptr;
+}
+
 void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, const LogicalType &type,
                     const ArrowOptions &options) {
 	switch (type.id()) {
@@ -135,12 +144,7 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		break;
 	case LogicalTypeId::TIMESTAMP_TZ: {
 		string format = "tsu:" + options.time_zone;
-		auto format_ptr = make_unsafe_uniq_array<char>(format.size() + 1);
-		for (size_t i = 0; i < format.size(); i++) {
-			format_ptr[i] = format[i];
-		}
-		format_ptr[format.size()] = '\0';
-		root_holder.owned_type_names.push_back(std::move(format_ptr));
+		root_holder.owned_type_names.push_back(AddName(format));
 		child.format = root_holder.owned_type_names.back().get();
 		break;
 	}
@@ -160,12 +164,7 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		uint8_t width, scale;
 		type.GetDecimalProperties(width, scale);
 		string format = "d:" + to_string(width) + "," + to_string(scale);
-		auto format_ptr = make_unsafe_uniq_array<char>(format.size() + 1);
-		for (size_t i = 0; i < format.size(); i++) {
-			format_ptr[i] = format[i];
-		}
-		format_ptr[format.size()] = '\0';
-		root_holder.owned_type_names.push_back(std::move(format_ptr));
+		root_holder.owned_type_names.push_back(AddName(format));
 		child.format = root_holder.owned_type_names.back().get();
 		break;
 	}
@@ -211,13 +210,7 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 
 			InitializeChild(*child.children[type_idx]);
 
-			auto &struct_col_name = child_types[type_idx].first;
-			auto name_ptr = make_unsafe_uniq_array<char>(struct_col_name.size() + 1);
-			for (size_t i = 0; i < struct_col_name.size(); i++) {
-				name_ptr[i] = struct_col_name[i];
-			}
-			name_ptr[struct_col_name.size()] = '\0';
-			root_holder.owned_type_names.push_back(std::move(name_ptr));
+			root_holder.owned_type_names.push_back(AddName(child_types[type_idx].first));
 
 			child.children[type_idx]->name = root_holder.owned_type_names.back().get();
 			SetArrowFormat(root_holder, *child.children[type_idx], child_types[type_idx].second, options);
@@ -228,6 +221,38 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		SetArrowMapFormat(root_holder, child, type, options);
 		break;
 	}
+	case LogicalTypeId::UNION: {
+		std::string format = "+us:";
+
+		auto &child_types = UnionType::CopyMemberTypes(type);
+		child.n_children = child_types.size();
+		root_holder.nested_children.emplace_back();
+		root_holder.nested_children.back().resize(child_types.size());
+		root_holder.nested_children_ptr.emplace_back();
+		root_holder.nested_children_ptr.back().resize(child_types.size());
+		for (idx_t type_idx = 0; type_idx < child_types.size(); type_idx++) {
+			root_holder.nested_children_ptr.back()[type_idx] = &root_holder.nested_children.back()[type_idx];
+		}
+		child.children = &root_holder.nested_children_ptr.back()[0];
+		for (size_t type_idx = 0; type_idx < child_types.size(); type_idx++) {
+
+			InitializeChild(*child.children[type_idx]);
+
+			root_holder.owned_type_names.push_back(AddName(child_types[type_idx].first));
+
+			child.children[type_idx]->name = root_holder.owned_type_names.back().get();
+			SetArrowFormat(root_holder, *child.children[type_idx], child_types[type_idx].second, options);
+
+			format += to_string(type_idx) + ",";
+		}
+
+		format.pop_back();
+
+		root_holder.owned_type_names.push_back(AddName(format));
+		child.format = root_holder.owned_type_names.back().get();
+
+		break;
+	}
 	case LogicalTypeId::ENUM: {
 		// TODO what do we do with pointer enums here?
 		switch (EnumType::GetPhysicalType(type)) {

package/src/duckdb/src/common/assert.cpp

@@ -4,6 +4,9 @@
 namespace duckdb {
 
 void DuckDBAssertInternal(bool condition, const char *condition_name, const char *file, int linenr) {
+#ifdef DISABLE_ASSERTIONS
+	return;
+#endif
 	if (condition) {
 		return;
 	}