Merge #60832

60832: geo: add support for storing spatial types with Z and M dimensions r=otan a=andyyang890 Previously, CockroachDB only supported storing spatial types with X and Y dimensions (2D). This patch enhances CockroachDB with support for storing Z and M (the 3rd and 4th dimensions). Based on #53089. Resolves #53091. Release note (sql change): CockroachDB now supports storing spatial type objects with Z and M dimensions (e.g. POINTZ, LINESTRINGM). Co-authored-by: Andy Yang <[email protected]>
cockroachdb · Feb 21, 2021 · 682582f · 682582f
2 parents d485973 + 54a6ac4
commit 682582f
Show file tree

Hide file tree

Showing 16 changed files with 450 additions and 143 deletions.
diff --git a/docs/generated/sql/bnf/stmt_block.bnf b/docs/generated/sql/bnf/stmt_block.bnf
@@ -2982,13 +2982,37 @@ opt_timezone ::=
 
 geo_shape_type ::=
 	'POINT'
+	| 'POINTM'
+	| 'POINTZ'
+	| 'POINTZM'
 	| 'LINESTRING'
+	| 'LINESTRINGM'
+	| 'LINESTRINGZ'
+	| 'LINESTRINGZM'
 	| 'POLYGON'
+	| 'POLYGONM'
+	| 'POLYGONZ'
+	| 'POLYGONZM'
 	| 'MULTIPOINT'
+	| 'MULTIPOINTM'
+	| 'MULTIPOINTZ'
+	| 'MULTIPOINTZM'
 	| 'MULTILINESTRING'
+	| 'MULTILINESTRINGM'
+	| 'MULTILINESTRINGZ'
+	| 'MULTILINESTRINGZM'
 	| 'MULTIPOLYGON'
+	| 'MULTIPOLYGONM'
+	| 'MULTIPOLYGONZ'
+	| 'MULTIPOLYGONZM'
 	| 'GEOMETRYCOLLECTION'
+	| 'GEOMETRYCOLLECTIONM'
+	| 'GEOMETRYCOLLECTIONZ'
+	| 'GEOMETRYCOLLECTIONZM'
 	| 'GEOMETRY'
+	| 'GEOMETRYM'
+	| 'GEOMETRYZ'
+	| 'GEOMETRYZM'
 
 char_aliases ::=
 	'CHAR'

