Merge #139450

139450: builtins: implement ST_3DLength r=yuzefovich a=qburke

This patch allows measuring the length of 3D geometries when there was no builtin function to do this before.

Fixes: #60866

Release note (sql change): The `ST_3DLength` function is now available for use.

Co-authored-by: Quin Burke <[email protected]>

Author: craig[bot]

docs/generated/sql/functions.md

@@ -1711,6 +1711,11 @@ the locality flag on node startup. Returns an error if no region is set.</p>
 </span></td><td>Immutable</td></tr>
 <tr><td><a name="postgis_wagyu_version"></a><code>postgis_wagyu_version() &rarr; <a href="string.html">string</a></code></td><td><span class="funcdesc"><p>Compatibility placeholder function with PostGIS. Returns a fixed string based on PostGIS 3.0.1, with minor edits.</p>
 </span></td><td>Immutable</td></tr>
+<tr><td><a name="st_3dlength"></a><code>st_3dlength(geometry: geometry) &rarr; <a href="float.html">float</a></code></td><td><span class="funcdesc"><p>Returns the 3-dimensional or 2-dimensional length of the geometry.</p>
+<p>Note ST_3DLength is only valid for LineString or MultiLineString.
+For 2-D lines it will return the 2-D length (same as ST_Length and ST_Length2D)</p>
+<p>This function utilizes the GEOS module.</p>
+</span></td><td>Immutable</td></tr>
 <tr><td><a name="st_addmeasure"></a><code>st_addmeasure(geometry: geometry, start: <a href="float.html">float</a>, end: <a href="float.html">float</a>) &rarr; geometry</code></td><td><span class="funcdesc"><p>Returns a copy of a LineString or MultiLineString with measure coordinates linearly interpolated between the specified start and end values. Any existing M coordinates will be overwritten.</p>
 </span></td><td>Immutable</td></tr>
 <tr><td><a name="st_addpoint"></a><code>st_addpoint(line_string: geometry, point: geometry) &rarr; geometry</code></td><td><span class="funcdesc"><p>Adds a Point to the end of a LineString.</p>

pkg/geo/geomfn/unary_operators.go

@@ -6,6 +6,8 @@
 package geomfn
 
 import (
+	"math"
+
 	"github.com/cockroachdb/cockroach/pkg/geo"
 	"github.com/cockroachdb/cockroach/pkg/geo/geos"
 	"github.com/cockroachdb/errors"
@@ -242,3 +244,86 @@ func CountVertices(t geom.T) int {
 		return len(t.FlatCoords()) / t.Stride()
 	}
 }
