Merge #49554 #49613

49554: roachtest: fix typo in tpchvec/perf r=yuzefovich a=yuzefovich

TPCH queries have semicolon at the end, so when trying to do `SELECT url
FROM [EXPLAIN ANALYZE %s];` with query replacing the `%s`, we would have
a semicolon within square brackets which would produce a syntax error.
Fix this by reading both `automatic` boolean and `url` string and
outputting the latter.

Release note: None

49613: geo/geomfn: implement ST_Centroid for Geometry and string r=sumeerbhola a=otan

Also had to add hacks in to handle `POINT EMPTY` decoding -- this is
currently a problem with both GEOS and twpayne/go-geom, which I have
circumvented for now. An issue has been filed to track this.

Resolves #48893.
Resolves #48805.

Release note (sql change): Implement the ST_Centroid builtin which works
for Geometry and string arguments.

Co-authored-by: Yahor Yuzefovich <[email protected]>
Co-authored-by: Oliver Tan <[email protected]>

docs/generated/sql/functions.md

@@ -737,6 +737,12 @@ has no relationship with the commit order of concurrent transactions.</p>
 </span></td></tr>
 <tr><td><a name="st_astext"></a><code>st_astext(geometry: geometry) &rarr; <a href="string.html">string</a></code></td><td><span class="funcdesc"><p>Returns the WKT representation of a given Geometry.</p>
 </span></td></tr>
+<tr><td><a name="st_centroid"></a><code>st_centroid(geometry: geometry) &rarr; geometry</code></td><td><span class="funcdesc"><p>Returns the centroid of the given geometry.</p>
+<p>This function utilizes the GEOS module.</p>
+</span></td></tr>
+<tr><td><a name="st_centroid"></a><code>st_centroid(val: <a href="string.html">string</a>) &rarr; geometry</code></td><td><span class="funcdesc"><p>Returns the centroid of the given string, which will be parsed as a geometry object.</p>
+<p>This function utilizes the GEOS module.</p>
+</span></td></tr>
 <tr><td><a name="st_contains"></a><code>st_contains(geometry_a: geometry, geometry_b: geometry) &rarr; <a href="bool.html">bool</a></code></td><td><span class="funcdesc"><p>Returns true if no points of geometry_b lie in the exterior of geometry_a, and there is at least one point in the interior of geometry_b that lies in the interior of geometry_a.</p>
 <p>This function utilizes the GEOS module.</p>
 <p>This function will automatically use any available index.</p>

pkg/cmd/roachtest/tpchvec.go

