github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/geo/geoindex/utils_test.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  	"fmt"
    15  	"sort"
    16  	"strings"
    17  	"testing"
    18  
    19  	"github.com/cockroachdb/datadriven"
    20  	"github.com/golang/geo/s2"
    21  )
    22  
    23  func nameArg(t *testing.T, d *datadriven.TestData) string {
    24  	var name string
    25  	d.ScanArgs(t, "name", &name)
    26  	return name
    27  }
    28  
    29  func s2Config(t *testing.T, d *datadriven.TestData) S2Config {
    30  	var minLevel, maxLevel, maxCells int
    31  	d.ScanArgs(t, "minlevel", &minLevel)
    32  	d.ScanArgs(t, "maxlevel", &maxLevel)
    33  	d.ScanArgs(t, "maxcells", &maxCells)
    34  	return S2Config{
    35  		MinLevel: int32(minLevel),
    36  		MaxLevel: int32(maxLevel),
    37  		LevelMod: 1,
    38  		MaxCells: int32(maxCells),
    39  	}
    40  }
    41  
    42  func keysToString(keys []Key, err error) string {
    43  	if err != nil {
    44  		return err.Error()
    45  	}
    46  	var cells []string
    47  	for _, k := range keys {
    48  		cells = append(cells, k.String())
    49  	}
    50  	return strings.Join(cells, ", ")
    51  }
    52  
    53  func cellUnionToString(cells s2.CellUnion) string {
    54  	var strs []string
    55  	for _, c := range cells {
    56  		strs = append(strs, Key(c).String())
    57  	}
    58  	return strings.Join(strs, ", ")
    59  }
    60  
    61  func spansToString(spans UnionKeySpans, err error) string {
    62  	if err != nil {
    63  		return err.Error()
    64  	}
    65  	return spans.toString(60)
    66  }
    67  
    68  // Intersection of unioned keys.
    69  type evaluatedExpr [][]Key
    70  
    71  type unionSorter []Key
    72  
    73  func (s unionSorter) Len() int {
    74  	return len(s)
    75  }
    76  
    77  func (s unionSorter) Less(i, j int) bool {
    78  	return s2.CellID(s[i]).Level() > s2.CellID(s[j]).Level()
    79  }
    80  
    81  func (s unionSorter) Swap(i, j int) {
    82  	s[i], s[j] = s[j], s[i]
    83  }
    84  
    85  // Checks that the factored expression does not repeat a cell
    86  // and prints out the unfactored expression with each unioned
    87  // sub-expression in high-to-low level order, so that it is
    88  // easy to read and validate.
    89  func checkExprAndToString(expr RPKeyExpr, err error) string {
    90  	if err != nil {
    91  		return err.Error()
    92  	}
    93  	keys := make(map[Key]struct{})
    94  	for _, elem := range expr {
    95  		switch k := elem.(type) {
    96  		case Key:
    97  			if _, ok := keys[k]; ok {
    98  				return fmt.Sprintf("duplicate key: %s", k)
    99  			}
   100  			keys[k] = struct{}{}
   101  		}
   102  	}
   103  	var stack []evaluatedExpr
   104  	for _, elem := range expr {
   105  		switch e := elem.(type) {
   106  		case Key:
   107  			stack = append(stack, evaluatedExpr{[]Key{e}})
   108  		case RPSetOperator:
   109  			op0, op1 := stack[len(stack)-1], stack[len(stack)-2]
   110  			stack = stack[:len(stack)-2]
   111  			switch e {
   112  			case RPSetIntersection:
   113  				op0 = append(op0, op1...)
   114  			case RPSetUnion:
   115  				if len(op1) != 1 {
   116  					op0, op1 = op1, op0
   117  				}
   118  				if len(op1) != 1 {
   119  					return "error in expression"
   120  				}
   121  				for i := range op0 {
   122  					op0[i] = append(op0[i], op1[0]...)
   123  				}
   124  			}
   125  			stack = append(stack, op0)
   126  		}
   127  	}
   128  	if len(stack) != 1 {
   129  		return fmt.Sprintf("stack has unexpected length: %d", len(stack))
   130  	}
   131  	b := newStringBuilderWithWrap(&strings.Builder{}, 60)
   132  	for i, cells := range stack[0] {
   133  		sort.Sort(unionSorter(cells))
   134  		fmt.Fprintf(b, "%d: ", i)
   135  		for i, c := range cells {
   136  			fmt.Fprintf(b, "%s", c)
   137  			if i != len(cells)-1 {
   138  				b.WriteString(", ")
   139  			}
   140  			b.tryWrap()
   141  		}
   142  		b.doWrap()
   143  	}
   144  
   145  	return b.String()
   146  }