github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/opt/invertedexpr/geo_expression.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 invertedexpr
    12  
    13  import (
    14  	"github.com/cockroachdb/cockroach/pkg/geo/geoindex"
    15  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    16  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    17  	"github.com/cockroachdb/cockroach/pkg/util/encoding"
    18  	"github.com/cockroachdb/errors"
    19  )
    20  
    21  // This file contains functions to encode geoindex.{UnionKeySpans, RPKeyExpr}
    22  // into a SpanExpressionProto. These functions are in this package since they
    23  // need to use sqlbase.EncodeTableKey to convert geoindex.Key to
    24  // invertedexpr.EncInvertedVal and that cannot be done in the geoindex package
    25  // as it introduces a circular dependency.
    26  //
    27  // TODO(sumeer): change geoindex to produce SpanExpressionProtos directly.
    28  
    29  func geoKeyToEncInvertedVal(k geoindex.Key) EncInvertedVal {
    30  	dint := tree.DInt(k)
    31  	encoded, err := sqlbase.EncodeTableKey(nil, &dint, encoding.Ascending)
    32  	if err != nil {
    33  		panic(errors.NewAssertionErrorWithWrappedErrf(err, "unexpected encoding error: %d", k))
    34  	}
    35  	return encoded
    36  }
    37  
    38  func geoToSpan(span geoindex.KeySpan) SpanExpressionProto_Span {
    39  	return SpanExpressionProto_Span{
    40  		Start: geoKeyToEncInvertedVal(span.Start),
    41  		End:   geoKeyToEncInvertedVal(span.End),
    42  	}
    43  }
    44  
    45  // GeoUnionKeySpansToProto converts geoindex.UnionKeySpans to
    46  // SpanExpressionProto.
    47  func GeoUnionKeySpansToProto(ukSpans geoindex.UnionKeySpans) *SpanExpressionProto {
    48  	if len(ukSpans) == 0 {
    49  		return nil
    50  	}
    51  	spans := make([]SpanExpressionProto_Span, len(ukSpans))
    52  	for i, ukSpan := range ukSpans {
    53  		spans[i] = geoToSpan(ukSpan)
    54  	}
    55  	return &SpanExpressionProto{
    56  		SpansToRead: spans,
    57  		Node: SpanExpressionProto_Node{
    58  			FactoredUnionSpans: spans,
    59  		},
    60  	}
    61  }
    62  
    63  // GeoRPKeyExprToProto converts geoindex.RPKeyExpr to SpanExpressionProto.
    64  func GeoRPKeyExprToProto(rpExpr geoindex.RPKeyExpr) (*SpanExpressionProto, error) {
    65  	if len(rpExpr) == 0 {
    66  		return nil, nil
    67  	}
    68  	spansToRead := make([]SpanExpressionProto_Span, 0, len(rpExpr))
    69  	for _, elem := range rpExpr {
    70  		// The keys in the RPKeyExpr are unique.
    71  		if key, ok := elem.(geoindex.Key); ok {
    72  			spansToRead = append(spansToRead, geoToSpan(geoindex.KeySpan{Start: key, End: key + 1}))
    73  		}
    74  	}
    75  	var stack []*SpanExpressionProto_Node
    76  	for _, elem := range rpExpr {
    77  		switch e := elem.(type) {
    78  		case geoindex.Key:
    79  			stack = append(stack, &SpanExpressionProto_Node{
    80  				FactoredUnionSpans: []SpanExpressionProto_Span{
    81  					geoToSpan(geoindex.KeySpan{Start: e, End: e + 1})},
    82  			})
    83  		case geoindex.RPSetOperator:
    84  			if len(stack) < 2 {
    85  				return nil, errors.Errorf("malformed expression: %s", rpExpr)
    86  			}
    87  			node0, node1 := stack[len(stack)-1], stack[len(stack)-2]
    88  			var node *SpanExpressionProto_Node
    89  			stack = stack[:len(stack)-2]
    90  			switch e {
    91  			case geoindex.RPSetIntersection:
    92  				node = &SpanExpressionProto_Node{
    93  					Operator: SetIntersection,
    94  					Left:     node0,
    95  					Right:    node1,
    96  				}
    97  			case geoindex.RPSetUnion:
    98  				if node0.Operator == None {
    99  					node0, node1 = node1, node0
   100  				}
   101  				if node1.Operator == None {
   102  					node = node0
   103  					node.FactoredUnionSpans = append(node.FactoredUnionSpans, node1.FactoredUnionSpans...)
   104  				} else {
   105  					node = &SpanExpressionProto_Node{
   106  						Operator: SetUnion,
   107  						Left:     node0,
   108  						Right:    node1,
   109  					}
   110  				}
   111  			}
   112  			stack = append(stack, node)
   113  		}
   114  	}
   115  	if len(stack) != 1 {
   116  		return nil, errors.Errorf("malformed expression: %s", rpExpr)
   117  	}
   118  	return &SpanExpressionProto{
   119  		SpansToRead: spansToRead,
   120  		Node:        *stack[0],
   121  	}, nil
   122  }