sql: support specifying custom index parameters for geospatial indexes

pkg/geo/geoindex/geoindex.go

@@ -687,7 +687,8 @@ func (b *stringBuilderWithWrap) doWrap() {
 	b.lastWrap = b.Len()
 }
 
-func defaultS2Config() *S2Config {
+// DefaultS2Config returns the default S2Config to initialize.
+func DefaultS2Config() *S2Config {
 	return &S2Config{
 		MinLevel: 0,
 		MaxLevel: 30,

pkg/geo/geoindex/s2_geography_index.go

@@ -50,7 +50,7 @@ func NewS2GeographyIndex(cfg S2GeographyConfig) GeographyIndex {
 // DefaultGeographyIndexConfig returns a default config for a geography index.
 func DefaultGeographyIndexConfig() *Config {
 	return &Config{
-		S2Geography: &S2GeographyConfig{S2Config: defaultS2Config()},
+		S2Geography: &S2GeographyConfig{S2Config: DefaultS2Config()},
 	}
 }
 

pkg/geo/geoindex/s2_geometry_index.go

@@ -72,7 +72,7 @@ func DefaultGeometryIndexConfig() *Config {
 			MaxX:     10000,
 			MinY:     -10000,
 			MaxY:     10000,
-			S2Config: defaultS2Config()},
+			S2Config: DefaultS2Config()},
 	}
 }
 
@@ -122,7 +122,7 @@ func GeometryIndexConfigForSRID(srid geopb.SRID) (*Config, error) {
 			MaxX:     maxX + deltaX,
 			MinY:     minY - deltaY,
 			MaxY:     maxY + deltaY,
-			S2Config: defaultS2Config()},
+			S2Config: DefaultS2Config()},
 	}, nil
 }
 

pkg/sql/create_index.go

@@ -35,16 +35,89 @@ type createIndexNode struct {
 	tableDesc *sqlbase.MutableTableDescriptor
 }
 
-type indexStorageParamObserver struct{}
+type indexStorageParamObserver struct {
+	indexDesc *descpb.IndexDescriptor
+}
 
 var _ storageParamObserver = (*indexStorageParamObserver)(nil)
 
