diff --git a/src/geoarrow_geos/geoarrow_geos.c b/src/geoarrow_geos/geoarrow_geos.c index 49cb6cc..cc84589 100644 --- a/src/geoarrow_geos/geoarrow_geos.c +++ b/src/geoarrow_geos/geoarrow_geos.c @@ -363,8 +363,43 @@ struct GeoArrowGEOSArrayReader { // faster than GEOS' own readers. GEOSWKTReader* wkt_reader; GEOSWKBReader* wkb_reader; + // In-progress items that we might need to clean up if an error was returned + int64_t n_geoms[2]; + GEOSGeometry** geoms[2]; }; +static GeoArrowErrorCode GeoArrowGEOSArrayReaderEnsureScratch( + struct GeoArrowGEOSArrayReader* reader, int64_t n_geoms, int level) { + if (n_geoms <= reader->n_geoms[level]) { + return GEOARROW_OK; + } + + if ((reader->n_geoms[level] * 2) > n_geoms) { + n_geoms = reader->n_geoms[level] * 2; + } + + reader->geoms[level] = + (GEOSGeometry**)realloc(reader->geoms[level], n_geoms * sizeof(GEOSGeometry*)); + if (reader->geoms[level] == NULL) { + reader->n_geoms[level] = 0; + return ENOMEM; + } + + memset(reader->geoms[level], 0, n_geoms * sizeof(GEOSGeometry*)); + return GEOARROW_OK; +} + +static void GeoArrowGEOSArrayReaderResetScratch(struct GeoArrowGEOSArrayReader* reader) { + for (int level = 0; level < 2; level++) { + for (int64_t i = 0; i < reader->n_geoms[level]; i++) { + if (reader->geoms[level][i] != NULL) { + GEOSGeom_destroy_r(reader->handle, reader->geoms[level][i]); + reader->geoms[level][i] = NULL; + } + } + } +} + GeoArrowGEOSErrorCode GeoArrowGEOSArrayReaderCreate( GEOSContextHandle_t handle, struct ArrowSchema* schema, struct GeoArrowGEOSArrayReader** out) { @@ -478,9 +513,104 @@ static GeoArrowErrorCode MakeLinestrings(struct GeoArrowGEOSArrayReader* reader, return GEOARROW_OK; } +static GeoArrowErrorCode MakeLinearrings(struct GeoArrowGEOSArrayReader* reader, + size_t offset, size_t length, + GEOSGeometry** out) { + offset += reader->array_view.offset[reader->array_view.n_offsets - 1]; + const int32_t* coord_offsets = + reader->array_view.offsets[reader->array_view.n_offsets - 1]; + + GEOSCoordSequence* seq = NULL; + for (size_t i = 0; i < length; i++) { + GEOARROW_RETURN_NOT_OK( + MakeCoordSeq(reader, coord_offsets[offset + i], + coord_offsets[offset + i + 1] - coord_offsets[offset + i], &seq)); + out[i] = GEOSGeom_createLinearRing_r(reader->handle, seq); + if (out[i] == NULL) { + GEOSCoordSeq_destroy_r(reader->handle, seq); + GeoArrowErrorSet(&reader->error, "[%ld] GEOSGeom_createLinearRing_r() failed", + (long)i); + return ENOMEM; + } + } + + return GEOARROW_OK; +} + +static GeoArrowErrorCode MakePolygons(struct GeoArrowGEOSArrayReader* reader, + size_t offset, size_t length, GEOSGeometry** out) { + offset += reader->array_view.offset[reader->array_view.n_offsets - 2]; + const int32_t* ring_offsets = + reader->array_view.offsets[reader->array_view.n_offsets - 2]; + + for (size_t i = 0; i < length; i++) { + int64_t ring_offset = ring_offsets[offset + i]; + int64_t n_rings = ring_offsets[offset + i + 1] - ring_offset; + + if (n_rings == 0) { + out[i] = GEOSGeom_createEmptyPolygon_r(reader->handle); + } else { + GEOARROW_RETURN_NOT_OK(GeoArrowGEOSArrayReaderEnsureScratch(reader, n_rings, 0)); + GEOARROW_RETURN_NOT_OK( + MakeLinearrings(reader, ring_offset, n_rings, reader->geoms[0])); + out[i] = GEOSGeom_createPolygon_r(reader->handle, reader->geoms[0][0], + reader->geoms[0] + 1, n_rings - 1); + memset(reader->geoms[0], 0, n_rings * sizeof(GEOSGeometry*)); + } + + if (out[i] == NULL) { + GeoArrowErrorSet(&reader->error, "[%ld] GEOSGeom_createPolygon_r() failed", + (long)i); + return ENOMEM; + } + } + + return GEOARROW_OK; +} + +typedef GeoArrowErrorCode (*GeoArrowGEOSPartMaker)(struct GeoArrowGEOSArrayReader* reader, + size_t offset, size_t length, + GEOSGeometry** out); + +static GeoArrowErrorCode MakeCollection(struct GeoArrowGEOSArrayReader* reader, + size_t offset, size_t length, GEOSGeometry** out, + int geom_level, int offset_level, int geos_type, + GeoArrowGEOSPartMaker part_maker) { + offset += reader->array_view.offset[reader->array_view.n_offsets - offset_level]; + const int32_t* part_offsets = + reader->array_view.offsets[reader->array_view.n_offsets - offset_level]; + + for (size_t i = 0; i < length; i++) { + int64_t part_offset = part_offsets[offset + i]; + int64_t n_parts = part_offsets[offset + i + 1] - part_offset; + + if (n_parts == 0) { + out[i] = GEOSGeom_createEmptyCollection_r(reader->handle, geos_type); + } else { + GEOARROW_RETURN_NOT_OK( + GeoArrowGEOSArrayReaderEnsureScratch(reader, n_parts, geom_level)); + GEOARROW_RETURN_NOT_OK( + part_maker(reader, part_offset, n_parts, reader->geoms[geom_level])); + out[i] = GEOSGeom_createCollection_r(reader->handle, geos_type, + reader->geoms[geom_level], n_parts); + memset(reader->geoms[geom_level], 0, n_parts * sizeof(GEOSGeometry*)); + } + + if (out[i] == NULL) { + GeoArrowErrorSet(&reader->error, "[%ld] GEOSGeom_createEmptyCollection_r() failed", + (long)i); + return ENOMEM; + } + } + + return GEOARROW_OK; +} + GeoArrowGEOSErrorCode GeoArrowGEOSArrayReaderRead(struct GeoArrowGEOSArrayReader* reader, struct ArrowArray* array, size_t offset, size_t length, GEOSGeometry** out) { + GeoArrowGEOSArrayReaderResetScratch(reader); + GEOARROW_RETURN_NOT_OK( GeoArrowArrayViewSetArray(&reader->array_view, array, &reader->error)); @@ -494,6 +624,20 @@ GeoArrowGEOSErrorCode GeoArrowGEOSArrayReaderRead(struct GeoArrowGEOSArrayReader case GEOARROW_GEOMETRY_TYPE_LINESTRING: result = MakeLinestrings(reader, offset, length, out); break; + case GEOARROW_GEOMETRY_TYPE_POLYGON: + result = MakePolygons(reader, offset, length, out); + break; + case GEOARROW_GEOMETRY_TYPE_MULTIPOINT: + result = MakeCollection(reader, offset, length, out, 0, 1, GEOS_MULTIPOINT, &MakePoints); + break; + case GEOARROW_GEOMETRY_TYPE_MULTILINESTRING: + result = MakeCollection(reader, offset, length, out, 0, 2, GEOS_MULTILINESTRING, + &MakeLinestrings); + break; + case GEOARROW_GEOMETRY_TYPE_MULTIPOLYGON: + result = MakeCollection(reader, offset, length, out, 1, 3, GEOS_MULTIPOLYGON, + &MakePolygons); + break; default: GeoArrowErrorSet(&reader->error, "GeoArrowGEOSArrayReaderRead not implemented for geometry type"); @@ -522,5 +666,13 @@ void GeoArrowGEOSArrayReaderDestroy(struct GeoArrowGEOSArrayReader* reader) { GEOSWKBReader_destroy_r(reader->handle, reader->wkb_reader); } + GeoArrowGEOSArrayReaderResetScratch(reader); + + for (int i = 0; i < 2; i++) { + if (reader->geoms[i] != NULL) { + free(reader->geoms[i]); + } + } + free(reader); } diff --git a/src/geoarrow_geos/geoarrow_geos_test.cc b/src/geoarrow_geos/geoarrow_geos_test.cc index 6f9fa1c..6bd81c4 100644 --- a/src/geoarrow_geos/geoarrow_geos_test.cc +++ b/src/geoarrow_geos/geoarrow_geos_test.cc @@ -28,6 +28,26 @@ class GEOSCppGeometry { } }; +class GEOSCppGeometryVec { + public: + std::vector ptrs; + GEOSContextHandle_t handle; + + GEOSCppGeometryVec(GEOSContextHandle_t handle) : handle(handle) {} + + GEOSGeometry** data() { return ptrs.data(); } + + const GEOSGeometry** const_data() { return const_cast(data()); } + + ~GEOSCppGeometryVec() { + for (const auto& ptr : ptrs) { + if (ptr != nullptr) { + GEOSGeom_destroy_r(handle, ptr); + } + } + } +}; + class GEOSCppWKTReader { public: GEOSWKTReader* ptr; @@ -156,7 +176,7 @@ TEST(GeoArrowGEOSTest, TestArrayBuilderRoundtripWKTCollection) { TestBuilderRoundtripWKT("MULTIPOINT (30 10, 40 30, 20 20)"); } -void TestReaderRoundtripWKT(const std::string& wkt, int wkb_type) { +void TestReaderRoundtripWKTVec(const std::vector& wkt, int wkb_type) { GEOSCppHandle handle; GeoArrowGEOSCppArrayBuilder builder(handle.handle); GeoArrowGEOSCppArrayReader reader(handle.handle); @@ -169,13 +189,22 @@ void TestReaderRoundtripWKT(const std::string& wkt, int wkb_type) { ASSERT_EQ(builder.Init(schema.get()), GEOARROW_GEOS_OK); GEOSCppWKTReader wkt_reader(handle.handle); - GEOSCppGeometry geom(handle.handle); - ASSERT_EQ(wkt_reader.Read(wkt, &geom.ptr), GEOARROW_GEOS_OK); + + GEOSCppGeometryVec geoms_in(handle.handle); + GEOSCppGeometryVec geoms_out(handle.handle); + for (const auto& wkt_item : wkt) { + geoms_in.ptrs.push_back(nullptr); + geoms_out.ptrs.push_back(nullptr); + + ASSERT_EQ(wkt_reader.Read(wkt_item, &geoms_in.ptrs.back()), GEOARROW_GEOS_OK) + << "Failed to append " << wkt_item; + } size_t n = 0; - const GEOSGeometry* geom_const = geom.ptr; - ASSERT_EQ(GeoArrowGEOSArrayBuilderAppend(builder.ptr, &geom_const, 1, &n), - GEOARROW_GEOS_OK); + ASSERT_EQ( + GeoArrowGEOSArrayBuilderAppend(builder.ptr, geoms_in.const_data(), wkt.size(), &n), + GEOARROW_GEOS_OK); + ASSERT_EQ(n, wkt.size()); nanoarrow::UniqueArray array; ASSERT_EQ(GeoArrowGEOSArrayBuilderFinish(builder.ptr, array.get()), GEOARROW_GEOS_OK); @@ -183,16 +212,21 @@ void TestReaderRoundtripWKT(const std::string& wkt, int wkb_type) { // Read it back! ASSERT_EQ(reader.Init(schema.get()), GEOARROW_GEOS_OK); - GEOSCppGeometry geom_out(handle.handle); - ASSERT_EQ(GeoArrowGEOSArrayReaderRead(reader.ptr, array.get(), 0, 1, &geom_out.ptr), + ASSERT_EQ(GeoArrowGEOSArrayReaderRead(reader.ptr, array.get(), 0, array->length, + geoms_out.data()), GEOARROW_GEOS_OK) - << "WKT: " << wkt + << "WKT[0]: " << wkt[0] << " n = " << n << "\n Error: " << GeoArrowGEOSArrayReaderGetLastError(reader.ptr); // Check for GEOS equality - EXPECT_EQ(GEOSEqualsExact_r(handle.handle, geom_out.ptr, geom.ptr, 0), 1) - << "WKT: " << wkt - << "\n Error: " << GeoArrowGEOSArrayReaderGetLastError(reader.ptr); + for (size_t i = 0; i < n; i++) { + EXPECT_EQ(GEOSEqualsExact_r(handle.handle, geoms_out.ptrs[i], geoms_in.ptrs[i], 0), 1) + << "WKT: " << wkt[i] << " at index " << i; + } +} + +void TestReaderRoundtripWKT(const std::string& wkt, int wkb_type) { + TestReaderRoundtripWKTVec({wkt}, wkb_type); } TEST(GeoArrowGEOSTest, TestArrayReaderPoint) { @@ -200,11 +234,134 @@ TEST(GeoArrowGEOSTest, TestArrayReaderPoint) { TestReaderRoundtripWKT("POINT (0 1)", 1); TestReaderRoundtripWKT("POINT Z EMPTY", 1001); TestReaderRoundtripWKT("POINT Z (0 1 2)", 1001); + + TestReaderRoundtripWKTVec({}, 1); + TestReaderRoundtripWKTVec({}, 1001); + TestReaderRoundtripWKTVec({"POINT EMPTY", "POINT (0 1)", "POINT (2 3)", "POINT EMPTY"}, + 1); + TestReaderRoundtripWKTVec( + {"POINT Z EMPTY", "POINT Z (0 1 2)", "POINT Z (3 4 5)", "POINT Z EMPTY"}, 1001); } TEST(GeoArrowGEOSTest, TestArrayReaderLinestring) { TestReaderRoundtripWKT("LINESTRING EMPTY", 2); TestReaderRoundtripWKT("LINESTRING (0 1, 2 3)", 2); - // LINESTRING Z EMPTY doesn't seem to roundtrip through GEOSGeometry + TestReaderRoundtripWKT("LINESTRING Z EMPTY", 2); TestReaderRoundtripWKT("LINESTRING Z (0 1 2, 3 4 5)", 1002); + + TestReaderRoundtripWKTVec({}, 2); + TestReaderRoundtripWKTVec({}, 1002); + TestReaderRoundtripWKTVec({"LINESTRING EMPTY", "LINESTRING (0 1, 2 3)", + "LINESTRING (4 5, 6 7, 8 9)", "LINESTRING EMPTY"}, + 2); + TestReaderRoundtripWKTVec( + {"LINESTRING Z EMPTY", "LINESTRING Z (0 1 2, 3 4 5)", + "LINESTRING Z (6 7 8, 9 10 11, 12 13 14)", "LINESTRING Z EMPTY"}, + 1002); +} + +TEST(GeoArrowGEOSTest, TestArrayReaderPolygon) { + TestReaderRoundtripWKT("POLYGON EMPTY", 3); + TestReaderRoundtripWKT("POLYGON ((30 10, 40 40, 20 40, 10 20, 30 10))", 3); + TestReaderRoundtripWKT( + "POLYGON ((35 10, 45 45, 15 40, 10 20, 35 10), (20 30, 35 35, 30 20, 20 30))", 3); + TestReaderRoundtripWKT("POLYGON Z EMPTY", 1003); + TestReaderRoundtripWKT("POLYGON Z ((30 10 40, 40 40 80, 20 40 60, 10 20 30, 30 10 40))", + 1003); + TestReaderRoundtripWKT( + "POLYGON Z ((35 10 45, 45 45 90, 15 40 55, 10 20 30, 35 10 45), (20 30 50, 35 35 " + "70, 30 20 50, 20 30 50))", + 1003); + TestReaderRoundtripWKT( + "POLYGON Z ((35 10 45, 45 45 90, 15 40 55, 10 20 30, 35 10 45), (20 30 50, 35 35 " + "70, 30 20 50, 20 30 50))", + 1003); + + TestReaderRoundtripWKTVec({}, 3); + TestReaderRoundtripWKTVec({}, 1003); + TestReaderRoundtripWKTVec( + {"POLYGON ((30 10, 40 40, 20 40, 10 20, 30 10))", + "POLYGON ((35 10, 45 45, 15 40, 10 20, 35 10), (20 30, 35 35, 30 20, 20 30))", + "POLYGON EMPTY"}, + 3); + + TestReaderRoundtripWKTVec( + {"POLYGON Z ((30 10 40, 40 40 80, 20 40 60, 10 20 30, 30 10 40))", + "POLYGON Z ((35 10 45, 45 45 90, 15 40 55, 10 20 30, 35 10 45), (20 30 50, 35 35 " + "70, 30 20 50, 20 30 50))", + "POLYGON Z EMPTY"}, + 1003); +} + +TEST(GeoArrowGEOSTest, TestArrayReaderMultipoint) { + TestReaderRoundtripWKT("MULTIPOINT EMPTY", 4); + TestReaderRoundtripWKT("MULTIPOINT (10 40, 40 30, 20 20, 30 10)", 4); + TestReaderRoundtripWKT("MULTIPOINT (30 10)", 4); + + TestReaderRoundtripWKTVec({}, 4); + TestReaderRoundtripWKTVec({}, 1004); + TestReaderRoundtripWKTVec( + {"MULTIPOINT ((30 10))", "MULTIPOINT ((10 40), (40 30), (20 20), (30 10))", + "MULTIPOINT ((10 40), (40 30), (20 20), (30 10))"}, + 4); + + TestReaderRoundtripWKTVec( + {"MULTIPOINT Z ((30 10 40))", + "MULTIPOINT Z ((10 40 50), (40 30 70), (20 20 40), (30 10 40))", + "MULTIPOINT Z ((10 40 50), (40 30 70), (20 20 40), (30 10 40))", + "MULTIPOINT Z EMPTY"}, + 1004); +} + +TEST(GeoArrowGEOSTest, TestArrayReaderMultilinestring) { + TestReaderRoundtripWKT("MULTILINESTRING EMPTY", 5); + TestReaderRoundtripWKT("MULTILINESTRING ((30 10, 10 30, 40 40))", 5); + TestReaderRoundtripWKT( + "MULTILINESTRING ((10 10, 20 20, 10 40), (40 40, 30 30, 40 20, 30 10))", 5); + + TestReaderRoundtripWKTVec( + {"MULTILINESTRING ((30 10, 10 30, 40 40))", + "MULTILINESTRING ((10 10, 20 20, 10 40), (40 40, 30 30, 40 20, 30 10))", + "MULTILINESTRING EMPTY"}, + 5); + + TestReaderRoundtripWKTVec({}, 5); + TestReaderRoundtripWKTVec({}, 1005); + TestReaderRoundtripWKTVec({"MULTILINESTRING Z ((30 10 40, 10 30 40, 40 40 80))", + "MULTILINESTRING Z ((10 10 20, 20 20 40, 10 40 50), (40 40 " + "80, 30 30 60, 40 20 60, 30 10 40))", + "MULTILINESTRING Z EMPTY"}, + 1005); +} + +TEST(GeoArrowGEOSTest, TestArrayReaderMultipolygon) { + TestReaderRoundtripWKT("MULTIPOLYGON EMPTY", 6); + TestReaderRoundtripWKT( + "MULTIPOLYGON (((30 20, 45 40, 10 40, 30 20)), ((15 5, 40 10, 10 20, 5 10, 15 5)))", + 6); + TestReaderRoundtripWKT( + "MULTIPOLYGON (((40 40, 20 45, 45 30, 40 40)), ((20 35, 10 30, 10 10, 30 5, 45 20, " + "20 35), (30 20, 20 15, 20 25, 30 20)))", + 6); + + TestReaderRoundtripWKTVec({}, 6); + TestReaderRoundtripWKTVec({}, 1006); + TestReaderRoundtripWKTVec( + {"MULTIPOLYGON (((30 10, 40 40, 20 40, 10 20, 30 10)))", + "MULTIPOLYGON (((30 20, 45 40, 10 40, 30 20)), ((15 5, 40 10, 10 20, 5 10, 15 " + "5)))", + "MULTIPOLYGON (((40 40, 20 45, 45 30, 40 40)), ((20 35, 10 30, 10 10, 30 5, 45 " + "20, 20 35), (30 20, 20 15, 20 25, 30 20)))", + "MULTIPOLYGON EMPTY"}, + 6); + + TestReaderRoundtripWKTVec( + {"MULTIPOLYGON Z (((30 10 40, 40 40 80, 20 40 60, 10 20 30, 30 10 40)))", + "MULTIPOLYGON Z (((30 20 50, 45 40 85, 10 40 50, 30 20 50)), ((15 5 20, 40 10 50, " + "10 20 30, 5 10 15, 15 5 20)))", + "MULTIPOLYGON Z (((40 40 80, 20 45 65, 45 30 75, 40 40 80)), ((20 35 55, 10 30 " + "40, 10 10 20, 30 5 35, 45 20 65, 20 35 55), (30 20 50, 20 15 35, 20 25 45, 30 20 " + "50)))", + "MULTIPOLYGON Z EMPTY"}, + 1006); }