duckdb 1.3.3-dev4.0 → 1.3.3

package/src/duckdb/extension/parquet/geo_parquet.cpp

@@ -16,168 +16,155 @@ namespace duckdb {
 
 using namespace duckdb_yyjson; // NOLINT
 
+const char *WKBGeometryTypes::ToString(WKBGeometryType type) {
+	switch (type) {
+	case WKBGeometryType::POINT:
+		return "Point";
+	case WKBGeometryType::LINESTRING:
+		return "LineString";
+	case WKBGeometryType::POLYGON:
+		return "Polygon";
+	case WKBGeometryType::MULTIPOINT:
+		return "MultiPoint";
+	case WKBGeometryType::MULTILINESTRING:
+		return "MultiLineString";
+	case WKBGeometryType::MULTIPOLYGON:
+		return "MultiPolygon";
+	case WKBGeometryType::GEOMETRYCOLLECTION:
+		return "GeometryCollection";
+	case WKBGeometryType::POINT_Z:
+		return "Point Z";
+	case WKBGeometryType::LINESTRING_Z:
+		return "LineString Z";
+	case WKBGeometryType::POLYGON_Z:
+		return "Polygon Z";
+	case WKBGeometryType::MULTIPOINT_Z:
+		return "MultiPoint Z";
+	case WKBGeometryType::MULTILINESTRING_Z:
+		return "MultiLineString Z";
+	case WKBGeometryType::MULTIPOLYGON_Z:
+		return "MultiPolygon Z";
+	case WKBGeometryType::GEOMETRYCOLLECTION_Z:
+		return "GeometryCollection Z";
+	default:
+		throw NotImplementedException("Unsupported geometry type");
+	}
+}
+
 //------------------------------------------------------------------------------
-// WKB stats
+// GeoParquetColumnMetadataWriter
 //------------------------------------------------------------------------------
-namespace {
+GeoParquetColumnMetadataWriter::GeoParquetColumnMetadataWriter(ClientContext &context) {
+	executor = make_uniq<ExpressionExecutor>(context);
 
-class BinaryReader {
-public:
-	const char *beg;
-	const char *end;
-	const char *ptr;
+	auto &catalog = Catalog::GetSystemCatalog(context);
 
-	BinaryReader(const char *beg, uint32_t len) : beg(beg), end(beg + len), ptr(beg) {
-	}
+	// These functions are required to extract the geometry type, ZM flag and bounding box from a WKB blob
+	auto &type_func_set = catalog.GetEntry<ScalarFunctionCatalogEntry>(context, DEFAULT_SCHEMA, "st_geometrytype");
+	auto &flag_func_set = catalog.GetEntry<ScalarFunctionCatalogEntry>(context, DEFAULT_SCHEMA, "st_zmflag");
+	auto &bbox_func_set = catalog.GetEntry<ScalarFunctionCatalogEntry>(context, DEFAULT_SCHEMA, "st_extent");
 
-	template <class T>
-	T Read() {
-		if (ptr + sizeof(T) > end) {
-			throw InvalidInputException("Unexpected end of WKB data");
-		}
-		T val;
-		memcpy(&val, ptr, sizeof(T));
-		ptr += sizeof(T);
-		return val;
-	}
+	auto wkb_type = LogicalType(LogicalTypeId::BLOB);
+	wkb_type.SetAlias("WKB_BLOB");
 
-	void Skip(idx_t len) {
-		if (ptr + len > end) {
-			throw InvalidInputException("Unexpected end of WKB data");
-		}
-		ptr += len;
-	}
+	auto type_func = type_func_set.functions.GetFunctionByArguments(context, {wkb_type});
+	auto flag_func = flag_func_set.functions.GetFunctionByArguments(context, {wkb_type});
+	auto bbox_func = bbox_func_set.functions.GetFunctionByArguments(context, {wkb_type});
 
-	const char *Reserve(idx_t len) {
-		if (ptr + len > end) {
-			throw InvalidInputException("Unexpected end of WKB data");
-		}
-		auto ret = ptr;
-		ptr += len;
-		return ret;
-	}
+	auto type_type = LogicalType::UTINYINT;
+	auto flag_type = flag_func.return_type;
+	auto bbox_type = bbox_func.return_type;
 
-	bool IsAtEnd() const {
-		return ptr >= end;
-	}
-};
-
-} // namespace
-
-static void UpdateBoundsFromVertexArray(GeometryExtent &bbox, uint32_t flag, const char *vert_array,
-                                        uint32_t vert_count) {
-	switch (flag) {
-	case 0: { // XY
-		constexpr auto vert_width = sizeof(double) * 2;
-		for (uint32_t vert_idx = 0; vert_idx < vert_count; vert_idx++) {
-			double vert[2];
-			memcpy(vert, vert_array + vert_idx * vert_width, vert_width);
-			bbox.ExtendX(vert[0]);
-			bbox.ExtendY(vert[1]);
-		}
-	} break;
-	case 1: { // XYZ
-		constexpr auto vert_width = sizeof(double) * 3;
-		for (uint32_t vert_idx = 0; vert_idx < vert_count; vert_idx++) {
-			double vert[3];
-			memcpy(vert, vert_array + vert_idx * vert_width, vert_width);
-			bbox.ExtendX(vert[0]);
-			bbox.ExtendY(vert[1]);
-			bbox.ExtendZ(vert[2]);
-		}
-	} break;
-	case 2: { // XYM
-		constexpr auto vert_width = sizeof(double) * 3;
-		for (uint32_t vert_idx = 0; vert_idx < vert_count; vert_idx++) {
-			double vert[3];
-			memcpy(vert, vert_array + vert_idx * vert_width, vert_width);
-			bbox.ExtendX(vert[0]);
-			bbox.ExtendY(vert[1]);
-			bbox.ExtendM(vert[2]);
-		}
-	} break;
-	case 3: { // XYZM
-		constexpr auto vert_width = sizeof(double) * 4;
-		for (uint32_t vert_idx = 0; vert_idx < vert_count; vert_idx++) {
-			double vert[4];
-			memcpy(vert, vert_array + vert_idx * vert_width, vert_width);
-			bbox.ExtendX(vert[0]);
-			bbox.ExtendY(vert[1]);
-			bbox.ExtendZ(vert[2]);
-			bbox.ExtendM(vert[3]);
-		}
-	} break;
-	default:
-		break;
-	}
+	vector<unique_ptr<Expression>> type_args;
+	type_args.push_back(make_uniq<BoundReferenceExpression>(wkb_type, 0));
+
+	vector<unique_ptr<Expression>> flag_args;
+	flag_args.push_back(make_uniq<BoundReferenceExpression>(wkb_type, 0));
+
+	vector<unique_ptr<Expression>> bbox_args;
+	bbox_args.push_back(make_uniq<BoundReferenceExpression>(wkb_type, 0));
+
+	type_expr = make_uniq<BoundFunctionExpression>(type_type, type_func, std::move(type_args), nullptr);
+	flag_expr = make_uniq<BoundFunctionExpression>(flag_type, flag_func, std::move(flag_args), nullptr);
+	bbox_expr = make_uniq<BoundFunctionExpression>(bbox_type, bbox_func, std::move(bbox_args), nullptr);
+
+	// Add the expressions to the executor
+	executor->AddExpression(*type_expr);
+	executor->AddExpression(*flag_expr);
+	executor->AddExpression(*bbox_expr);
+
+	// Initialize the input and result chunks
+	// The input chunk should be empty, as we always reference the input vector
+	input_chunk.InitializeEmpty({wkb_type});
+	result_chunk.Initialize(context, {type_type, flag_type, bbox_type});
 }
 
-void GeometryStats::Update(const string_t &wkb) {
-	BinaryReader reader(wkb.GetData(), wkb.GetSize());
-
-	bool first_geom = true;
-	while (!reader.IsAtEnd()) {
-		reader.Read<uint8_t>(); // byte order
-		auto type = reader.Read<uint32_t>();
-		auto kind = type % 1000;
-		auto flag = type / 1000;
-		const auto hasz = (flag & 0x01) != 0;
-		const auto hasm = (flag & 0x02) != 0;
-
-		if (first_geom) {
-			// Only add the top-level geometry type
-			types.Add(type);
-			first_geom = false;
+void GeoParquetColumnMetadataWriter::Update(GeoParquetColumnMetadata &meta, Vector &vector, idx_t count) {
+	input_chunk.Reset();
+	result_chunk.Reset();
+
+	// Reference the vector
+	input_chunk.data[0].Reference(vector);
+	input_chunk.SetCardinality(count);
+
+	// Execute the expression
+	executor->Execute(input_chunk, result_chunk);
+
+	// The first column is the geometry type
+	// The second column is the zm flag
+	// The third column is the bounding box
+
+	UnifiedVectorFormat type_format;
+	UnifiedVectorFormat flag_format;
+	UnifiedVectorFormat bbox_format;
+
+	result_chunk.data[0].ToUnifiedFormat(count, type_format);
+	result_chunk.data[1].ToUnifiedFormat(count, flag_format);
+	result_chunk.data[2].ToUnifiedFormat(count, bbox_format);
+
+	const auto &bbox_components = StructVector::GetEntries(result_chunk.data[2]);
+	D_ASSERT(bbox_components.size() == 4);
+
+	UnifiedVectorFormat xmin_format;
+	UnifiedVectorFormat ymin_format;
+	UnifiedVectorFormat xmax_format;
+	UnifiedVectorFormat ymax_format;
+
+	bbox_components[0]->ToUnifiedFormat(count, xmin_format);
+	bbox_components[1]->ToUnifiedFormat(count, ymin_format);
+	bbox_components[2]->ToUnifiedFormat(count, xmax_format);
+	bbox_components[3]->ToUnifiedFormat(count, ymax_format);
+
+	for (idx_t in_idx = 0; in_idx < count; in_idx++) {
+		const auto type_idx = type_format.sel->get_index(in_idx);
+		const auto flag_idx = flag_format.sel->get_index(in_idx);
+		const auto bbox_idx = bbox_format.sel->get_index(in_idx);
+
+		const auto type_valid = type_format.validity.RowIsValid(type_idx);
+		const auto flag_valid = flag_format.validity.RowIsValid(flag_idx);
+		const auto bbox_valid = bbox_format.validity.RowIsValid(bbox_idx);
+
+		if (!type_valid || !flag_valid || !bbox_valid) {
+			continue;
 		}
 
-		const auto vert_width = sizeof(double) * (2 + (hasz ? 1 : 0) + (hasm ? 1 : 0));
-
-		switch (kind) {
-		case 1: { // POINT
-
-			// Point are special in that they are considered "empty" if they are all-nan
-			const auto vert_array = reader.Reserve(vert_width);
-			const auto dims_count = 2 + (hasz ? 1 : 0) + (hasm ? 1 : 0);
-			double vert_point[4] = {0, 0, 0, 0};
-
-			memcpy(vert_point, vert_array, vert_width);
-
-			for (auto dim_idx = 0; dim_idx < dims_count; dim_idx++) {
-				if (!std::isnan(vert_point[dim_idx])) {
-					bbox.ExtendX(vert_point[0]);
-					bbox.ExtendY(vert_point[1]);
-					if (hasz && hasm) {
-						bbox.ExtendZ(vert_point[2]);
-						bbox.ExtendM(vert_point[3]);
-					} else if (hasz) {
-						bbox.ExtendZ(vert_point[2]);
-					} else if (hasm) {
-						bbox.ExtendM(vert_point[2]);
-					}
-					break;
-				}
-			}
-		} break;
-		case 2: { // LINESTRING
-			const auto vert_count = reader.Read<uint32_t>();
-			const auto vert_array = reader.Reserve(vert_count * vert_width);
-			UpdateBoundsFromVertexArray(bbox, flag, vert_array, vert_count);
-		} break;
-		case 3: { // POLYGON
-			const auto ring_count = reader.Read<uint32_t>();
-			for (uint32_t ring_idx = 0; ring_idx < ring_count; ring_idx++) {
-				const auto vert_count = reader.Read<uint32_t>();
-				const auto vert_array = reader.Reserve(vert_count * vert_width);
-				UpdateBoundsFromVertexArray(bbox, flag, vert_array, vert_count);
-			}
-		} break;
-		case 4:   // MULTIPOINT
-		case 5:   // MULTILINESTRING
-		case 6:   // MULTIPOLYGON
-		case 7: { // GEOMETRYCOLLECTION
-			reader.Skip(sizeof(uint32_t));
-		} break;
+		// Update the geometry type
+		const auto flag = UnifiedVectorFormat::GetData<uint8_t>(flag_format)[flag_idx];
+		const auto type = UnifiedVectorFormat::GetData<uint8_t>(type_format)[type_idx];
+		if (flag == 1 || flag == 3) {
+			// M or ZM
+			throw InvalidInputException("Geoparquet does not support geometries with M coordinates");
 		}
+		const auto has_z = flag == 2;
+		auto wkb_type = static_cast<WKBGeometryType>((type + 1) + (has_z ? 1000 : 0));
+		meta.geometry_types.insert(wkb_type);
+
+		// Update the bounding box
+		const auto min_x = UnifiedVectorFormat::GetData<double>(xmin_format)[bbox_idx];
+		const auto min_y = UnifiedVectorFormat::GetData<double>(ymin_format)[bbox_idx];
+		const auto max_x = UnifiedVectorFormat::GetData<double>(xmax_format)[bbox_idx];
+		const auto max_y = UnifiedVectorFormat::GetData<double>(ymax_format)[bbox_idx];
+		meta.bbox.Combine(min_x, max_x, min_y, max_y);
 	}
 }
 
@@ -221,6 +208,13 @@ unique_ptr<GeoParquetFileMetadata> GeoParquetFileMetadata::TryRead(const duckdb_
 					throw InvalidInputException("Geoparquet version %s is not supported", result->version);
 				}
 
+				// Check and parse the primary geometry column
+				const auto primary_geometry_column_val = yyjson_obj_get(root, "primary_column");
+				if (!yyjson_is_str(primary_geometry_column_val)) {
+					throw InvalidInputException("Geoparquet metadata does not have a primary column");
+				}
+				result->primary_geometry_column = yyjson_get_str(primary_geometry_column_val);
+
 				// Check and parse the geometry columns
 				const auto columns_val = yyjson_obj_get(root, "columns");
 				if (!yyjson_is_obj(columns_val)) {
@@ -249,18 +243,6 @@ unique_ptr<GeoParquetFileMetadata> GeoParquetFileMetadata::TryRead(const duckdb_
 					const auto encoding_str = yyjson_get_str(encoding_val);
 					if (strcmp(encoding_str, "WKB") == 0) {
 						column.geometry_encoding = GeoParquetColumnEncoding::WKB;
-					} else if (strcmp(encoding_str, "point") == 0) {
-						column.geometry_encoding = GeoParquetColumnEncoding::POINT;
-					} else if (strcmp(encoding_str, "linestring") == 0) {
-						column.geometry_encoding = GeoParquetColumnEncoding::LINESTRING;
-					} else if (strcmp(encoding_str, "polygon") == 0) {
-						column.geometry_encoding = GeoParquetColumnEncoding::POLYGON;
-					} else if (strcmp(encoding_str, "multipoint") == 0) {
-						column.geometry_encoding = GeoParquetColumnEncoding::MULTIPOINT;
-					} else if (strcmp(encoding_str, "multilinestring") == 0) {
-						column.geometry_encoding = GeoParquetColumnEncoding::MULTILINESTRING;
-					} else if (strcmp(encoding_str, "multipolygon") == 0) {
-						column.geometry_encoding = GeoParquetColumnEncoding::MULTIPOLYGON;
 					} else {
 						throw InvalidInputException("Geoparquet column '%s' has an unsupported encoding", column_name);
 					}
@@ -291,53 +273,18 @@ unique_ptr<GeoParquetFileMetadata> GeoParquetFileMetadata::TryRead(const duckdb_
 	return nullptr;
 }
 
-void GeoParquetFileMetadata::AddGeoParquetStats(const string &column_name, const LogicalType &type,
-                                                const GeometryStats &stats) {
-
+void GeoParquetFileMetadata::FlushColumnMeta(const string &column_name, const GeoParquetColumnMetadata &meta) {
 	// Lock the metadata
 	lock_guard<mutex> glock(write_lock);
 
-	auto it = geometry_columns.find(column_name);
-	if (it == geometry_columns.end()) {
-		auto &column = geometry_columns[column_name];
+	auto &column = geometry_columns[column_name];
 
-		column.stats.types.Combine(stats.types);
-		column.stats.bbox.Combine(stats.bbox);
-		column.insertion_index = geometry_columns.size() - 1;
-	} else {
-		it->second.stats.types.Combine(stats.types);
-		it->second.stats.bbox.Combine(stats.bbox);
-	}
+	// Combine the metadata
+	column.geometry_types.insert(meta.geometry_types.begin(), meta.geometry_types.end());
+	column.bbox.Combine(meta.bbox);
 }
 
-void GeoParquetFileMetadata::Write(duckdb_parquet::FileMetaData &file_meta_data) {
-
-	// GeoParquet does not support M or ZM coordinates. So remove any columns that have them.
-	unordered_set<string> invalid_columns;
-	for (auto &column : geometry_columns) {
-		if (column.second.stats.bbox.HasM()) {
-			invalid_columns.insert(column.first);
-		}
-	}
-	for (auto &col_name : invalid_columns) {
-		geometry_columns.erase(col_name);
-	}
-	// No columns remaining, nothing to write
-	if (geometry_columns.empty()) {
-		return;
-	}
-
-	// Find the primary geometry column
-	const auto &random_first_column = *geometry_columns.begin();
-	auto primary_geometry_column = random_first_column.first;
-	auto primary_insertion_index = random_first_column.second.insertion_index;
-
-	for (auto &column : geometry_columns) {
-		if (column.second.insertion_index < primary_insertion_index) {
-			primary_insertion_index = column.second.insertion_index;
-			primary_geometry_column = column.first;
-		}
-	}
+void GeoParquetFileMetadata::Write(duckdb_parquet::FileMetaData &file_meta_data) const {
 
 	yyjson_mut_doc *doc = yyjson_mut_doc_new(nullptr);
 	yyjson_mut_val *root = yyjson_mut_obj(doc);
@@ -354,34 +301,18 @@ void GeoParquetFileMetadata::Write(duckdb_parquet::FileMetaData &file_meta_data)
 	const auto json_columns = yyjson_mut_obj_add_obj(doc, root, "columns");
 
 	for (auto &column : geometry_columns) {
-
 		const auto column_json = yyjson_mut_obj_add_obj(doc, json_columns, column.first.c_str());
 		yyjson_mut_obj_add_str(doc, column_json, "encoding", "WKB");
 		const auto geometry_types = yyjson_mut_obj_add_arr(doc, column_json, "geometry_types");
-		for (auto &type_name : column.second.stats.types.ToString(false)) {
-			yyjson_mut_arr_add_strcpy(doc, geometry_types, type_name.c_str());
-		}
-
-		const auto &bbox = column.second.stats.bbox;
-
-		if (bbox.IsSet()) {
-
-			const auto bbox_arr = yyjson_mut_obj_add_arr(doc, column_json, "bbox");
-
-			if (!column.second.stats.bbox.HasZ()) {
-				yyjson_mut_arr_add_real(doc, bbox_arr, bbox.xmin);
-				yyjson_mut_arr_add_real(doc, bbox_arr, bbox.ymin);
-				yyjson_mut_arr_add_real(doc, bbox_arr, bbox.xmax);
-				yyjson_mut_arr_add_real(doc, bbox_arr, bbox.ymax);
-			} else {
-				yyjson_mut_arr_add_real(doc, bbox_arr, bbox.xmin);
-				yyjson_mut_arr_add_real(doc, bbox_arr, bbox.ymin);
-				yyjson_mut_arr_add_real(doc, bbox_arr, bbox.zmin);
-				yyjson_mut_arr_add_real(doc, bbox_arr, bbox.xmax);
-				yyjson_mut_arr_add_real(doc, bbox_arr, bbox.ymax);
-				yyjson_mut_arr_add_real(doc, bbox_arr, bbox.zmax);
-			}
+		for (auto &geometry_type : column.second.geometry_types) {
+			const auto type_name = WKBGeometryTypes::ToString(geometry_type);
+			yyjson_mut_arr_add_str(doc, geometry_types, type_name);
 		}
+		const auto bbox = yyjson_mut_obj_add_arr(doc, column_json, "bbox");
+		yyjson_mut_arr_add_real(doc, bbox, column.second.bbox.min_x);
+		yyjson_mut_arr_add_real(doc, bbox, column.second.bbox.min_y);
+		yyjson_mut_arr_add_real(doc, bbox, column.second.bbox.max_x);
+		yyjson_mut_arr_add_real(doc, bbox, column.second.bbox.max_y);
 
 		// If the CRS is present, add it
 		if (!column.second.projjson.empty()) {
@@ -423,6 +354,14 @@ bool GeoParquetFileMetadata::IsGeometryColumn(const string &column_name) const {
 	return geometry_columns.find(column_name) != geometry_columns.end();
 }
 
+void GeoParquetFileMetadata::RegisterGeometryColumn(const string &column_name) {
+	lock_guard<mutex> glock(write_lock);
+	if (primary_geometry_column.empty()) {
+		primary_geometry_column = column_name;
+	}
+	geometry_columns[column_name] = GeoParquetColumnMetadata();
+}
+
 bool GeoParquetFileMetadata::IsGeoParquetConversionEnabled(const ClientContext &context) {
 	Value geoparquet_enabled;
 	if (!context.TryGetCurrentSetting("enable_geoparquet_conversion", geoparquet_enabled)) {
@@ -445,19 +384,20 @@ LogicalType GeoParquetFileMetadata::GeometryType() {
 	return blob_type;
 }
 
-const unordered_map<string, GeoParquetColumnMetadata> &GeoParquetFileMetadata::GetColumnMeta() const {
-	return geometry_columns;
-}
-
 unique_ptr<ColumnReader> GeoParquetFileMetadata::CreateColumnReader(ParquetReader &reader,
                                                                     const ParquetColumnSchema &schema,
                                                                     ClientContext &context) {
 
+	D_ASSERT(IsGeometryColumn(schema.name));
+
+	const auto &column = geometry_columns[schema.name];
+
 	// Get the catalog
 	auto &catalog = Catalog::GetSystemCatalog(context);
 
 	// WKB encoding
-	if (schema.children[0].type.id() == LogicalTypeId::BLOB) {
+	if (schema.children[0].type.id() == LogicalTypeId::BLOB &&
+	    column.geometry_encoding == GeoParquetColumnEncoding::WKB) {
 		// Look for a conversion function in the catalog
 		auto &conversion_func_set =
 		    catalog.GetEntry<ScalarFunctionCatalogEntry>(context, DEFAULT_SCHEMA, "st_geomfromwkb");

package/src/duckdb/extension/parquet/include/column_reader.hpp

@@ -21,11 +21,13 @@
 #include "decoder/delta_length_byte_array_decoder.hpp"
 #include "decoder/delta_byte_array_decoder.hpp"
 #include "parquet_column_schema.hpp"
+#ifndef DUCKDB_AMALGAMATION
 
 #include "duckdb/common/operator/cast_operators.hpp"
 #include "duckdb/common/types/string_type.hpp"
 #include "duckdb/common/types/vector.hpp"
 #include "duckdb/common/types/vector_cache.hpp"
+#endif
 
 namespace duckdb {
 class ParquetReader;

package/src/duckdb/extension/parquet/include/column_writer.hpp

@@ -119,8 +119,7 @@ public:
 		throw NotImplementedException("Writer does not need analysis");
 	}
 
-	virtual void Prepare(ColumnWriterState &state, ColumnWriterState *parent, Vector &vector, idx_t count,
-	                     bool vector_can_span_multiple_pages) = 0;
+	virtual void Prepare(ColumnWriterState &state, ColumnWriterState *parent, Vector &vector, idx_t count) = 0;
 
 	virtual void BeginWrite(ColumnWriterState &state) = 0;
 	virtual void Write(ColumnWriterState &state, Vector &vector, idx_t count) = 0;