@@ -267,15 +267,20 @@ func (p *tpchVecPerfTest) postTestRunHook(t *test, conn *gosql.DB, version crdbV
 				}
 				for i := 0; i < tpchPerfTestNumRunsPerQuery; i++ {
 					rows, err := conn.Query(fmt.Sprintf(
-						"SELECT url FROM [EXPLAIN ANALYZE %s];", tpch.QueriesByNumber[queryNum],
+						"EXPLAIN ANALYZE %s;", tpch.QueriesByNumber[queryNum],
 					))
 					if err != nil {
 						t.Fatal(err)
 					}
-					defer rows.Close()
-					var url string
+					var (
+						automatic bool
+						url       string
+					)
 					rows.Next()
-					if err = rows.Scan(&url); err != nil {
+					if err = rows.Scan(&automatic, &url); err != nil {
+						t.Fatal(err)
+					}
+					if err = rows.Close(); err != nil {
 						t.Fatal(err)
 					}
 					t.Status(fmt.Sprintf("EXPLAIN ANALYZE with vectorize=%s url:\n%s", vectorizeSetting, url))

pkg/geo/encode.go

@@ -27,6 +27,11 @@ import (
 
 // EWKBToWKT transforms a given EWKB to WKT.
 func EWKBToWKT(b geopb.EWKB) (geopb.WKT, error) {
+	// twpayne/go-geom doesn't seem to handle POINT EMPTY just yet. Add this hack in.
+	// Remove after #49209 is resolved.
+	if bytes.Equal(b, []byte{0x01, 0x01, 0x00, 0x00, 0x00}) {
+		return geopb.WKT("POINT EMPTY"), nil
+	}
 	t, err := ewkb.Unmarshal([]byte(b))
 	if err != nil {
 		return "", err
@@ -37,6 +42,11 @@ func EWKBToWKT(b geopb.EWKB) (geopb.WKT, error) {
 
 // EWKBToEWKT transforms a given EWKB to EWKT.
 func EWKBToEWKT(b geopb.EWKB) (geopb.EWKT, error) {
+	// twpayne/go-geom doesn't seem to handle POINT EMPTY just yet. Add this hack in.
+	// Remove after #49209 is resolved.
+	if bytes.Equal(b, []byte{0x01, 0x01, 0x00, 0x00, 0x00}) {
+		return geopb.EWKT("POINT EMPTY"), nil
+	}
 	t, err := ewkb.Unmarshal([]byte(b))
 	if err != nil {
 		return "", err

pkg/geo/geomfn/unary_operators.go

@@ -19,6 +19,25 @@ import (
 	"github.com/twpayne/go-geom/encoding/ewkb"
 )
 
+// Centroid returns the Centroid of a given Geometry.
+func Centroid(g *geo.Geometry) (*geo.Geometry, error) {
+	// Empty geometries do not react well in GEOS, so we have to
+	// convert and check beforehand.
+	// Remove after #49209 is resolved.
+	t, err := g.AsGeomT()
+	if err != nil {
+		return nil, err
+	}
+	if t.Empty() {
+		return geo.NewGeometryFromGeom(geom.NewPointEmpty(geom.XY))
+	}
+	centroidEWKB, err := geos.Centroid(g.EWKB())
+	if err != nil {
+		return nil, err
+	}
+	return geo.ParseGeometryFromEWKB(centroidEWKB)
+}
+
 // Length returns the length of a given Geometry.
 // Note only (MULTI)LINESTRING objects have a length.
 // (MULTI)POLYGON objects should use Perimeter.

pkg/geo/geomfn/unary_operators_test.go

@@ -15,8 +15,48 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/geo"
 	"github.com/stretchr/testify/require"
+	"github.com/twpayne/go-geom"
 )
 
+func TestCentroid(t *testing.T) {
+	testCases := []struct {
+		wkt      string
+		expected string
+	}{
+		{"POINT(1.0 1.0)", "POINT (1.0 1.0)"},
+		{"SRID=4326;POINT(1.0 1.0)", "SRID=4326;POINT (1.0 1.0)"},
+		{"LINESTRING(1.0 1.0, 2.0 2.0, 3.0 3.0)", "POINT (2.0 2.0)"},
+		{"POLYGON((0.0 0.0, 1.0 0.0, 1.0 1.0, 0.0 0.0))", "POINT (0.666666666666667 0.333333333333333)"},
+		{"POLYGON((0.0 0.0, 1.0 0.0, 1.0 1.0, 0.0 0.0), (0.1 0.1, 0.2 0.1, 0.2 0.2, 0.1 0.1))", "POINT (0.671717171717172 0.335353535353535)"},
+		{"MULTIPOINT((1.0 1.0), (2.0 2.0))", "POINT (1.5 1.5)"},
+		{"MULTILINESTRING((1.0 1.0, 2.0 2.0, 3.0 3.0), (6.0 6.0, 7.0 6.0))", "POINT (3.17541743733684 3.04481549985497)"},
+		{"MULTIPOLYGON(((3.0 3.0, 4.0 3.0, 4.0 4.0, 3.0 3.0)), ((0.0 0.0, 1.0 0.0, 1.0 1.0, 0.0 0.0), (0.1 0.1, 0.2 0.1, 0.2 0.2, 0.1 0.1)))", "POINT (2.17671691792295 1.84187604690117)"},
+		{"GEOMETRYCOLLECTION (POINT (40 10),LINESTRING (10 10, 20 20, 10 40),POLYGON ((40 40, 20 45, 45 30, 40 40)))", "POINT (35 38.3333333333333)"},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.wkt, func(t *testing.T) {
+			g, err := geo.ParseGeometry(tc.wkt)
+			require.NoError(t, err)
+			ret, err := Centroid(g)
+			require.NoError(t, err)
+
+			retAsGeomT, err := ret.AsGeomT()
+			require.NoError(t, err)
+
+			expected, err := geo.ParseGeometry(tc.expected)
+			require.NoError(t, err)
+			expectedAsGeomT, err := expected.AsGeomT()
+			require.NoError(t, err)
+
+			// Ensure points are close in terms of precision.
+			require.InEpsilon(t, expectedAsGeomT.(*geom.Point).X(), retAsGeomT.(*geom.Point).X(), 2e-10)
+			require.InEpsilon(t, expectedAsGeomT.(*geom.Point).Y(), retAsGeomT.(*geom.Point).Y(), 2e-10)
+			require.Equal(t, expected.SRID(), ret.SRID())
+		})
+	}
+}
+
 func TestLength(t *testing.T) {
 	testCases := []struct {
 		wkt      string

pkg/geo/geos/geos.cc

@@ -65,6 +65,7 @@ typedef void (*CR_GEOS_WKBReader_destroy_r)(CR_GEOS_Handle, CR_GEOS_WKBReader);
 
 typedef int (*CR_GEOS_Area_r)(CR_GEOS_Handle, CR_GEOS_Geometry, double*);
 typedef int (*CR_GEOS_Length_r)(CR_GEOS_Handle, CR_GEOS_Geometry, double*);
+typedef CR_GEOS_Geometry (*CR_GEOS_Centroid_r)(CR_GEOS_Handle, CR_GEOS_Geometry);
 
 typedef int (*CR_GEOS_Distance_r)(CR_GEOS_Handle, CR_GEOS_Geometry, CR_GEOS_Geometry, double*);
 
@@ -117,6 +118,7 @@ struct CR_GEOS {
 
   CR_GEOS_Area_r GEOSArea_r;
   CR_GEOS_Length_r GEOSLength_r;
+  CR_GEOS_Centroid_r GEOSGetCentroid_r;
 
   CR_GEOS_Distance_r GEOSDistance_r;
 
@@ -166,6 +168,7 @@ struct CR_GEOS {
     INIT(GEOSGetSRID_r);
     INIT(GEOSArea_r);
     INIT(GEOSLength_r);
+    INIT(GEOSGetCentroid_r);
     INIT(GEOSDistance_r);
     INIT(GEOSCovers_r);
     INIT(GEOSCoveredBy_r);
@@ -364,6 +367,24 @@ CR_GEOS_Status CR_GEOS_Length(CR_GEOS* lib, CR_GEOS_Slice a, double *ret) {
   return CR_GEOS_UnaryOperator(lib, lib->GEOSLength_r, a, ret);
 }
 
+CR_GEOS_Status CR_GEOS_Centroid(CR_GEOS* lib, CR_GEOS_Slice a, CR_GEOS_String *centroidEWKB) {
+  std::string error;
+  auto handle = initHandleWithErrorBuffer(lib, &error);
+  auto geom = CR_GEOS_GeometryFromSlice(lib, handle, a);
+  *centroidEWKB = {.data = NULL, .len = 0};
+  if (geom != nullptr) {
+    auto centroidGeom = lib->GEOSGetCentroid_r(handle, geom);
+    if (centroidGeom != nullptr) {
+      auto srid = lib->GEOSGetSRID_r(handle, geom);
+      CR_GEOS_writeGeomToEWKB(lib, handle, centroidGeom, centroidEWKB, srid);
+      lib->GEOSGeom_destroy_r(handle, centroidGeom);
+    }
+    lib->GEOSGeom_destroy_r(handle, geom);
+  }
+  lib->GEOS_finish_r(handle);
+  return toGEOSString(error.data(), error.length());
+}
+
 CR_GEOS_Status CR_GEOS_Distance(CR_GEOS* lib, CR_GEOS_Slice a, CR_GEOS_Slice b, double *ret) {
   return CR_GEOS_BinaryOperator(lib, lib->GEOSDistance_r, a, b, ret);
 }

pkg/geo/geos/geos.go

@@ -254,6 +254,19 @@ func Length(ewkb geopb.EWKB) (float64, error) {
 	return float64(length), nil
 }
 
+// Centroid returns the centroid of an EWKB.
+func Centroid(ewkb geopb.EWKB) (geopb.EWKB, error) {
+	g, err := ensureInitInternal()
+	if err != nil {
+		return nil, err
+	}
+	var cEWKB C.CR_GEOS_String
+	if err := statusToError(C.CR_GEOS_Centroid(g, goToCSlice(ewkb), &cEWKB)); err != nil {
+		return nil, err
+	}
+	return cStringToSafeGoBytes(cEWKB), nil
+}
+
 // MinDistance returns the minimum distance between two EWKBs.
 func MinDistance(a geopb.EWKB, b geopb.EWKB) (float64, error) {
 	g, err := ensureInitInternal()

pkg/geo/geos/geos.h

@@ -63,6 +63,7 @@ CR_GEOS_Status CR_GEOS_ClipEWKBByRect(CR_GEOS* lib, CR_GEOS_Slice wkb, double xm
 
 CR_GEOS_Status CR_GEOS_Area(CR_GEOS* lib, CR_GEOS_Slice a, double *ret);
 CR_GEOS_Status CR_GEOS_Length(CR_GEOS* lib, CR_GEOS_Slice a, double *ret);
+CR_GEOS_Status CR_GEOS_Centroid(CR_GEOS* lib, CR_GEOS_Slice a, CR_GEOS_String *centroidEWKB);
 
 //
 // Binary operators.

pkg/sql/logictest/testdata/logic_test/geospatial

@@ -325,6 +325,43 @@ Square (left)                             1     0     4
 Square (right)                            1     0     4
 Square overlapping left and right square  1.1   0     4.2
 
+query TT
+SELECT
+  a.dsc,
+  ST_AsEWKT(ST_Centroid(a.geom))
+FROM geom_operators_test a
+ORDER BY a.dsc
+----
+Empty GeometryCollection                  POINT EMPTY
+Empty LineString                          POINT EMPTY
+Faraway point                             POINT (5 5)
+Line going through left and right square  POINT (0 0.5)
+NULL                                      NULL
+Point middle of Left Square               POINT (-0.5 0.5)
+Point middle of Right Square              POINT (0.5 0.5)
+Square (left)                             POINT (-0.5 0.5)
+Square (right)                            POINT (0.5 0.5)
+Square overlapping left and right square  POINT (0.4499999999999999 0.5)
+
+# Functions which take in strings as input as well.
+query TT
+SELECT
+  dsc,
+  ST_AsEWKT(ST_Centroid(ewkt))
+FROM [SELECT dsc, ST_AsEWKT(a.geom) ewkt FROM geom_operators_test a]
+ORDER BY dsc ASC
+----
+Empty GeometryCollection                  POINT EMPTY
+Empty LineString                          POINT EMPTY
+Faraway point                             POINT (5 5)
+Line going through left and right square  POINT (0 0.5)
+NULL                                      NULL
+Point middle of Left Square               POINT (-0.5 0.5)
+Point middle of Right Square              POINT (0.5 0.5)
+Square (left)                             POINT (-0.5 0.5)
+Square (right)                            POINT (0.5 0.5)
+Square overlapping left and right square  POINT (0.4499999999999999 0.5)
+
 # Binary operators
 query TTRR
 SELECT

pkg/sql/sem/builtins/geo_builtins.go

@@ -1241,6 +1241,43 @@ Note ST_Perimeter is only valid for Polygon - use ST_Length for LineString.`,
 			tree.VolatilityImmutable,
 		),
 	),
+	"st_centroid": makeBuiltin(
+		defProps(),
+		geometryOverload1(
+			func(ctx *tree.EvalContext, g *tree.DGeometry) (tree.Datum, error) {
+				centroid, err := geomfn.Centroid(g.Geometry)
+				if err != nil {
+					return nil, err
+				}
+				return tree.NewDGeometry(centroid), err
+			},
+			types.Geometry,
+			infoBuilder{
+				info:         "Returns the centroid of the given geometry.",
+				libraryUsage: usesGEOS,
+			},
+			tree.VolatilityImmutable,
+		),
+		stringOverload1(
+			func(ctx *tree.EvalContext, s string) (tree.Datum, error) {
+				g, err := geo.ParseGeometry(s)
+				if err != nil {
+					return nil, err
+				}
+				centroid, err := geomfn.Centroid(g)
+				if err != nil {
+					return nil, err
+				}
+				return tree.NewDGeometry(centroid), err
+			},
+			types.Geometry,
+			infoBuilder{
+				info:         "Returns the centroid of the given string, which will be parsed as a geometry object.",
+				libraryUsage: usesGEOS,
+			}.String(),
+			tree.VolatilityImmutable,
+		),
+	),
 
 	//
 	// Binary functions