+
+// length3DLineString returns the length of a
+// given 3D LINESTRING. Returns an error if
+// lineString does not have a z-coordinate.
+func length3DLineString(lineString *geom.LineString) (float64, error) {
+	lineCoords := lineString.Coords()
+	prevPoint := lineCoords[0]
+	lineLength := float64(0)
+	zIndex := lineString.Layout().ZIndex()
+	if zIndex < 0 || zIndex >= lineString.Stride() {
+		return 0, errors.AssertionFailedf("Z-Index for LINESTRING is out-of-bounds")
+	}
+	for i := 0; i < lineString.NumCoords(); i++ {
+		curPoint := lineCoords[i]
+		deltaX := curPoint.X() - prevPoint.X()
+		deltaY := curPoint.Y() - prevPoint.Y()
+		deltaZ := curPoint[zIndex] - prevPoint[zIndex]
+		distBetweenPoints := math.Sqrt(deltaX*deltaX + deltaY*deltaY + deltaZ*deltaZ)
+		lineLength += distBetweenPoints
+		prevPoint = curPoint
+	}
+
+	return lineLength, nil
+}
+
+// length3DMultiLineString returns the length of a
+// given 3D MULTILINESTRING. Returns an error if
+// multiLineString is not 3D.
+func length3DMultiLineString(multiLineString *geom.MultiLineString) (float64, error) {
+	multiLineLength := 0.0
+	for i := 0; i < multiLineString.NumLineStrings(); i++ {
+		lineLength, err := length3DLineString(multiLineString.LineString(i))
+		if err != nil {
+			return 0, err
+		}
+		multiLineLength += lineLength
+	}
+
+	return multiLineLength, nil
+}
+
+// Length3D returns the length of a given Geometry.
+// Compatible with 3D geometries.
+// Note only (MULTI)LINESTRING objects have a length.
+func Length3D(g geo.Geometry) (float64, error) {
+	geomRepr, err := g.AsGeomT()
+	if err != nil {
+		return 0, err
+	}
+
+	switch geomRepr.Layout() {
+	case geom.XYZ, geom.XYZM:
+		return length3DFromGeomT(geomRepr)
+	}
+	// Call default length
+	return lengthFromGeomT(geomRepr)
+}
+
+// length3DFromGeomT returns the length from a geom.T, recursing down
+// GeometryCollections if required.
+// Compatible with 3D geometries.
+func length3DFromGeomT(geomRepr geom.T) (float64, error) {
+	switch geomRepr := geomRepr.(type) {
+	case *geom.Point, *geom.MultiPoint, *geom.Polygon, *geom.MultiPolygon:
+		return 0, nil
+	case *geom.LineString:
+		return length3DLineString(geomRepr)
+	case *geom.MultiLineString:
+		return length3DMultiLineString(geomRepr)
+	case *geom.GeometryCollection:
+		total := float64(0)
+		for _, subG := range geomRepr.Geoms() {
+			subLength, err := length3DFromGeomT(subG)
+			if err != nil {
+				return 0, err
+			}
+			total += subLength
+		}
+		return total, nil
+	default:
+		return 0, errors.AssertionFailedf("unknown geometry type: %T", geomRepr)
+	}
+}

pkg/geo/geomfn/unary_operators_test.go

@@ -252,3 +252,34 @@ func TestMinimumClearanceLine(t *testing.T) {
 		})
 	}
 }
+
+func TestLength3D(t *testing.T) {
+	testCases := []struct {
+		wkt      string
+		expected float64
+	}{
+		{"POINT(1.0 1.0)", 0},
+		{"POINT(1.0 1.0 1.0)", 0},
+		{"POLYGON((0.0 0.0, 1.0 0.0, 1.0 1.0, 0.0 0.0))", 0},
+		{"POLYGON((0.0 0.0 0.0, 1.0 0.0 0.0, 1.0 1.0 0.0, 1.0 1.0 1.0, 0.0 0.0 0.0))", 0},
+		{"LINESTRING(1.0 1.0, 2.0 2.0, 3.0 3.0)", 2.8284271247461903},
+		{"MULTILINESTRING((1.0 1.0, 2.0 2.0, 3.0 3.0), (6.0 6.0, 7.0 6.0))", 3.8284271247461903},
+		{"GEOMETRYCOLLECTION (POINT (40 10),LINESTRING (10 10, 20 20, 10 40),POLYGON ((40 40, 20 45, 45 30, 40 40)))", 36.50281539872885},
+		{"GEOMETRYCOLLECTION (GEOMETRYCOLLECTION(POINT (40 10),LINESTRING (10 10, 20 20, 10 40),POLYGON ((40 40, 20 45, 45 30, 40 40))))", 36.50281539872885},
+		{"LINESTRING(743238 2967416 1,743238 2967450 1,743265 2967450 3, 743265.625 2967416 3,743238 2967416 3)", 122.70471674145682},
+		{"LINESTRING(0 0 0, 1 1 1, 0 0 0)", 3.4641016151377544},
+		{"MULTILINESTRING((0 0 0, 1 1 1, 0 0 0), (743238 2967416 1,743238 2967450 1,743265 2967450 3, 743265.625 2967416 3,743238 2967416 3))", 126.16881835659457},
+		{"GEOMETRYCOLLECTION (POINT (0 0 0), LINESTRING(0 0 0, 1 1 1, 0 0 0), POLYGON((0.0 0.0 0.0, 1.0 0.0 0.0, 1.0 1.0 0.0, 1.0 1.0 1.0, 0.0 0.0 0.0)))", 3.4641016151377544},
+		{"GEOMETRYCOLLECTION (GEOMETRYCOLLECTION(POINT (0 0 0), LINESTRING(0 0 0, 1 1 1, 0 0 0), POLYGON((0.0 0.0 0.0, 1.0 0.0 0.0, 1.0 1.0 0.0, 1.0 1.0 1.0, 0.0 0.0 0.0))))", 3.4641016151377544},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.wkt, func(t *testing.T) {
+			g, err := geo.ParseGeometry(tc.wkt)
+			require.NoError(t, err)
+			ret, err := Length3D(g)
+			require.NoError(t, err)
+			require.Equal(t, tc.expected, ret)
+		})
+	}
+}

