github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/geo/geoindex/s2_geography_index.go (about)

     1  // Copyright 2020 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package geoindex
    12  
    13  import (
    14  	"context"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/geo"
    17  	"github.com/golang/geo/s2"
    18  )
    19  
    20  // s2GeographyIndex is an implementation of GeographyIndex that uses the S2 geometry
    21  // library.
    22  type s2GeographyIndex struct {
    23  	rc *s2.RegionCoverer
    24  }
    25  
    26  var _ GeographyIndex = (*s2GeographyIndex)(nil)
    27  
    28  // NewS2GeographyIndex returns an index with the given configuration. The
    29  // configuration of an index cannot be changed without rewriting the index
    30  // since deletes could miss some index entries. Currently, reads could use a
    31  // different configuration, but that is subject to change if we manage to
    32  // strengthen the covering invariants (see the todo in covers() in index.go).
    33  func NewS2GeographyIndex(cfg S2GeographyConfig) GeographyIndex {
    34  	// TODO(sumeer): Sanity check cfg.
    35  	return &s2GeographyIndex{
    36  		rc: &s2.RegionCoverer{
    37  			MinLevel: int(cfg.S2Config.MinLevel),
    38  			MaxLevel: int(cfg.S2Config.MaxLevel),
    39  			LevelMod: int(cfg.S2Config.LevelMod),
    40  			MaxCells: int(cfg.S2Config.MaxCells),
    41  		},
    42  	}
    43  }
    44  
    45  // DefaultGeographyIndexConfig returns a default config for a geography index.
    46  func DefaultGeographyIndexConfig() *Config {
    47  	return &Config{
    48  		S2Geography: &S2GeographyConfig{S2Config: defaultS2Config()},
    49  	}
    50  }
    51  
    52  // InvertedIndexKeys implements the GeographyIndex interface.
    53  func (i *s2GeographyIndex) InvertedIndexKeys(c context.Context, g *geo.Geography) ([]Key, error) {
    54  	r, err := g.AsS2(geo.EmptyBehaviorOmit)
    55  	if err != nil {
    56  		return nil, err
    57  	}
    58  	return invertedIndexKeys(c, i.rc, r), nil
    59  }
    60  
    61  // Covers implements the GeographyIndex interface.
    62  func (i *s2GeographyIndex) Covers(c context.Context, g *geo.Geography) (UnionKeySpans, error) {
    63  	r, err := g.AsS2(geo.EmptyBehaviorOmit)
    64  	if err != nil {
    65  		return nil, err
    66  	}
    67  	return covers(c, i.rc, r), nil
    68  }
    69  
    70  // CoveredBy implements the GeographyIndex interface.
    71  func (i *s2GeographyIndex) CoveredBy(c context.Context, g *geo.Geography) (RPKeyExpr, error) {
    72  	r, err := g.AsS2(geo.EmptyBehaviorOmit)
    73  	if err != nil {
    74  		return nil, err
    75  	}
    76  	return coveredBy(c, i.rc, r), nil
    77  }
    78  
    79  // Intersects implements the GeographyIndex interface.
    80  func (i *s2GeographyIndex) Intersects(c context.Context, g *geo.Geography) (UnionKeySpans, error) {
    81  	r, err := g.AsS2(geo.EmptyBehaviorOmit)
    82  	if err != nil {
    83  		return nil, err
    84  	}
    85  	return intersects(c, i.rc, r), nil
    86  }
    87  
    88  func (i *s2GeographyIndex) TestingInnerCovering(g *geo.Geography) s2.CellUnion {
    89  	r, _ := g.AsS2(geo.EmptyBehaviorOmit)
    90  	if r == nil {
    91  		return nil
    92  	}
    93  	return innerCovering(i.rc, r)
    94  }