Update vendored DuckDB sources to 44b706b2b7

Author: duckdblabs-bot

src/duckdb/extension/icu/third_party/icu/common/unicode/ucnv.h

@@ -0,0 +1,11 @@
+/**
+ * Converter option for EBCDIC SBCS or mixed-SBCS/DBCS (stateful) codepages.
+ * Swaps Unicode mappings for EBCDIC LF and NL codes, as used on
+ * S/390 (z/OS) Unix System Services (Open Edition).
+ * For example, ucnv_open("ibm-1047,swaplfnl", &errorCode);
+ * See convrtrs.txt.
+ *
+ * @see ucnv_open
+ * @stable ICU 2.4
+ */
+#define UCNV_SWAP_LFNL_OPTION_STRING ",swaplfnl"

src/duckdb/extension/parquet/column_reader.cpp

@@ -895,9 +895,6 @@ unique_ptr<ColumnReader> ColumnReader::CreateReader(ParquetReader &reader, const
 		default:
 			throw NotImplementedException("Unrecognized Parquet type for Decimal");
 		}
-	case LogicalTypeId::GEOMETRY:
-		// TODO: Make GeometryColumnReader
-		return make_uniq<StringColumnReader>(reader, schema);
 	case LogicalTypeId::UUID:
 		return make_uniq<UUIDColumnReader>(reader, schema);
 	case LogicalTypeId::INTERVAL:

src/duckdb/extension/parquet/include/parquet_column_schema.hpp