pkg/sql/logictest/testdata/logic_test/geospatial_zm

@@ -383,6 +383,11 @@ SELECT st_length('LINESTRING M(0 0 -25, 1 1 -50, 2 2 0)')
 ----
 2.8284271247461903
 
+query R
+SELECT ST_3DLength('LINESTRING(743238 2967416 1,743238 2967450 1,743265 2967450 3, 743265.625 2967416 3,743238 2967416 3)')
+----
+122.70471674145682
+
 query TTT nosort
 SELECT
   encode(ST_AsTWKB(t, 5), 'hex'),

pkg/sql/sem/builtins/fixed_oids.go

@@ -2652,6 +2652,7 @@ var builtinOidsArray = []string{
 	2689: `jsonb_path_exists(target: jsonb, path: jsonpath) -> bool`,
 	2690: `jsonb_path_exists(target: jsonb, path: jsonpath, vars: jsonb) -> bool`,
 	2691: `jsonb_path_exists(target: jsonb, path: jsonpath, vars: jsonb, silent: bool) -> bool`,
+	2692: `st_3dlength(geometry: geometry) -> float`,
 }
 
 var builtinOidsBySignature map[string]oid.Oid

pkg/sql/sem/builtins/geo_builtins.go

@@ -258,6 +258,25 @@ Note ST_Length is only valid for LineString - use ST_Perimeter for Polygon.`,
 	volatility.Immutable,
 )
 
+var length3DOverloadGeometry1 = geometryOverload1(
+	func(_ context.Context, _ *eval.Context, g *tree.DGeometry) (tree.Datum, error) {
+		ret, err := geomfn.Length3D(g.Geometry)
+		if err != nil {
+			return nil, err
+		}
+		return tree.NewDFloat(tree.DFloat(ret)), nil
+	},
+	types.Float,
+	infoBuilder{
+		info: `Returns the 3-dimensional or 2-dimensional length of the geometry.
+
+Note ST_3DLength is only valid for LineString or MultiLineString.
+For 2-D lines it will return the 2-D length (same as ST_Length and ST_Length2D)`,
+		libraryUsage: usesGEOS,
+	},
+	volatility.Immutable,
+)
+
 var perimeterOverloadGeometry1 = geometryOverload1(
 	func(_ context.Context, _ *eval.Context, g *tree.DGeometry) (tree.Datum, error) {
 		ret, err := geomfn.Perimeter(g.Geometry)
@@ -3126,6 +3145,10 @@ The requested number of points must be not larger than 65336.`,
 		defProps(),
 		lengthOverloadGeometry1,
 	),
+	"st_3dlength": makeBuiltin(
+		defProps(),
+		length3DOverloadGeometry1,
+	),
 	"st_perimeter": makeBuiltin(
 		defProps(),
 		append(