Merge pull request #380 from Maxxen/main

spatial/src/spatial/core/functions/scalar/st_asgeojson.cpp

@@ -238,10 +238,10 @@ static void GeometryToGeoJSONFragmentFunction(DataChunk &args, ExpressionState &
 		Geometry::Match<ToGeoJSONFunctor>(geom, doc, obj);
 
 		size_t json_size = 0;
-		// TODO: YYJSON_WRITE_PRETTY
-		auto json_data = yyjson_mut_write(doc, 0, &json_size);
-		auto json_str = StringVector::AddString(result, json_data, json_size);
-		return json_str;
+		char *json_data = yyjson_mut_write_opts(doc, 0, json_allocator.GetYYJSONAllocator(), &json_size, nullptr);
+		// Because the arena allocator only resets after each pipeline invocation, we can safely just point into the
+		// arena here without needing to copy the data to the string heap with StringVector::AddString
+		return string_t(json_data, json_size);
 	});
 }
 

spatial/src/spatial/core/functions/scalar/st_extent.cpp

@@ -7,7 +7,6 @@
 #include "spatial/core/geometry/geometry.hpp"
 #include "spatial/core/types.hpp"
 
-
 namespace spatial {
 
 namespace core {
@@ -26,8 +25,9 @@ static uint32_t ReadInt(const bool le, Cursor &cursor) {
 }
 static void ReadWKB(Cursor &cursor, BoundingBox &bbox);
 
-static void ReadWKB(const bool le, const uint32_t type, const bool has_z, const bool has_m, Cursor &cursor, BoundingBox &bbox) {
-	switch(type) {
+static void ReadWKB(const bool le, const uint32_t type, const bool has_z, const bool has_m, Cursor &cursor,
+                    BoundingBox &bbox) {
+	switch (type) {
 	case 1: { // POINT
 		// Points are special in that they can be all-nan (empty)
 		bool all_nan = true;
@@ -44,41 +44,41 @@ static void ReadWKB(const bool le, const uint32_t type, const bool has_z, const
 	} break;
 	case 2: { // LINESTRING
 		const auto num_verts = ReadInt(le, cursor);
-		for(uint32_t i = 0; i < num_verts; i++) {
+		for (uint32_t i = 0; i < num_verts; i++) {
 			const auto x = ReadDouble(le, cursor);
 			const auto y = ReadDouble(le, cursor);
-			if(has_z) {
+			if (has_z) {
 				ReadDouble(le, cursor);
 			}
-			if(has_m) {
+			if (has_m) {
 				ReadDouble(le, cursor);
 			}
 			bbox.Stretch(x, y);
 		}
 	} break;
 	case 3: { // POLYGON
 		const auto num_rings = ReadInt(le, cursor);
-		for(uint32_t i = 0; i < num_rings; i++) {
+		for (uint32_t i = 0; i < num_rings; i++) {
 			const auto num_verts = ReadInt(le, cursor);
-			for(uint32_t j = 0; j < num_verts; j++) {
+			for (uint32_t j = 0; j < num_verts; j++) {
 				const auto x = ReadDouble(le, cursor);
 				const auto y = ReadDouble(le, cursor);
-				if(has_z) {
+				if (has_z) {
 					ReadDouble(le, cursor);
 				}
-				if(has_m) {
+				if (has_m) {
 					ReadDouble(le, cursor);
 				}
 				bbox.Stretch(x, y);
 			}
 		}
 	} break;
-	case 4: // MULTIPOINT
-	case 5: // MULTILINESTRING
-	case 6: // MULTIPOLYGON
+	case 4:   // MULTIPOINT
+	case 5:   // MULTILINESTRING
+	case 6:   // MULTIPOLYGON
 	case 7: { // GEOMETRYCOLLECTION
 		const auto num_items = ReadInt(le, cursor);
-		for(uint32_t i = 0; i < num_items; i++) {
+		for (uint32_t i = 0; i < num_items; i++) {
 			ReadWKB(cursor, bbox);
 		}
 	} break;
@@ -98,7 +98,7 @@ static void ReadWKB(Cursor &cursor, BoundingBox &bbox) {
 
 	// Skip SRID if present
 	const auto has_srid = (type & 0x20000000) != 0;
-	if(has_srid) {
+	if (has_srid) {
 		cursor.Skip(sizeof(uint32_t));
 	}
 

spatial/src/spatial/core/functions/scalar/st_geometrytype.cpp

@@ -55,9 +55,7 @@ static void GeometryTypeFunction(DataChunk &args, ExpressionState &state, Vector
 	auto &input = args.data[0];
 
 	UnaryExecutor::Execute<geometry_t, uint8_t>(
-	    input, result, count, [&](const geometry_t &geom) {
-		    return static_cast<uint8_t>(geom.GetType());
-	    });
+	    input, result, count, [&](const geometry_t &geom) { return static_cast<uint8_t>(geom.GetType()); });
 }
 
 //------------------------------------------------------------------------------
@@ -72,7 +70,7 @@ static void WKBTypeFunction(DataChunk &args, ExpressionState &state, Vector &res
 		const auto le = cursor.Read<uint8_t>();
 		const auto type = le ? cursor.Read<uint32_t>() : cursor.ReadBigEndian<uint32_t>();
 		const auto normalized_type = (type & 0xffff) % 1000;
-		if(normalized_type == 0 || normalized_type > 7) {
+		if (normalized_type == 0 || normalized_type > 7) {
 			throw InvalidInputException("WKB type '%d' is not a supported geometry type", type);
 		}
 
@@ -108,7 +106,7 @@ void CoreScalarFunctions::RegisterStGeometryType(DatabaseInstance &db) {
 	    ScalarFunction({GeoTypes::GEOMETRY()}, LogicalType::ANY, GeometryTypeFunction, GeometryTypeFunctionBind));
 
 	geometry_type_set.AddFunction(
-		ScalarFunction({GeoTypes::WKB_BLOB()}, LogicalType::ANY, WKBTypeFunction, GeometryTypeFunctionBind));
+	    ScalarFunction({GeoTypes::WKB_BLOB()}, LogicalType::ANY, WKBTypeFunction, GeometryTypeFunctionBind));
 
 	ExtensionUtil::RegisterFunction(db, geometry_type_set);
 	DocUtil::AddDocumentation(db, "ST_GeometryType", DOC_DESCRIPTION, DOC_EXAMPLE, DOC_TAGS);

spatial/src/spatial/core/functions/scalar/st_has.cpp

@@ -31,9 +31,15 @@ static void GeometryZMFlagFunction(DataChunk &args, ExpressionState &state, Vect
 		const auto has_z = props.HasZ();
 		const auto has_m = props.HasM();
 
-		if(has_z && has_m) { return 3; }
-		if(has_z) { return 2; }
-		if(has_m) { return 1; }
+		if (has_z && has_m) {
+			return 3;
+		}
+		if (has_z) {
+			return 2;
+		}
+		if (has_m) {
+			return 1;
+		}
 		return 0;
 	});
 }
@@ -68,9 +74,15 @@ static void WKBZMFlagFunction(DataChunk &args, ExpressionState &state, Vector &r
 		const auto has_z = (iso_wkb_props == 1) || (iso_wkb_props == 3) || ((type & 0x80000000) != 0);
 		const auto has_m = (iso_wkb_props == 2) || (iso_wkb_props == 3) || ((type & 0x40000000) != 0);
 
-		if(has_z && has_m) { return 3; }
-		if(has_z) { return 2; }
-		if(has_m) { return 1; }
+		if (has_z && has_m) {
+			return 3;
+		}
+		if (has_z) {
+			return 2;
+		}
+		if (has_m) {
+			return 1;
+		}
 		return 0;
 	});
 }
@@ -170,25 +182,16 @@ static constexpr const char *ZMFLAG_EXAMPLE = R"(
 //------------------------------------------------------------------------------
 void CoreScalarFunctions::RegisterStHas(DatabaseInstance &db) {
 	ScalarFunctionSet st_hasz("ST_HasZ");
-	st_hasz.AddFunction(ScalarFunction({GeoTypes::GEOMETRY()}, LogicalType::BOOLEAN,
-		GeometryHasFunction<true>));
-	st_hasz.AddFunction(ScalarFunction({GeoTypes::WKB_BLOB()}, LogicalType::BOOLEAN,
-		WKBHasFunction<true>));
-
+	st_hasz.AddFunction(ScalarFunction({GeoTypes::GEOMETRY()}, LogicalType::BOOLEAN, GeometryHasFunction<true>));
+	st_hasz.AddFunction(ScalarFunction({GeoTypes::WKB_BLOB()}, LogicalType::BOOLEAN, WKBHasFunction<true>));
 
 	ScalarFunctionSet st_hasm("ST_HasM");
-	st_hasm.AddFunction(ScalarFunction({GeoTypes::GEOMETRY()}, LogicalType::BOOLEAN,
-		GeometryHasFunction<false>));
-	st_hasm.AddFunction(ScalarFunction({GeoTypes::WKB_BLOB()}, LogicalType::BOOLEAN,
-		WKBHasFunction<false>));
-
+	st_hasm.AddFunction(ScalarFunction({GeoTypes::GEOMETRY()}, LogicalType::BOOLEAN, GeometryHasFunction<false>));
+	st_hasm.AddFunction(ScalarFunction({GeoTypes::WKB_BLOB()}, LogicalType::BOOLEAN, WKBHasFunction<false>));
 
 	ScalarFunctionSet st_zmflag("ST_ZMFlag");
-	st_zmflag.AddFunction(ScalarFunction({GeoTypes::GEOMETRY()}, LogicalType::UTINYINT,
-		GeometryZMFlagFunction));
-	st_zmflag.AddFunction(ScalarFunction({GeoTypes::WKB_BLOB()}, LogicalType::UTINYINT,
-		WKBZMFlagFunction));
-
+	st_zmflag.AddFunction(ScalarFunction({GeoTypes::GEOMETRY()}, LogicalType::UTINYINT, GeometryZMFlagFunction));
+	st_zmflag.AddFunction(ScalarFunction({GeoTypes::WKB_BLOB()}, LogicalType::UTINYINT, WKBZMFlagFunction));
 
 	ExtensionUtil::RegisterFunction(db, st_hasz);
 	ExtensionUtil::RegisterFunction(db, st_hasm);

spatial/src/spatial/core/io/osm/st_read_osm.cpp

@@ -284,6 +284,20 @@ struct LocalState : LocalTableFunctionState {
 			string_table.push_back(string_table_reader.get_string());
 		}
 
+		// Need to read ahead without advancing block_reader
+		auto reader_copy = block_reader;
+
+		// Read the granularity and optional offsets
+		if (reader_copy.next(17)) {
+			granularity = reader_copy.get_int32();
+		}
+		if (reader_copy.next(19)) {
+			lat_offset = reader_copy.get_int64();
+		}
+		if (reader_copy.next(20)) {
+			lon_offset = reader_copy.get_int64();
+		}
+
 		state = ParseState::Block;
 	}
 
@@ -310,17 +324,6 @@ struct LocalState : LocalTableFunctionState {
 			case ParseState::Block:
 				if (block_reader.next(2)) {
 					group_reader = block_reader.get_message();
-
-					// Read the granularity and optional offsets
-					if (block_reader.next(17)) {
-						granularity = block_reader.get_int32();
-					}
-					if (block_reader.next(19)) {
-						lat_offset = block_reader.get_int64();
-					}
-					if (block_reader.next(20)) {
-						lon_offset = block_reader.get_int64();
-					}
 					state = ParseState::Group;
 				} else {
 					state = ParseState::End;

spatial/src/spatial/core/io/shapefile/read_shapefile.cpp

@@ -37,8 +37,8 @@ struct ShapefileBindData : TableFunctionData {
 	vector<LogicalType> attribute_types;
 
 	explicit ShapefileBindData(string file_name_p)
-	    : file_name(std::move(file_name_p)), shape_count(0), shape_type(0), min_bound {0, 0, 0, 0},
-	      max_bound {0, 0, 0, 0}, attribute_encoding(AttributeEncoding::LATIN1) {
+	    : file_name(std::move(file_name_p)), shape_count(0),
+	      shape_type(0), min_bound {0, 0, 0, 0}, max_bound {0, 0, 0, 0}, attribute_encoding(AttributeEncoding::LATIN1) {
 	}
 };
 

spatial/src/spatial/gdal/file_handler.cpp

@@ -385,10 +385,8 @@ string GDALClientContextState::GetPrefix(const string &value) const {
 }
 
 GDALClientContextState &GDALClientContextState::GetOrCreate(ClientContext &context) {
-	if (!context.registered_state["gdal"]) {
-		context.registered_state["gdal"] = make_uniq<GDALClientContextState>(context);
-	}
-	return *dynamic_cast<GDALClientContextState *>(context.registered_state["gdal"].get());
+	auto gdal_state = context.registered_state->GetOrCreate<GDALClientContextState>("gdal", context);
+	return *gdal_state;
 }
 
 } // namespace gdal