+func getS2ConfigFromIndex(indexDesc *descpb.IndexDescriptor) *geoindex.S2Config {
+	var s2Config *geoindex.S2Config
+	if indexDesc.GeoConfig.S2Geometry != nil {
+		s2Config = indexDesc.GeoConfig.S2Geometry.S2Config
+	}
+	if indexDesc.GeoConfig.S2Geography != nil {
+		s2Config = indexDesc.GeoConfig.S2Geography.S2Config
+	}
+	return s2Config
+}
+
+func (a *indexStorageParamObserver) applyS2ConfigSetting(
+	evalCtx *tree.EvalContext, key string, expr tree.Datum, min int64, max int64,
+) error {
+	s2Config := getS2ConfigFromIndex(a.indexDesc)
+	if s2Config == nil {
+		return errors.Newf("index setting %q can only be set on GEOMETRY or GEOGRAPHY spatial indexes", key)
+	}
+
+	val, err := datumAsInt(evalCtx, key, expr)
+	if err != nil {
+		return errors.Wrapf(err, "error decoding %q", key)
+	}
+	if val < min || val > max {
+		return errors.Newf("%q value must be between %d and %d inclusive", key, min, max)
+	}
+	switch key {
+	case `s2_max_level`:
+		s2Config.MaxLevel = int32(val)
+	case `s2_level_mod`:
+		s2Config.LevelMod = int32(val)
+	case `s2_max_cells`:
+		s2Config.MaxCells = int32(val)
+	}
+
+	return nil
+}
+
+func (a *indexStorageParamObserver) applyGeometryIndexSetting(
+	evalCtx *tree.EvalContext, key string, expr tree.Datum,
+) error {
+	if a.indexDesc.GeoConfig.S2Geometry == nil {
+		return errors.Newf("%q can only be applied to GEOMETRY spatial indexes", key)
+	}
+	val, err := datumAsFloat(evalCtx, key, expr)
+	if err != nil {
+		return errors.Wrapf(err, "error decoding %q", key)
+	}
+	switch key {
+	case `geometry_min_x`:
+		a.indexDesc.GeoConfig.S2Geometry.MinX = val
+	case `geometry_max_x`:
+		a.indexDesc.GeoConfig.S2Geometry.MaxX = val
+	case `geometry_min_y`:
+		a.indexDesc.GeoConfig.S2Geometry.MinY = val
+	case `geometry_max_y`:
+		a.indexDesc.GeoConfig.S2Geometry.MaxY = val
+	default:
+		return errors.Newf("unknown key: %q", key)
+	}
+	return nil
+}
+
 func (a *indexStorageParamObserver) apply(
 	evalCtx *tree.EvalContext, key string, expr tree.Datum,
 ) error {
 	switch key {
 	case `fillfactor`:
 		return applyFillFactorStorageParam(evalCtx, key, expr)
+	case `s2_max_level`:
+		return a.applyS2ConfigSetting(evalCtx, key, expr, 0, 30)
+	case `s2_level_mod`:
+		return a.applyS2ConfigSetting(evalCtx, key, expr, 1, 3)
+	case `s2_max_cells`:
+		return a.applyS2ConfigSetting(evalCtx, key, expr, 1, 32)
+	case `geometry_min_x`, `geometry_max_x`, `geometry_min_y`, `geometry_max_y`:
+		return a.applyGeometryIndexSetting(evalCtx, key, expr)
 	case `vacuum_cleanup_index_scale_factor`,
 		`buffering`,
 		`fastupdate`,
@@ -56,6 +129,37 @@ func (a *indexStorageParamObserver) apply(
 	return errors.Errorf("invalid storage parameter %q", key)
 }
 
+func (a *indexStorageParamObserver) runPostChecks() error {
+	s2Config := getS2ConfigFromIndex(a.indexDesc)
+	if s2Config != nil {
+		if (s2Config.MaxLevel)%s2Config.LevelMod != 0 {
+			return errors.Newf(
+				"s2_max_level (%d) must be divisible by s2_level_mod (%d)",
+				s2Config.MaxLevel,
+				s2Config.LevelMod,
+			)
+		}
+	}
+
+	if cfg := a.indexDesc.GeoConfig.S2Geometry; cfg != nil {
+		if cfg.MaxX <= cfg.MinX {
+			return errors.Newf(
+				"geometry_max_x (%f) must be greater than geometry_min_x (%f)",
+				cfg.MaxX,
+				cfg.MinX,
+			)
+		}
+		if cfg.MaxY <= cfg.MinY {
+			return errors.Newf(
+				"geometry_max_y (%f) must be greater than geometry_min_y (%f)",
+				cfg.MaxY,
+				cfg.MinY,
+			)
+		}
+	}
+	return nil
+}
+
 // CreateIndex creates an index.
 // Privileges: CREATE on table.
 //   notes: postgres requires CREATE on the table.
@@ -251,7 +355,7 @@ func MakeIndexDescriptor(
 		params.p.SemaCtx(),
 		params.EvalContext(),
 		n.StorageParams,
-		&indexStorageParamObserver{},
+		&indexStorageParamObserver{indexDesc: &indexDesc},
 	); err != nil {
 		return nil, err
 	}

pkg/sql/create_table.go

@@ -76,6 +76,7 @@ type createTableRun struct {
 // storageParamObserver applies a storage parameter to an underlying item.
 type storageParamObserver interface {
 	apply(evalCtx *tree.EvalContext, key string, datum tree.Datum) error
+	runPostChecks() error
 }
 
 type tableStorageParamObserver struct{}
@@ -99,6 +100,10 @@ func applyFillFactorStorageParam(evalCtx *tree.EvalContext, key string, datum tr
 	return nil
 }
 
+func (a *tableStorageParamObserver) runPostChecks() error {
+	return nil
+}
+
 func (a *tableStorageParamObserver) apply(
 	evalCtx *tree.EvalContext, key string, datum tree.Datum,
 ) error {
@@ -1542,6 +1547,15 @@ func MakeTableDesc(
 				}
 				idx.Predicate = expr
 			}
+			if err := applyStorageParameters(
+				ctx,
+				semaCtx,
+				evalCtx,
+				d.StorageParams,
+				&indexStorageParamObserver{indexDesc: &idx},
+			); err != nil {
+				return desc, err
+			}
 
 			if err := desc.AddIndex(idx, false); err != nil {
 				return desc, err
@@ -1852,7 +1866,7 @@ func applyStorageParameters(
 			return err
 		}
 	}
-	return nil
+	return paramObserver.runPostChecks()
 }
 
 // makeTableDesc creates a table descriptor from a CreateTable statement.

pkg/sql/logictest/testdata/logic_test/geospatial_index

@@ -0,0 +1,89 @@
+statement ok
+CREATE TABLE geo_table(
+  id int primary key,
+  geog geography(geometry, 4326),
+  geom geometry(geometry, 3857),
+  FAMILY fam_0_geog (geog),
+  FAMILY fam_1_geom (geom),
+  FAMILY fam_2_id (id)
+)
+
+statement error index setting "s2_max_cells" can only be set on GEOMETRY or GEOGRAPHY spatial indexes
+CREATE INDEX bad_idx ON geo_table(id) WITH (s2_max_cells=15)
+
+statement error "s2_max_cells" value must be between 1 and 32 inclusive
+CREATE INDEX bad_idx ON geo_table USING GIST(geom) WITH (s2_max_cells=42)
+
+statement error s2_max_level \(29\) must be divisible by s2_level_mod \(2\)
+CREATE INDEX bad_idx ON geo_table USING GIST(geom) WITH (s2_max_level=29, s2_level_mod=2)
+
+statement error "geometry_min_x" can only be applied to GEOMETRY spatial indexes
+CREATE INDEX bad_idx ON geo_table USING GIST(geog) WITH (geometry_min_x=0)
+
+statement error geometry_max_x \(0\.000000\) must be greater than geometry_min_x \(10\.000000\)
+CREATE INDEX bad_idx ON geo_table USING GIST(geom) WITH (geometry_min_x=10, geometry_max_x=0)
+
+statement error geometry_max_y \(0\.000000\) must be greater than geometry_min_y \(10\.000000\)
+CREATE INDEX bad_idx ON geo_table USING GIST(geom) WITH (geometry_min_y=10, geometry_max_y=0)
+
+statement ok
+CREATE INDEX geom_idx_1 ON geo_table USING GIST(geom) WITH (geometry_min_x=0, s2_max_level=15)
+
+statement ok
+CREATE INDEX geom_idx_2 ON geo_table USING GIST(geom) WITH (geometry_min_x=0)
+
+statement ok
+CREATE INDEX geom_idx_3 ON geo_table USING GIST(geom) WITH (s2_max_level=10)
+
+statement ok
+CREATE INDEX geom_idx_4 ON geo_table USING GIST(geom)
+
+statement ok
+CREATE INDEX geog_idx_1 ON geo_table USING GIST(geog) WITH (s2_level_mod=2)
+
+statement ok
+CREATE INDEX geog_idx_2 ON geo_table USING GIST(geog)
+
+query T
+SELECT create_statement FROM [SHOW CREATE TABLE geo_table]
+----
+CREATE TABLE public.geo_table (
+   id INT8 NOT NULL,
+   geog GEOGRAPHY(GEOMETRY,4326) NULL,
+   geom GEOMETRY(GEOMETRY,3857) NULL,
+   CONSTRAINT "primary" PRIMARY KEY (id ASC),
+   INVERTED INDEX geom_idx_1 (geom) WITH (s2_max_level=15, geometry_min_x=0),
+   INVERTED INDEX geom_idx_2 (geom) WITH (geometry_min_x=0),
+   INVERTED INDEX geom_idx_3 (geom) WITH (s2_max_level=10),
+   INVERTED INDEX geom_idx_4 (geom),
+   INVERTED INDEX geog_idx_1 (geog) WITH (s2_level_mod=2),
+   INVERTED INDEX geog_idx_2 (geog),
+   FAMILY fam_0_geog (geog),
+   FAMILY fam_1_geom (geom),
+   FAMILY fam_2_id (id)
+)
+
+let $create_table
+SELECT create_statement FROM [SHOW CREATE TABLE geo_table]
+
+statement ok
+DROP TABLE geo_table; $create_table
+
+query T
+SELECT create_statement FROM [SHOW CREATE TABLE geo_table]
+----
+CREATE TABLE public.geo_table (
+   id INT8 NOT NULL,
+   geog GEOGRAPHY(GEOMETRY,4326) NULL,
+   geom GEOMETRY(GEOMETRY,3857) NULL,
+   CONSTRAINT "primary" PRIMARY KEY (id ASC),
+   INVERTED INDEX geom_idx_1 (geom) WITH (s2_max_level=15, geometry_min_x=0),
+   INVERTED INDEX geom_idx_2 (geom) WITH (geometry_min_x=0),
+   INVERTED INDEX geom_idx_3 (geom) WITH (s2_max_level=10),
+   INVERTED INDEX geom_idx_4 (geom),
+   INVERTED INDEX geog_idx_1 (geog) WITH (s2_level_mod=2),
+   INVERTED INDEX geog_idx_2 (geog),
+   FAMILY fam_0_geog (geog),
+   FAMILY fam_1_geom (geom),
+   FAMILY fam_2_id (id)
+)

pkg/sql/parser/sql.y

@@ -5650,39 +5650,42 @@ generated_as:
 
 
 index_def:
-  INDEX opt_index_name '(' index_params ')' opt_hash_sharded opt_storing opt_interleave opt_partition_by opt_where_clause
+  INDEX opt_index_name '(' index_params ')' opt_hash_sharded opt_storing opt_interleave opt_partition_by opt_with_storage_parameter_list opt_where_clause
   {
     $$.val = &tree.IndexTableDef{
-      Name:    tree.Name($2),
-      Columns: $4.idxElems(),
-      Sharded: $6.shardedIndexDef(),
-      Storing: $7.nameList(),
-      Interleave: $8.interleave(),
-      PartitionBy: $9.partitionBy(),
-      Predicate: $10.expr(),
+      Name:          tree.Name($2),
+      Columns:       $4.idxElems(),
+      Sharded:       $6.shardedIndexDef(),
+      Storing:       $7.nameList(),
+      Interleave:    $8.interleave(),
+      PartitionBy:   $9.partitionBy(),
+      StorageParams: $10.storageParams(),
+      Predicate:     $11.expr(),
     }
   }
-| UNIQUE INDEX opt_index_name '(' index_params ')' opt_hash_sharded opt_storing opt_interleave opt_partition_by opt_where_clause
+| UNIQUE INDEX opt_index_name '(' index_params ')' opt_hash_sharded opt_storing opt_interleave opt_partition_by opt_with_storage_parameter_list opt_where_clause
   {
     $$.val = &tree.UniqueConstraintTableDef{
       IndexTableDef: tree.IndexTableDef {
-        Name:    tree.Name($3),
-        Columns: $5.idxElems(),
-        Sharded: $7.shardedIndexDef(),
-        Storing: $8.nameList(),
-        Interleave: $9.interleave(),
-        PartitionBy: $10.partitionBy(),
-        Predicate: $11.expr(),
+        Name:          tree.Name($3),
+        Columns:       $5.idxElems(),
+        Sharded:       $7.shardedIndexDef(),
+        Storing:       $8.nameList(),
+        Interleave:    $9.interleave(),
+        PartitionBy:   $10.partitionBy(),
+        StorageParams: $11.storageParams(),
+        Predicate:     $12.expr(),
       },
     }
   }
-| INVERTED INDEX opt_name '(' index_params ')' opt_where_clause
+| INVERTED INDEX opt_name '(' index_params ')' opt_with_storage_parameter_list opt_where_clause
   {
     $$.val = &tree.IndexTableDef{
-      Name:    tree.Name($3),
-      Columns: $5.idxElems(),
-      Inverted: true,
-      Predicate: $7.expr(),
+      Name:          tree.Name($3),
+      Columns:       $5.idxElems(),
+      Inverted:      true,
+      StorageParams: $7.storageParams(),
+      Predicate:     $8.expr(),
     }
   }