github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/split.go (about)

     1  // Copyright 2016 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 sql
    12  
    13  import (
    14  	"context"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/keys"
    17  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    19  	"github.com/cockroachdb/cockroach/pkg/util/hlc"
    20  	"github.com/cockroachdb/errors"
    21  )
    22  
    23  type splitNode struct {
    24  	optColumnsSlot
    25  
    26  	tableDesc      *sqlbase.TableDescriptor
    27  	index          *sqlbase.IndexDescriptor
    28  	rows           planNode
    29  	run            splitRun
    30  	expirationTime hlc.Timestamp
    31  }
    32  
    33  // splitRun contains the run-time state of splitNode during local execution.
    34  type splitRun struct {
    35  	lastSplitKey       []byte
    36  	lastExpirationTime hlc.Timestamp
    37  }
    38  
    39  func (n *splitNode) startExec(params runParams) error {
    40  	return nil
    41  }
    42  
    43  func (n *splitNode) Next(params runParams) (bool, error) {
    44  	// TODO(radu): instead of performing the splits sequentially, accumulate all
    45  	// the split keys and then perform the splits in parallel (e.g. split at the
    46  	// middle key and recursively to the left and right).
    47  
    48  	if ok, err := n.rows.Next(params); err != nil || !ok {
    49  		return ok, err
    50  	}
    51  
    52  	rowKey, err := getRowKey(params.ExecCfg().Codec, n.tableDesc, n.index, n.rows.Values())
    53  	if err != nil {
    54  		return false, err
    55  	}
    56  
    57  	if err := params.ExecCfg().DB.AdminSplit(params.ctx, rowKey, rowKey, n.expirationTime); err != nil {
    58  		return false, err
    59  	}
    60  
    61  	n.run.lastSplitKey = rowKey
    62  	n.run.lastExpirationTime = n.expirationTime
    63  
    64  	return true, nil
    65  }
    66  
    67  func (n *splitNode) Values() tree.Datums {
    68  	splitEnforcedUntil := tree.DNull
    69  	if (n.run.lastExpirationTime != hlc.Timestamp{}) {
    70  		splitEnforcedUntil = tree.TimestampToInexactDTimestamp(n.run.lastExpirationTime)
    71  	}
    72  	return tree.Datums{
    73  		tree.NewDBytes(tree.DBytes(n.run.lastSplitKey)),
    74  		tree.NewDString(keys.PrettyPrint(nil /* valDirs */, n.run.lastSplitKey)),
    75  		splitEnforcedUntil,
    76  	}
    77  }
    78  
    79  func (n *splitNode) Close(ctx context.Context) {
    80  	n.rows.Close(ctx)
    81  }
    82  
    83  // getRowKey generates a key that corresponds to a row (or prefix of a row) in a table or index.
    84  // Both tableDesc and index are required (index can be the primary index).
    85  func getRowKey(
    86  	codec keys.SQLCodec,
    87  	tableDesc *sqlbase.TableDescriptor,
    88  	index *sqlbase.IndexDescriptor,
    89  	values []tree.Datum,
    90  ) ([]byte, error) {
    91  	colMap := make(map[sqlbase.ColumnID]int)
    92  	for i := range values {
    93  		colMap[index.ColumnIDs[i]] = i
    94  	}
    95  	prefix := sqlbase.MakeIndexKeyPrefix(codec, tableDesc, index.ID)
    96  	key, _, err := sqlbase.EncodePartialIndexKey(
    97  		tableDesc, index, len(values), colMap, values, prefix,
    98  	)
    99  	if err != nil {
   100  		return nil, err
   101  	}
   102  	return key, nil
   103  }
   104  
   105  // parseExpriationTime parses an expression into a hlc.Timestamp representing
   106  // the expiration time of the split.
   107  func parseExpirationTime(
   108  	evalCtx *tree.EvalContext, expireExpr tree.TypedExpr,
   109  ) (hlc.Timestamp, error) {
   110  	if !tree.IsConst(evalCtx, expireExpr) {
   111  		return hlc.Timestamp{}, errors.Errorf("SPLIT AT: only constant expressions are allowed for expiration")
   112  	}
   113  	d, err := expireExpr.Eval(evalCtx)
   114  	if err != nil {
   115  		return hlc.Timestamp{}, err
   116  	}
   117  	if d == tree.DNull {
   118  		return hlc.MaxTimestamp, nil
   119  	}
   120  	stmtTimestamp := evalCtx.GetStmtTimestamp()
   121  	ts, err := tree.DatumToHLC(evalCtx, stmtTimestamp, d)
   122  	if err != nil {
   123  		return ts, errors.Wrap(err, "SPLIT AT")
   124  	}
   125  	if ts.GoTime().Before(stmtTimestamp) {
   126  		return ts, errors.Errorf("SPLIT AT: expiration time should be greater than or equal to current time")
   127  	}
   128  	return ts, nil
   129  }