tree: correct sizes for geospatial datums

pkg/geo/geopb/geopb.go

@@ -10,13 +10,25 @@
 
 package geopb
 
-import "fmt"
+import (
+	"fmt"
+	"unsafe"
+)
 
 // EWKBHex returns the EWKB-hex version of this data type
 func (b *SpatialObject) EWKBHex() string {
 	return fmt.Sprintf("%X", b.EWKB)
 }
 
+// MemSize returns the size of the spatial object in memory.
+func (b *SpatialObject) MemSize() uintptr {
+	var bboxSize uintptr
+	if bbox := b.BoundingBox; bbox != nil {
+		bboxSize = unsafe.Sizeof(*bbox)
+	}
+	return unsafe.Sizeof(*b) + bboxSize + uintptr(len(b.EWKB))
+}
+
 // MultiType returns the corresponding multi-type for a shape type, or unset
 // if there is no multi-type.
 func (s ShapeType) MultiType() ShapeType {

pkg/sql/sem/tree/datum.go

@@ -2826,7 +2826,7 @@ func (d *DGeography) Format(ctx *FmtCtx) {
 
 // Size implements the Datum interface.
 func (d *DGeography) Size() uintptr {
-	return unsafe.Sizeof(*d)
+	return d.Geography.SpatialObjectRef().MemSize()
 }
 
 // DGeometry is the Geometry Datum.
@@ -2934,7 +2934,7 @@ func (d *DGeometry) Format(ctx *FmtCtx) {
 
 // Size implements the Datum interface.
 func (d *DGeometry) Size() uintptr {
-	return unsafe.Sizeof(*d)
+	return d.Geometry.SpatialObjectRef().MemSize()
 }
 
 // DBox2D is the Datum representation of the Box2D type.

pkg/sql/sem/tree/datum_test.go

@@ -1128,3 +1128,29 @@ var _ tree.ParseTimeContext = testParseTimeContext{}
 func (t testParseTimeContext) GetRelativeParseTime() time.Time {
 	return time.Time(t)
 }
+
+func TestGeospatialSize(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	testCases := []struct {
+		wkt      string
+		expected uintptr
+	}{
+		{"SRID=4004;POINT EMPTY", 73},
+		{"SRID=4326;LINESTRING(0 0, 10 0)", 125},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.wkt, func(t *testing.T) {
+			t.Run("geometry", func(t *testing.T) {
+				g, err := tree.ParseDGeometry(tc.wkt)
+				require.NoError(t, err)
+				require.Equal(t, tc.expected, g.Size())
+			})
+			t.Run("geography", func(t *testing.T) {
+				g, err := tree.ParseDGeography(tc.wkt)
+				require.NoError(t, err)
+				require.Equal(t, tc.expected, g.Size())
+			})
+		})
+	}
+}