diff --git a/pkg/BUILD.bazel b/pkg/BUILD.bazel
@@ -60,6 +60,7 @@ ALL_TESTS = [
     "//pkg/geo/geographiclib:geographiclib_test",
     "//pkg/geo/geoindex:geoindex_test",
     "//pkg/geo/geomfn:geomfn_test",
+    "//pkg/geo/geopb:geopb_test",
     "//pkg/geo/geoproj:geoproj_test",
     "//pkg/geo/geoprojbase:geoprojbase_test",
     "//pkg/geo/geos:geos_test",

diff --git a/pkg/geo/BUILD.bazel b/pkg/geo/BUILD.bazel
@@ -22,7 +22,6 @@ go_library(
         "//pkg/geo/geoprojbase",
         "//pkg/geo/wkt",
         "//pkg/util",
-        "//pkg/util/errorutil/unimplemented",
         "//pkg/util/protoutil",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_golang_geo//r1",

diff --git a/pkg/geo/geo.go b/pkg/geo/geo.go
@@ -20,7 +20,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/geo/geographiclib"
 	"github.com/cockroachdb/cockroach/pkg/geo/geopb"
 	"github.com/cockroachdb/cockroach/pkg/geo/geoprojbase"
-	"github.com/cockroachdb/cockroach/pkg/util/errorutil/unimplemented"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 	"github.com/cockroachdb/errors"
 	"github.com/golang/geo/r1"
@@ -71,14 +70,39 @@ func SpatialObjectFitsColumnMetadata(
 	if srid != 0 && so.SRID != srid {
 		return errors.Newf("object SRID %d does not match column SRID %d", so.SRID, srid)
 	}
-	// Shape_Geometry/Shape_Unset can take in any kind of shape.
+	// Shape_Unset can take in any kind of shape.
+	// Shape_Geometry[ZM] must match dimensions.
 	// Otherwise, shapes must match.
-	if shapeType != geopb.ShapeType_Unset && shapeType != geopb.ShapeType_Geometry && shapeType != so.ShapeType {
-		return errors.Newf("object type %s does not match column type %s", so.ShapeType, shapeType)
+	switch shapeType {
+	case geopb.ShapeType_Unset:
+		break
+	case geopb.ShapeType_Geometry, geopb.ShapeType_GeometryM, geopb.ShapeType_GeometryZ, geopb.ShapeType_GeometryZM:
+		if ShapeTypeToLayout(shapeType) != ShapeTypeToLayout(so.ShapeType) {
+			return errors.Newf("object type %s does not match column dimensionality %s", so.ShapeType, shapeType)
+		}
+	default:
+		if shapeType != so.ShapeType {
+			return errors.Newf("object type %s does not match column type %s", so.ShapeType, shapeType)
+		}
 	}
 	return nil
 }
 
+// ShapeTypeToLayout returns the geom.Layout of the given ShapeType.
+// Note this is not a definition on ShapeType to prevent geopb from importing twpayne/go-geom.
+func ShapeTypeToLayout(s geopb.ShapeType) geom.Layout {
+	switch {
+	case (s&geopb.MShapeTypeFlag > 0) && (s&geopb.ZShapeTypeFlag > 0):
+		return geom.XYZM
+	case s&geopb.ZShapeTypeFlag > 0:
+		return geom.XYZ
+	case s&geopb.MShapeTypeFlag > 0:
+		return geom.XYM
+	default:
+		return geom.XY
+	}
+}
+
 //
 // Geometry
 //
@@ -859,17 +883,6 @@ func spatialObjectFromGeomT(t geom.T, soType geopb.SpatialObjectType) (geopb.Spa
 	if err != nil {
 		return geopb.SpatialObject{}, err
 	}
-	switch t.Layout() {
-	case geom.XY:
-	case geom.NoLayout:
-		if gc, ok := t.(*geom.GeometryCollection); !ok || !gc.Empty() {
-			return geopb.SpatialObject{}, errors.Newf("no layout found on object")
-		}
-	case geom.XYM, geom.XYZ, geom.XYZM:
-		return geopb.SpatialObject{}, unimplemented.NewWithIssueDetailf(53091, t.Layout().String()+"_datum", "dimension %s is not currently supported", t.Layout())
-	default:
-		return geopb.SpatialObject{}, errors.Newf("unexpected layout: %s", geom.XY)
-	}
 	bbox, err := boundingBoxFromGeomT(t, soType)
 	if err != nil {
 		return geopb.SpatialObject{}, err
@@ -884,24 +897,42 @@ func spatialObjectFromGeomT(t geom.T, soType geopb.SpatialObjectType) (geopb.Spa
 }
 
 func shapeTypeFromGeomT(t geom.T) (geopb.ShapeType, error) {
+	var shapeType geopb.ShapeType
 	switch t := t.(type) {
 	case *geom.Point:
-		return geopb.ShapeType_Point, nil
+		shapeType = geopb.ShapeType_Point
 	case *geom.LineString:
-		return geopb.ShapeType_LineString, nil
+		shapeType = geopb.ShapeType_LineString
 	case *geom.Polygon:
-		return geopb.ShapeType_Polygon, nil
+		shapeType = geopb.ShapeType_Polygon
 	case *geom.MultiPoint:
-		return geopb.ShapeType_MultiPoint, nil
+		shapeType = geopb.ShapeType_MultiPoint
 	case *geom.MultiLineString:
-		return geopb.ShapeType_MultiLineString, nil
+		shapeType = geopb.ShapeType_MultiLineString
 	case *geom.MultiPolygon:
-		return geopb.ShapeType_MultiPolygon, nil
+		shapeType = geopb.ShapeType_MultiPolygon
 	case *geom.GeometryCollection:
-		return geopb.ShapeType_GeometryCollection, nil
+		shapeType = geopb.ShapeType_GeometryCollection
 	default:
 		return geopb.ShapeType_Unset, errors.Newf("unknown shape: %T", t)
 	}
+	switch t.Layout() {
+	case geom.NoLayout:
+		if gc, ok := t.(*geom.GeometryCollection); !ok || !gc.Empty() {
+			return geopb.ShapeType_Unset, errors.Newf("no layout found on object")
+		}
+	case geom.XY:
+		break
+	case geom.XYM:
+		shapeType = shapeType | geopb.MShapeTypeFlag
+	case geom.XYZ:
+		shapeType = shapeType | geopb.ZShapeTypeFlag
+	case geom.XYZM:
+		shapeType = shapeType | geopb.ZShapeTypeFlag | geopb.MShapeTypeFlag
+	default:
+		return geopb.ShapeType_Unset, errors.Newf("unknown layout: %s", t.Layout())
+	}
+	return shapeType, nil
 }
 
 // GeomTContainsEmpty returns whether a geom.T contains any empty element.

diff --git a/pkg/geo/geo_test.go b/pkg/geo/geo_test.go
@@ -104,6 +104,9 @@ func TestSpatialObjectFitsColumnMetadata(t *testing.T) {
 		errorContains string
 	}{
 		{MustParseGeometry("POINT(1.0 1.0)"), 0, geopb.ShapeType_Geometry, ""},
+		{MustParseGeometry("POINT Z(1.0 1.0 1.0)"), 0, geopb.ShapeType_GeometryZ, ""},
+		{MustParseGeometry("POINT ZM(1.0 1.0 1.0 1.0)"), 0, geopb.ShapeType_GeometryZM, ""},
+		{MustParseGeometry("POINT M(1.0 1.0 1.0)"), 0, geopb.ShapeType_GeometryM, ""},
 		{MustParseGeometry("POINT(1.0 1.0)"), 0, geopb.ShapeType_Unset, ""},
 		{MustParseGeometry("SRID=4326;POINT(1.0 1.0)"), 0, geopb.ShapeType_Geometry, ""},
 		{MustParseGeometry("SRID=4326;POINT(1.0 1.0)"), 0, geopb.ShapeType_Unset, ""},

diff --git a/pkg/geo/geopb/BUILD.bazel b/pkg/geo/geopb/BUILD.bazel
@@ -1,6 +1,6 @@
 load("@rules_proto//proto:defs.bzl", "proto_library")
 load("@io_bazel_rules_go//proto:def.bzl", "go_proto_library")
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "geopb",
@@ -29,3 +29,10 @@ go_proto_library(
     visibility = ["//visibility:public"],
     deps = ["@com_github_gogo_protobuf//gogoproto"],
 )
+
+go_test(
+    name = "geopb_test",
+    srcs = ["types_test.go"],
+    embed = [":geopb"],
+    deps = ["@com_github_stretchr_testify//require"],
+)