duckdb 1.1.4-dev2.0 → 1.1.4-dev9.0

package/src/duckdb/src/catalog/duck_catalog.cpp

@@ -8,15 +8,13 @@
 #include "duckdb/function/built_in_functions.hpp"
 #include "duckdb/main/attached_database.hpp"
 #include "duckdb/transaction/duck_transaction_manager.hpp"
-#ifndef DISABLE_CORE_FUNCTIONS_EXTENSION
-#include "duckdb/core_functions/core_functions.hpp"
-#endif
+#include "duckdb/function/function_list.hpp"
 
 namespace duckdb {
 
 DuckCatalog::DuckCatalog(AttachedDatabase &db)
     : Catalog(db), dependency_manager(make_uniq<DependencyManager>(*this)),
-      schemas(make_uniq<CatalogSet>(*this, make_uniq<DefaultSchemaGenerator>(*this))) {
+      schemas(make_uniq<CatalogSet>(*this, IsSystemCatalog() ? make_uniq<DefaultSchemaGenerator>(*this) : nullptr)) {
 }
 
 DuckCatalog::~DuckCatalog() {
@@ -35,13 +33,11 @@ void DuckCatalog::Initialize(bool load_builtin) {
 	CreateSchema(data, info);
 
 	if (load_builtin) {
-		// initialize default functions
 		BuiltinFunctions builtin(data, *this);
 		builtin.Initialize();
 
-#ifndef DISABLE_CORE_FUNCTIONS_EXTENSION
-		CoreFunctions::RegisterFunctions(*this, data);
-#endif
+		// initialize default functions
+		FunctionList::RegisterFunctions(*this, data);
 	}
 
 	Verify();
@@ -51,17 +47,25 @@ bool DuckCatalog::IsDuckCatalog() {
 	return true;
 }
 
+optional_ptr<DependencyManager> DuckCatalog::GetDependencyManager() {
+	return dependency_manager.get();
+}
+
 //===--------------------------------------------------------------------===//
 // Schema
 //===--------------------------------------------------------------------===//
 optional_ptr<CatalogEntry> DuckCatalog::CreateSchemaInternal(CatalogTransaction transaction, CreateSchemaInfo &info) {
 	LogicalDependencyList dependencies;
+
+	if (!info.internal && DefaultSchemaGenerator::IsDefaultSchema(info.schema)) {
+		return nullptr;
+	}
 	auto entry = make_uniq<DuckSchemaEntry>(*this, info);
 	auto result = entry.get();
 	if (!schemas->CreateEntry(transaction, info.schema, std::move(entry), dependencies)) {
 		return nullptr;
 	}
-	return (CatalogEntry *)result;
+	return result;
 }
 
 optional_ptr<CatalogEntry> DuckCatalog::CreateSchema(CatalogTransaction transaction, CreateSchemaInfo &info) {
@@ -115,6 +119,10 @@ void DuckCatalog::ScanSchemas(std::function<void(SchemaCatalogEntry &)> callback
 	schemas->Scan([&](CatalogEntry &entry) { callback(entry.Cast<SchemaCatalogEntry>()); });
 }
 
+CatalogSet &DuckCatalog::GetSchemaCatalogSet() {
+	return *schemas;
+}
+
 optional_ptr<SchemaCatalogEntry> DuckCatalog::GetSchema(CatalogTransaction transaction, const string &schema_name,
                                                         OnEntryNotFound if_not_found, QueryErrorContext error_context) {
 	D_ASSERT(!schema_name.empty());

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

@@ -172,6 +172,15 @@ AdbcStatusCode StatementSetSubstraitPlan(struct AdbcStatement *statement, const
 		return ADBC_STATUS_INVALID_ARGUMENT;
 	}
 	auto wrapper = static_cast<DuckDBAdbcStatementWrapper *>(statement->private_data);
+	if (wrapper->ingestion_stream.release) {
+		// Release any resources currently held by the ingestion stream before we overwrite it
+		wrapper->ingestion_stream.release(&wrapper->ingestion_stream);
+		wrapper->ingestion_stream.release = nullptr;
+	}
+	if (wrapper->statement) {
+		duckdb_destroy_prepare(&wrapper->statement);
+		wrapper->statement = nullptr;
+	}
 	wrapper->substrait_plan = static_cast<uint8_t *>(malloc(sizeof(uint8_t) * length));
 	wrapper->plan_length = length;
 	memcpy(wrapper->substrait_plan, plan, length);
@@ -912,6 +921,15 @@ AdbcStatusCode StatementSetSqlQuery(struct AdbcStatement *statement, const char
 	}
 
 	auto wrapper = static_cast<DuckDBAdbcStatementWrapper *>(statement->private_data);
+	if (wrapper->ingestion_stream.release) {
+		// Release any resources currently held by the ingestion stream before we overwrite it
+		wrapper->ingestion_stream.release(&wrapper->ingestion_stream);
+		wrapper->ingestion_stream.release = nullptr;
+	}
+	if (wrapper->statement) {
+		duckdb_destroy_prepare(&wrapper->statement);
+		wrapper->statement = nullptr;
+	}
 	auto res = duckdb_prepare(wrapper->connection, query, &wrapper->statement);
 	auto error_msg = duckdb_prepare_error(wrapper->statement);
 	return CheckResult(res, error, error_msg);

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

@@ -5,6 +5,7 @@
 #include "duckdb/common/exception.hpp"
 #include "duckdb/common/helper.hpp"
 #include "duckdb/common/numeric_utils.hpp"
+#include "duckdb/common/operator/cast_operators.hpp"
 #include "duckdb/common/types/timestamp.hpp"
 
 #include <cstdint>
@@ -242,7 +243,8 @@ static void MallocTrim(idx_t pad) {
 	static atomic<int64_t> LAST_TRIM_TIMESTAMP_MS {0};
 
 	int64_t last_trim_timestamp_ms = LAST_TRIM_TIMESTAMP_MS.load();
-	int64_t current_timestamp_ms = Timestamp::GetEpochMs(Timestamp::GetCurrentTimestamp());
+	auto current_ts = Timestamp::GetCurrentTimestamp();
+	auto current_timestamp_ms = Cast::Operation<timestamp_t, timestamp_ms_t>(current_ts).value;
 
 	if (current_timestamp_ms - last_trim_timestamp_ms < TRIM_INTERVAL_MS) {
 		return; // We trimmed less than TRIM_INTERVAL_MS ago

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

@@ -7,6 +7,7 @@
 #include "duckdb/function/table/arrow.hpp"
 #include "duckdb/common/arrow/appender/append_data.hpp"
 #include "duckdb/common/arrow/appender/list.hpp"
+#include "duckdb/function/table/arrow/arrow_duck_schema.hpp"
 
 namespace duckdb {
 
@@ -14,10 +15,17 @@ namespace duckdb {
 // ArrowAppender
 //===--------------------------------------------------------------------===//
 
-ArrowAppender::ArrowAppender(vector<LogicalType> types_p, const idx_t initial_capacity, ClientProperties options)
-    : types(std::move(types_p)) {
-	for (auto &type : types) {
-		auto entry = InitializeChild(type, initial_capacity, options);
+ArrowAppender::ArrowAppender(vector<LogicalType> types_p, const idx_t initial_capacity, ClientProperties options,
+                             unordered_map<idx_t, const shared_ptr<ArrowTypeExtensionData>> extension_type_cast)
+    : types(std::move(types_p)), options(options) {
+	for (idx_t i = 0; i < types.size(); i++) {
+		unique_ptr<ArrowAppendData> entry;
+		bool bitshift_boolean = types[i].id() == LogicalTypeId::BOOLEAN && !options.arrow_lossless_conversion;
+		if (extension_type_cast.find(i) != extension_type_cast.end() && !bitshift_boolean) {
+			entry = InitializeChild(types[i], initial_capacity, options, extension_type_cast[i]);
+		} else {
+			entry = InitializeChild(types[i], initial_capacity, options);
+		}
 		root_data.push_back(std::move(entry));
 	}
 }
@@ -26,11 +34,18 @@ ArrowAppender::~ArrowAppender() {
 }
 
 //! Append a data chunk to the underlying arrow array
-void ArrowAppender::Append(DataChunk &input, idx_t from, idx_t to, idx_t input_size) {
+void ArrowAppender::Append(DataChunk &input, const idx_t from, const idx_t to, const idx_t input_size) {
 	D_ASSERT(types == input.GetTypes());
 	D_ASSERT(to >= from);
 	for (idx_t i = 0; i < input.ColumnCount(); i++) {
-		root_data[i]->append_vector(*root_data[i], input.data[i], from, to, input_size);
+		if (root_data[i]->extension_data && root_data[i]->extension_data->duckdb_to_arrow) {
+			Vector input_data(root_data[i]->extension_data->GetInternalType());
+			root_data[i]->extension_data->duckdb_to_arrow(*options.client_context, input.data[i], input_data,
+			                                              input_size);
+			root_data[i]->append_vector(*root_data[i], input_data, from, to, input_size);
+		} else {
+			root_data[i]->append_vector(*root_data[i], input.data[i], from, to, input_size);
+		}
 	}
 	row_count += to - from;
 }
@@ -285,13 +300,19 @@ static void InitializeFunctionPointers(ArrowAppendData &append_data, const Logic
 }
 
 unique_ptr<ArrowAppendData> ArrowAppender::InitializeChild(const LogicalType &type, const idx_t capacity,
-                                                           ClientProperties &options) {
+                                                           ClientProperties &options,
+                                                           const shared_ptr<ArrowTypeExtensionData> &extension_type) {
 	auto result = make_uniq<ArrowAppendData>(options);
-	InitializeFunctionPointers(*result, type);
+	LogicalType array_type = type;
+	if (extension_type) {
+		array_type = extension_type->GetInternalType();
+	}
+	InitializeFunctionPointers(*result, array_type);
+	result->extension_data = extension_type;
 
 	const auto byte_count = (capacity + 7) / 8;
 	result->GetValidityBuffer().reserve(byte_count);
-	result->initialize(*result, type, capacity);
+	result->initialize(*result, array_type, capacity);
 	return result;
 }
 

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

@@ -9,14 +9,17 @@
 #include "duckdb/common/types/vector_cache.hpp"
 #include "duckdb/common/unordered_map.hpp"
 #include "duckdb/common/vector.hpp"
-#include <list>
+#include "duckdb/main/config.hpp"
+
 #include "duckdb/common/arrow/arrow_appender.hpp"
 #include "duckdb/common/arrow/schema_metadata.hpp"
-
+#include "duckdb/main/client_context.hpp"
 namespace duckdb {
 
-void ArrowConverter::ToArrowArray(DataChunk &input, ArrowArray *out_array, ClientProperties options) {
-	ArrowAppender appender(input.GetTypes(), input.size(), std::move(options));
+void ArrowConverter::ToArrowArray(
+    DataChunk &input, ArrowArray *out_array, ClientProperties options,
+    const unordered_map<idx_t, const shared_ptr<ArrowTypeExtensionData>> &extension_type_cast) {
+	ArrowAppender appender(input.GetTypes(), input.size(), std::move(options), extension_type_cast);
 	appender.Append(input, 0, input.size(), input.size());
 	*out_array = appender.Finalize();
 }
@@ -30,24 +33,6 @@ unsafe_unique_array<char> AddName(const string &name) {
 	return name_ptr;
 }
 
-//===--------------------------------------------------------------------===//
-// Arrow Schema
-//===--------------------------------------------------------------------===//
-struct DuckDBArrowSchemaHolder {
-	// unused in children
-	vector<ArrowSchema> children;
-	// unused in children
-	vector<ArrowSchema *> children_ptrs;
-	//! used for nested structures
-	std::list<vector<ArrowSchema>> nested_children;
-	std::list<vector<ArrowSchema *>> nested_children_ptr;
-	//! This holds strings created to represent decimal types
-	vector<unsafe_unique_array<char>> owned_type_names;
-	vector<unsafe_unique_array<char>> owned_column_names;
-	//! This holds any values created for metadata info
-	vector<unsafe_unique_array<char>> metadata_info;
-};
-
 static void ReleaseDuckDBArrowSchema(ArrowSchema *schema) {
 	if (!schema || !schema->release) {
 		return;
@@ -74,10 +59,10 @@ void InitializeChild(ArrowSchema &child, DuckDBArrowSchemaHolder &root_holder, c
 }
 
 void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, const LogicalType &type,
-                    const ClientProperties &options);
+                    ClientProperties &options, ClientContext &context);
 
 void SetArrowMapFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, const LogicalType &type,
-                       const ClientProperties &options) {
+                       ClientProperties &options, ClientContext &context) {
 	child.format = "+m";
 	//! Map has one child which is a struct
 	child.n_children = 1;
@@ -88,14 +73,38 @@ void SetArrowMapFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child,
 	InitializeChild(root_holder.nested_children.back()[0], root_holder);
 	child.children = &root_holder.nested_children_ptr.back()[0];
 	child.children[0]->name = "entries";
-	SetArrowFormat(root_holder, **child.children, ListType::GetChildType(type), options);
+	SetArrowFormat(root_holder, **child.children, ListType::GetChildType(type), options, context);
+}
+
+bool SetArrowExtension(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, const LogicalType &type,
+                       ClientContext &context) {
+	auto &config = DBConfig::GetConfig(context);
+	if (config.HasArrowExtension(type.id())) {
+		auto arrow_extension = config.GetArrowExtension(type.id());
+		arrow_extension.PopulateArrowSchema(root_holder, child, type, context, arrow_extension);
+		return true;
+	}
+	return false;
 }
 
 void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, const LogicalType &type,
-                    const ClientProperties &options) {
+                    ClientProperties &options, ClientContext &context) {
+	if (type.HasAlias()) {
+		// If it is a json type, we only export it as json if arrow_lossless_conversion = True
+		if (!(type.IsJSONType() && !options.arrow_lossless_conversion)) {
+			// If the type has an alias, we check if it is an Arrow-Type extension
+			if (SetArrowExtension(root_holder, child, type, context)) {
+				return;
+			}
+		}
+	}
 	switch (type.id()) {
 	case LogicalTypeId::BOOLEAN:
-		child.format = "b";
+		if (options.arrow_lossless_conversion) {
+			SetArrowExtension(root_holder, child, type, context);
+		} else {
+			child.format = "b";
+		}
 		break;
 	case LogicalTypeId::TINYINT:
 		child.format = "c";
@@ -126,43 +135,18 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		break;
 	case LogicalTypeId::HUGEINT: {
 		if (options.arrow_lossless_conversion) {
-			child.format = "w:16";
-			auto schema_metadata = ArrowSchemaMetadata::MetadataFromName("duckdb.hugeint");
-			root_holder.metadata_info.emplace_back(schema_metadata.SerializeMetadata());
-			child.metadata = root_holder.metadata_info.back().get();
+			SetArrowExtension(root_holder, child, type, context);
 		} else {
 			child.format = "d:38,0";
 		}
 		break;
 	}
-	case LogicalTypeId::UHUGEINT: {
-		child.format = "w:16";
-		auto schema_metadata = ArrowSchemaMetadata::MetadataFromName("duckdb.uhugeint");
-		root_holder.metadata_info.emplace_back(schema_metadata.SerializeMetadata());
-		child.metadata = root_holder.metadata_info.back().get();
-		break;
-	}
-	case LogicalTypeId::VARINT: {
-		if (options.arrow_offset_size == ArrowOffsetSize::LARGE) {
-			child.format = "Z";
-		} else {
-			child.format = "z";
-		}
-		auto schema_metadata = ArrowSchemaMetadata::MetadataFromName("duckdb.varint");
-		root_holder.metadata_info.emplace_back(schema_metadata.SerializeMetadata());
-		child.metadata = root_holder.metadata_info.back().get();
-		break;
-	}
 	case LogicalTypeId::DOUBLE:
 		child.format = "g";
 		break;
 	case LogicalTypeId::UUID: {
 		if (options.arrow_lossless_conversion) {
-			// This is a canonical extension, hence needs the "arrow." prefix
-			child.format = "w:16";
-			auto schema_metadata = ArrowSchemaMetadata::MetadataFromName("arrow.uuid");
-			root_holder.metadata_info.emplace_back(schema_metadata.SerializeMetadata());
-			child.metadata = root_holder.metadata_info.back().get();
+			SetArrowExtension(root_holder, child, type, context);
 		} else {
 			if (options.produce_arrow_string_view) {
 				child.format = "vu";
@@ -177,11 +161,6 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		break;
 	}
 	case LogicalTypeId::VARCHAR:
-		if (type.IsJSONType() && options.arrow_lossless_conversion) {
-			auto schema_metadata = ArrowSchemaMetadata::MetadataFromName("arrow.json");
-			root_holder.metadata_info.emplace_back(schema_metadata.SerializeMetadata());
-			child.metadata = root_holder.metadata_info.back().get();
-		}
 		if (options.produce_arrow_string_view) {
 			child.format = "vu";
 		} else {
@@ -197,10 +176,7 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		break;
 	case LogicalTypeId::TIME_TZ: {
 		if (options.arrow_lossless_conversion) {
-			child.format = "w:8";
-			auto schema_metadata = ArrowSchemaMetadata::MetadataFromName("duckdb.time_tz");
-			root_holder.metadata_info.emplace_back(schema_metadata.SerializeMetadata());
-			child.metadata = root_holder.metadata_info.back().get();
+			SetArrowExtension(root_holder, child, type, context);
 		} else {
 			child.format = "ttu";
 		}
@@ -250,16 +226,16 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		}
 		break;
 	case LogicalTypeId::BIT: {
-		if (options.arrow_offset_size == ArrowOffsetSize::LARGE) {
-			child.format = "Z";
-		} else {
-			child.format = "z";
-		}
 		if (options.arrow_lossless_conversion) {
-			auto schema_metadata = ArrowSchemaMetadata::MetadataFromName("duckdb.bit");
-			root_holder.metadata_info.emplace_back(schema_metadata.SerializeMetadata());
-			child.metadata = root_holder.metadata_info.back().get();
+			SetArrowExtension(root_holder, child, type, context);
+		} else {
+			if (options.arrow_offset_size == ArrowOffsetSize::LARGE) {
+				child.format = "Z";
+			} else {
+				child.format = "z";
+			}
 		}
+
 		break;
 	}
 	case LogicalTypeId::LIST: {
@@ -284,7 +260,7 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		InitializeChild(root_holder.nested_children.back()[0], root_holder);
 		child.children = &root_holder.nested_children_ptr.back()[0];
 		child.children[0]->name = "l";
-		SetArrowFormat(root_holder, **child.children, ListType::GetChildType(type), options);
+		SetArrowFormat(root_holder, **child.children, ListType::GetChildType(type), options, context);
 		break;
 	}
 	case LogicalTypeId::STRUCT: {
@@ -306,7 +282,7 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 			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);
+			SetArrowFormat(root_holder, *child.children[type_idx], child_types[type_idx].second, options, context);
 		}
 		break;
 	}
@@ -324,11 +300,11 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		root_holder.nested_children_ptr.back().push_back(&root_holder.nested_children.back()[0]);
 		InitializeChild(root_holder.nested_children.back()[0], root_holder);
 		child.children = &root_holder.nested_children_ptr.back()[0];
-		SetArrowFormat(root_holder, **child.children, child_type, options);
+		SetArrowFormat(root_holder, **child.children, child_type, options, context);
 		break;
 	}
 	case LogicalTypeId::MAP: {
-		SetArrowMapFormat(root_holder, child, type, options);
+		SetArrowMapFormat(root_holder, child, type, options, context);
 		break;
 	}
 	case LogicalTypeId::UNION: {
@@ -351,7 +327,7 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 			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);
+			SetArrowFormat(root_holder, *child.children[type_idx], child_types[type_idx].second, options, context);
 
 			format += to_string(type_idx) + ",";
 		}
@@ -387,17 +363,22 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		child.dictionary->format = "u";
 		break;
 	}
-	default:
-		throw NotImplementedException("Unsupported Arrow type " + type.ToString());
+	default: {
+		// It is possible we can export this type as a registered extension
+		auto success = SetArrowExtension(root_holder, child, type, context);
+		if (!success) {
+			throw NotImplementedException("Unsupported Arrow type %s", type.ToString());
+		}
+	}
 	}
 }
 
 void ArrowConverter::ToArrowSchema(ArrowSchema *out_schema, const vector<LogicalType> &types,
-                                   const vector<string> &names, const ClientProperties &options) {
+                                   const vector<string> &names, ClientProperties &options) {
 	D_ASSERT(out_schema);
 	D_ASSERT(types.size() == names.size());
-	idx_t column_count = types.size();
-	// Allocate as unique_ptr first to cleanup properly on error
+	const idx_t column_count = types.size();
+	// Allocate as unique_ptr first to clean-up properly on error
 	auto root_holder = make_uniq<DuckDBArrowSchemaHolder>();
 
 	// Allocate the children
@@ -421,7 +402,7 @@ void ArrowConverter::ToArrowSchema(ArrowSchema *out_schema, const vector<Logical
 		root_holder->owned_column_names.push_back(AddName(names[col_idx]));
 		auto &child = root_holder->children[col_idx];
 		InitializeChild(child, *root_holder, names[col_idx]);
-		SetArrowFormat(*root_holder, child, types[col_idx], options);
+		SetArrowFormat(*root_holder, child, types[col_idx], options, *options.client_context);
 	}
 
 	// Release ownership to caller

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

@@ -21,8 +21,9 @@ void ArrowBatchTask::ProduceRecordBatches() {
 	for (auto &index : record_batch_indices) {
 		auto &array = arrays[index];
 		D_ASSERT(array);
-		idx_t count;
-		count = ArrowUtil::FetchChunk(scan_state, arrow_options, batch_size, &array->arrow_array);
+		const idx_t count = ArrowUtil::FetchChunk(
+		    scan_state, arrow_options, batch_size, &array->arrow_array,
+		    ArrowTypeExtensionData::GetExtensionTypes(event->GetClientContext(), scan_state.Types()));
 		(void)count;
 		D_ASSERT(count != 0);
 	}
@@ -87,7 +88,7 @@ void ArrowMergeEvent::Schedule() {
 		idx_t tuples_for_task = 0;
 		idx_t start_index = transformer.GetIndex();
 		idx_t end_index = start_index;
-		while (tuples_for_task < Storage::ROW_GROUP_SIZE) {
+		while (tuples_for_task < DEFAULT_ROW_GROUP_SIZE) {
 			idx_t batch_size;
 			if (!transformer.TryGetNextBatchSize(batch_size)) {
 				finished = true;