@@ -15,7 +15,7 @@ namespace duckdb {
 using duckdb_parquet::FileMetaData;
 struct ParquetOptions;
 
-enum class ParquetColumnSchemaType { COLUMN, FILE_ROW_NUMBER, EXPRESSION, VARIANT };
+enum class ParquetColumnSchemaType { COLUMN, FILE_ROW_NUMBER, EXPRESSION, VARIANT, GEOMETRY };
 
 enum class ParquetExtraTypeInfo {
 	NONE,
@@ -35,7 +35,7 @@ struct ParquetColumnSchema {
 	                    ParquetColumnSchemaType schema_type = ParquetColumnSchemaType::COLUMN);
 	ParquetColumnSchema(string name, LogicalType type, idx_t max_define, idx_t max_repeat, idx_t schema_index,
 	                    idx_t column_index, ParquetColumnSchemaType schema_type = ParquetColumnSchemaType::COLUMN);
-	ParquetColumnSchema(ParquetColumnSchema parent, LogicalType result_type, ParquetColumnSchemaType schema_type);
+	ParquetColumnSchema(ParquetColumnSchema child, LogicalType result_type, ParquetColumnSchemaType schema_type);
 
 	ParquetColumnSchemaType schema_type;
 	string name;

src/duckdb/extension/parquet/include/parquet_geometry.hpp

@@ -21,7 +21,11 @@ struct ParquetColumnSchema;
 class ParquetReader;
 class ColumnReader;
 class ClientContext;
-class ExpressionExecutor;
+
+struct GeometryColumnReader {
+	static unique_ptr<ColumnReader> Create(ParquetReader &reader, const ParquetColumnSchema &schema,
+	                                       ClientContext &context);
+};
 
 enum class GeoParquetColumnEncoding : uint8_t {
 	WKB = 1,
@@ -90,9 +94,6 @@ class GeoParquetFileMetadata {
 	                                                  const ClientContext &context);
 	const unordered_map<string, GeoParquetColumnMetadata> &GetColumnMeta() const;
 
-	static unique_ptr<ColumnReader> CreateColumnReader(ParquetReader &reader, const ParquetColumnSchema &schema,
-	                                                   ClientContext &context);
-
 	bool IsGeometryColumn(const string &column_name) const;
 
 	static bool IsGeoParquetConversionEnabled(const ClientContext &context);

src/duckdb/extension/parquet/parquet_geometry.cpp

@@ -5,12 +5,14 @@
 #include "duckdb/catalog/catalog_entry/scalar_function_catalog_entry.hpp"
 #include "duckdb/execution/expression_executor.hpp"
 #include "duckdb/function/scalar_function.hpp"
+#include "duckdb/function/scalar/geometry_functions.hpp"
 #include "duckdb/planner/expression/bound_function_expression.hpp"
 #include "duckdb/planner/expression/bound_reference_expression.hpp"
 #include "duckdb/main/extension_helper.hpp"
 #include "reader/expression_column_reader.hpp"
 #include "parquet_reader.hpp"
 #include "yyjson.hpp"
+#include "reader/string_column_reader.hpp"
 
 namespace duckdb {
 
@@ -283,34 +285,24 @@ const unordered_map<string, GeoParquetColumnMetadata> &GeoParquetFileMetadata::G
 	return geometry_columns;
 }
 
-unique_ptr<ColumnReader> GeoParquetFileMetadata::CreateColumnReader(ParquetReader &reader,
-                                                                    const ParquetColumnSchema &schema,
-                                                                    ClientContext &context) {
-	// Get the catalog
-	auto &catalog = Catalog::GetSystemCatalog(context);
-
-	// WKB encoding
-	if (schema.children[0].type.id() == LogicalTypeId::BLOB) {
-		// Look for a conversion function in the catalog
-		auto &conversion_func_set =
-		    catalog.GetEntry<ScalarFunctionCatalogEntry>(context, DEFAULT_SCHEMA, "st_geomfromwkb");
-		auto conversion_func = conversion_func_set.functions.GetFunctionByArguments(context, {LogicalType::BLOB});
-
-		// Create a bound function call expression
-		auto args = vector<unique_ptr<Expression>>();
-		args.push_back(std::move(make_uniq<BoundReferenceExpression>(LogicalType::BLOB, 0)));
-		auto expr = make_uniq<BoundFunctionExpression>(conversion_func.GetReturnType(), conversion_func,
-		                                               std::move(args), nullptr);
-
-		// Create a child reader
-		auto child_reader = ColumnReader::CreateReader(reader, schema.children[0]);
-
-		// Create an expression reader that applies the conversion function to the child reader
-		return make_uniq<ExpressionColumnReader>(context, std::move(child_reader), std::move(expr), schema);
-	}
+unique_ptr<ColumnReader> GeometryColumnReader::Create(ParquetReader &reader, const ParquetColumnSchema &schema,
+                                                      ClientContext &context) {
+	D_ASSERT(schema.type.id() == LogicalTypeId::GEOMETRY);
+	D_ASSERT(schema.children.size() == 1 && schema.children[0].type.id() == LogicalTypeId::BLOB);
+
+	// Make a string reader for the underlying WKB data
+	auto string_reader = make_uniq<StringColumnReader>(reader, schema.children[0]);
+
+	// Wrap the string reader in a geometry reader
+	auto args = vector<unique_ptr<Expression>>();
+	auto ref = make_uniq_base<Expression, BoundReferenceExpression>(LogicalTypeId::BLOB, 0);
+	args.push_back(std::move(ref));
 
-	// Otherwise, unrecognized encoding
-	throw NotImplementedException("Unsupported geometry encoding");
+	// TODO: Pass the actual target type here so we get the CRS information too
+	auto func = StGeomfromwkbFun::GetFunction();
+	func.name = "ST_GeomFromWKB";
+	auto expr = make_uniq_base<Expression, BoundFunctionExpression>(schema.type, func, std::move(args), nullptr);
+	return make_uniq<ExpressionColumnReader>(context, std::move(string_reader), std::move(expr), schema);
 }
 
 } // namespace duckdb

src/duckdb/extension/parquet/parquet_reader.cpp

@@ -225,11 +225,6 @@ LogicalType ParquetReader::DeriveLogicalType(const SchemaElement &s_ele, Parquet
 				return LogicalType::TIME_TZ;
 			}
 			return LogicalType::TIME;
-		} else if (s_ele.logicalType.__isset.GEOMETRY) {
-			// TODO: Set CRS too
-			return LogicalType::GEOMETRY();
-		} else if (s_ele.logicalType.__isset.GEOGRAPHY) {
-			return LogicalType::GEOMETRY();
 		}
 	}
 	if (s_ele.__isset.converted_type) {
@@ -409,6 +404,9 @@ unique_ptr<ColumnReader> ParquetReader::CreateReaderRecursive(ClientContext &con
 	switch (schema.schema_type) {
 	case ParquetColumnSchemaType::FILE_ROW_NUMBER:
 		return make_uniq<RowNumberColumnReader>(*this, schema);
+	case ParquetColumnSchemaType::GEOMETRY: {
+		return GeometryColumnReader::Create(*this, schema, context);
+	}
 	case ParquetColumnSchemaType::COLUMN: {
 		if (schema.children.empty()) {
 			// leaf reader
@@ -486,11 +484,11 @@ ParquetColumnSchema::ParquetColumnSchema(string name_p, LogicalType type_p, idx_
       max_repeat(max_repeat), schema_index(schema_index), column_index(column_index) {
 }
 
-ParquetColumnSchema::ParquetColumnSchema(ParquetColumnSchema parent, LogicalType result_type,
+ParquetColumnSchema::ParquetColumnSchema(ParquetColumnSchema child, LogicalType result_type,
                                          ParquetColumnSchemaType schema_type)
-    : schema_type(schema_type), name(parent.name), type(std::move(result_type)), max_define(parent.max_define),
-      max_repeat(parent.max_repeat), schema_index(parent.schema_index), column_index(parent.column_index) {
-	children.push_back(std::move(parent));
+    : schema_type(schema_type), name(child.name), type(std::move(result_type)), max_define(child.max_define),
+      max_repeat(child.max_repeat), schema_index(child.schema_index), column_index(child.column_index) {
+	children.push_back(std::move(child));
 }
 
 unique_ptr<BaseStatistics> ParquetColumnSchema::Stats(const FileMetaData &file_meta_data,
@@ -517,6 +515,32 @@ unique_ptr<BaseStatistics> ParquetColumnSchema::Stats(const FileMetaData &file_m
 	return ParquetStatisticsUtils::TransformColumnStatistics(*this, columns, parquet_options.can_have_nan);
 }
 
+static bool IsGeometryType(const SchemaElement &s_ele, const ParquetFileMetadataCache &metadata, idx_t depth) {
+	const auto is_blob = s_ele.__isset.type && s_ele.type == Type::BYTE_ARRAY;
+	if (!is_blob) {
+		return false;
+	}
+
+	// TODO: Handle CRS in the future
+	const auto is_native_geom = s_ele.__isset.logicalType && s_ele.logicalType.__isset.GEOMETRY;
+	const auto is_native_geog = s_ele.__isset.logicalType && s_ele.logicalType.__isset.GEOGRAPHY;
+	if (is_native_geom || is_native_geog) {
+		return true;
+	}
+
+	// geoparquet types have to be at the root of the schema, and have to be present in the kv metadata.
+	const auto is_at_root = depth == 1;
+	const auto is_in_gpq_metadata = metadata.geo_metadata && metadata.geo_metadata->IsGeometryColumn(s_ele.name);
+	const auto is_leaf = s_ele.num_children == 0;
+	const auto is_geoparquet_geom = is_at_root && is_in_gpq_metadata && is_leaf;
+
+	if (is_geoparquet_geom) {
+		return true;
+	}
+
+	return false;
+}
+
 ParquetColumnSchema ParquetReader::ParseSchemaRecursive(idx_t depth, idx_t max_define, idx_t max_repeat,
                                                         idx_t &next_schema_idx, idx_t &next_file_idx,
                                                         ClientContext &context) {
@@ -540,16 +564,26 @@ ParquetColumnSchema ParquetReader::ParseSchemaRecursive(idx_t depth, idx_t max_d
 		max_repeat++;
 	}
 
-	// Check for geoparquet spatial types
-	if (depth == 1) {
-		// geoparquet types have to be at the root of the schema, and have to be present in the kv metadata.
-		// geoarrow types, although geometry columns, are structs and have children and are handled below.
-		if (metadata->geo_metadata && metadata->geo_metadata->IsGeometryColumn(s_ele.name) && s_ele.num_children == 0) {
-			auto geom_schema = ParseColumnSchema(s_ele, max_define, max_repeat, this_idx, next_file_idx++);
-			// overwrite the derived type with GEOMETRY
-			geom_schema.type = LogicalType::GEOMETRY();
-			return geom_schema;
-		}
+	// Check for geometry type
+	if (IsGeometryType(s_ele, *metadata, depth)) {
+		// Geometries in both GeoParquet and native parquet are stored as a WKB-encoded BLOB.
+		// Because we don't just want to validate that the WKB encoding is correct, but also transform it into
+		// little-endian if necessary, we cant just make use of the StringColumnReader without heavily modifying it.
+		// Therefore, we create a dedicated GEOMETRY parquet column schema type, which wraps the underlying BLOB column.
+		// This schema type gets instantiated as a ExpressionColumnReader on top of the standard Blob/String reader,
+		// which performs the WKB validation/transformation using the `ST_GeomFromWKB` function of DuckDB.
+		// This enables us to also support other geometry encodings (such as GeoArrow geometries) easier in the future.
+
+		// Inner BLOB schema
+		ParquetColumnSchema blob_schema(max_define, max_repeat, this_idx, next_file_idx++,
+		                                ParquetColumnSchemaType::COLUMN);
+		blob_schema.name = s_ele.name;
+		blob_schema.type = LogicalType::BLOB;
+
+		// Wrap in geometry schema
+		ParquetColumnSchema geom_schema(std::move(blob_schema), LogicalType::GEOMETRY(),
+		                                ParquetColumnSchemaType::GEOMETRY);
+		return geom_schema;
 	}
 
 	if (s_ele.__isset.num_children && s_ele.num_children > 0) { // inner node

src/duckdb/src/common/allocator.cpp

@@ -256,7 +256,7 @@ static void MallocTrim(idx_t pad) {
 		return; // Another thread has updated LAST_TRIM_TIMESTAMP_MS since we loaded it
 	}
 
-	// We succesfully updated LAST_TRIM_TIMESTAMP_MS, we can trim
+	// We successfully updated LAST_TRIM_TIMESTAMP_MS, we can trim
 	malloc_trim(pad);
 #endif
 }

src/duckdb/src/common/arrow/arrow_query_result.cpp

@@ -16,10 +16,7 @@ ArrowQueryResult::ArrowQueryResult(StatementType statement_type, StatementProper
 ArrowQueryResult::ArrowQueryResult(ErrorData error) : QueryResult(QueryResultType::ARROW_RESULT, std::move(error)) {
 }
 
-unique_ptr<DataChunk> ArrowQueryResult::Fetch() {
-	throw NotImplementedException("Can't 'Fetch' from ArrowQueryResult");
-}
-unique_ptr<DataChunk> ArrowQueryResult::FetchRaw() {
+unique_ptr<DataChunk> ArrowQueryResult::FetchInternal() {
 	throw NotImplementedException("Can't 'FetchRaw' from ArrowQueryResult");
 }
 

src/duckdb/src/common/arrow/arrow_type_extension.cpp

@@ -7,6 +7,8 @@
 #include "duckdb/common/arrow/schema_metadata.hpp"
 #include "duckdb/common/types/vector.hpp"
 
+#include "yyjson.hpp"
+
 namespace duckdb {
 
 ArrowTypeExtension::ArrowTypeExtension(string extension_name, string arrow_format,
@@ -365,6 +367,72 @@ struct ArrowBool8 {
 	}
 };
 
+struct ArrowGeometry {
+	static unique_ptr<ArrowType> GetType(const ArrowSchema &schema, const ArrowSchemaMetadata &schema_metadata) {
+		// Validate extension metadata. This metadata also contains a CRS, which we drop
+		// because the GEOMETRY type does not implement a CRS at the type level (yet).
+		const auto extension_metadata = schema_metadata.GetOption(ArrowSchemaMetadata::ARROW_METADATA_KEY);
+		if (!extension_metadata.empty()) {
+			unique_ptr<duckdb_yyjson::yyjson_doc, void (*)(duckdb_yyjson::yyjson_doc *)> doc(
+			    duckdb_yyjson::yyjson_read(extension_metadata.data(), extension_metadata.size(),
+			                               duckdb_yyjson::YYJSON_READ_NOFLAG),
+			    duckdb_yyjson::yyjson_doc_free);
+			if (!doc) {
+				throw SerializationException("Invalid JSON in GeoArrow metadata");
+			}
+
+			duckdb_yyjson::yyjson_val *val = yyjson_doc_get_root(doc.get());
+			if (!yyjson_is_obj(val)) {
+				throw SerializationException("Invalid GeoArrow metadata: not a JSON object");
+			}
+
+			duckdb_yyjson::yyjson_val *edges = yyjson_obj_get(val, "edges");
+			if (edges && yyjson_is_str(edges) && std::strcmp(yyjson_get_str(edges), "planar") != 0) {
+				throw NotImplementedException("Can't import non-planar edges");
+			}
+		}
+
+		const auto format = string(schema.format);
+		if (format == "z") {
+			return make_uniq<ArrowType>(LogicalType::GEOMETRY(),
+			                            make_uniq<ArrowStringInfo>(ArrowVariableSizeType::NORMAL));
+		}
+		if (format == "Z") {
+			return make_uniq<ArrowType>(LogicalType::GEOMETRY(),
+			                            make_uniq<ArrowStringInfo>(ArrowVariableSizeType::SUPER_SIZE));
+		}
+		if (format == "vz") {
+			return make_uniq<ArrowType>(LogicalType::GEOMETRY(),
+			                            make_uniq<ArrowStringInfo>(ArrowVariableSizeType::VIEW));
+		}
+		throw InvalidInputException("Arrow extension type \"%s\" not supported for geoarrow.wkb", format.c_str());
+	}
+
+	static void PopulateSchema(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &schema, const LogicalType &type,
+	                           ClientContext &context, const ArrowTypeExtension &extension) {
+		ArrowSchemaMetadata schema_metadata;
+		schema_metadata.AddOption(ArrowSchemaMetadata::ARROW_EXTENSION_NAME, "geoarrow.wkb");
+		schema_metadata.AddOption(ArrowSchemaMetadata::ARROW_METADATA_KEY, "{}");
+		root_holder.metadata_info.emplace_back(schema_metadata.SerializeMetadata());
+		schema.metadata = root_holder.metadata_info.back().get();
+
+		const auto options = context.GetClientProperties();
+		if (options.arrow_offset_size == ArrowOffsetSize::LARGE) {
+			schema.format = "Z";
+		} else {
+			schema.format = "z";
+		}
+	}
+
+	static void ArrowToDuck(ClientContext &, Vector &source, Vector &result, idx_t count) {
+		Geometry::FromBinary(source, result, count, true);
+	}
+
+	static void DuckToArrow(ClientContext &context, Vector &source, Vector &result, idx_t count) {
+		Geometry::ToBinary(source, result, count);
+	}
+};
+
 void ArrowTypeExtensionSet::Initialize(const DBConfig &config) {
 	// Types that are 1:1
 	config.RegisterArrowExtension({"arrow.uuid", "w:16", make_shared_ptr<ArrowTypeExtensionData>(LogicalType::UUID)});
@@ -380,6 +448,11 @@ void ArrowTypeExtensionSet::Initialize(const DBConfig &config) {
 	config.RegisterArrowExtension(
 	    {"DuckDB", "time_tz", "w:8", make_shared_ptr<ArrowTypeExtensionData>(LogicalType::TIME_TZ)});
 
+	config.RegisterArrowExtension(
+	    {"geoarrow.wkb", ArrowGeometry::PopulateSchema, ArrowGeometry::GetType,
+	     make_shared_ptr<ArrowTypeExtensionData>(LogicalType::GEOMETRY(), LogicalType::BLOB, ArrowGeometry::ArrowToDuck,
+	                                             ArrowGeometry::DuckToArrow)});
+
 	// Types that are 1:n
 	config.RegisterArrowExtension({"arrow.json", &ArrowJson::PopulateSchema, &ArrowJson::GetType,
 	                               make_shared_ptr<ArrowTypeExtensionData>(LogicalType::JSON())});

src/duckdb/src/common/encryption_key_manager.cpp

@@ -31,6 +31,8 @@ EncryptionKey::~EncryptionKey() {
 void EncryptionKey::LockEncryptionKey(data_ptr_t key, idx_t key_len) {
 #if defined(_WIN32)
 	VirtualLock(key, key_len);
+#elif defined(__MVS__)
+	__mlockall(_BPX_NONSWAP);
 #else
 	mlock(key, key_len);
 #endif
@@ -40,6 +42,8 @@ void EncryptionKey::UnlockEncryptionKey(data_ptr_t key, idx_t key_len) {
 	memset(key, 0, key_len);
 #if defined(_WIN32)
 	VirtualUnlock(key, key_len);
+#elif defined(__MVS__)
+	__mlockall(_BPX_SWAP);
 #else
 	munlock(key, key_len);
 #endif