// initScan sets up the rowFetcher and starts a scan. On error, sets n.err and
// returns false.
func (n *scanNode) initScan() (success bool) {
// TODO(radu): we could call init() just once, after the index and
// valNeededForCol are set.
err := n.fetcher.Init(&n.desc, n.colIdxMap, n.index, n.reverse, n.isSecondaryIndex,
n.valNeededForCol)
if err != nil {
n.err = err
return false
}
if len(n.spans) == 0 {
// If no spans were specified retrieve all of the keys that start with our
// index key prefix. This isn't needed for the fetcher, but it is for
// other external users of n.spans.
start := roachpb.Key(sqlbase.MakeIndexKeyPrefix(n.desc.ID, n.index.ID))
n.spans = append(n.spans, sqlbase.Span{Start: start, End: start.PrefixEnd()})
}
n.err = n.fetcher.StartScan(n.txn, n.spans, n.limitHint)
if n.err != nil {
return false
}
n.scanInitialized = true
return true
}
func (td *tableDeleter) deleteIndexFast(ctx context.Context, idx *sqlbase.IndexDescriptor) error {
indexPrefix := sqlbase.MakeIndexKeyPrefix(td.rd.helper.tableDesc, idx.ID)
indexStartKey := roachpb.Key(indexPrefix)
indexEndKey := indexStartKey.PrefixEnd()
if log.V(2) {
log.Infof(ctx, "DelRange %s - %s", indexStartKey, indexEndKey)
}
td.b.DelRange(indexStartKey, indexEndKey, false)
return td.finalize(ctx)
}
func TestDropIndex(t *testing.T) {
defer leaktest.AfterTest(t)()
params, _ := createTestServerParams()
s, sqlDB, kvDB := serverutils.StartServer(t, params)
defer s.Stopper().Stop()
if _, err := sqlDB.Exec(`
CREATE DATABASE t;
CREATE TABLE t.kv (k CHAR PRIMARY KEY, v CHAR);
CREATE INDEX foo on t.kv (v);
INSERT INTO t.kv VALUES ('c', 'e'), ('a', 'c'), ('b', 'd');
`); err != nil {
t.Fatal(err)
}
tableDesc := sqlbase.GetTableDescriptor(kvDB, "t", "kv")
status, i, err := tableDesc.FindIndexByName("foo")
if err != nil {
t.Fatal(err)
}
if status != sqlbase.DescriptorActive {
t.Fatal("Index 'foo' is not active.")
}
indexPrefix := sqlbase.MakeIndexKeyPrefix(tableDesc.ID, tableDesc.Indexes[i].ID)
indexStartKey := roachpb.Key(indexPrefix)
indexEndKey := indexStartKey.PrefixEnd()
if kvs, err := kvDB.Scan(indexStartKey, indexEndKey, 0); err != nil {
t.Fatal(err)
} else if l := 3; len(kvs) != l {
t.Fatalf("expected %d key value pairs, but got %d", l, len(kvs))
}
if _, err := sqlDB.Exec(`DROP INDEX t.kv@foo`); err != nil {
t.Fatal(err)
}
if kvs, err := kvDB.Scan(indexStartKey, indexEndKey, 0); err != nil {
t.Fatal(err)
} else if l := 0; len(kvs) != l {
t.Fatalf("expected %d key value pairs, but got %d", l, len(kvs))
}
tableDesc = sqlbase.GetTableDescriptor(kvDB, "t", "kv")
if _, _, err := tableDesc.FindIndexByName("foo"); err == nil {
t.Fatalf("table descriptor still contains index after index is dropped")
}
if err != nil {
t.Fatal(err)
}
}
func (tu *tableUpserter) init(txn *client.Txn) error {
tu.txn = txn
tu.tableDesc = tu.ri.helper.tableDesc
tu.indexKeyPrefix = sqlbase.MakeIndexKeyPrefix(tu.tableDesc, tu.tableDesc.PrimaryIndex.ID)
allColsIdentityExpr := len(tu.ri.insertCols) == len(tu.tableDesc.Columns) &&
tu.evaler != nil && tu.evaler.isIdentityEvaler()
if len(tu.tableDesc.Indexes) == 0 && allColsIdentityExpr {
tu.fastPathBatch = tu.txn.NewBatch()
tu.fastPathKeys = make(map[string]struct{})
return nil
}
// TODO(dan): This could be made tighter, just the rows needed for the ON
// CONFLICT exprs.
requestedCols := tu.tableDesc.Columns
var err error
if len(tu.updateCols) == 0 {
tu.fetchColIDtoRowIndex = colIDtoRowIndexFromCols(requestedCols)
} else {
tu.ru, err = makeRowUpdater(
txn, tu.tableDesc, tu.fkTables, tu.updateCols, requestedCols, rowUpdaterDefault,
)
if err != nil {
return err
}
tu.fetchColIDtoRowIndex = tu.ru.fetchColIDtoRowIndex
tu.updateColIDtoRowIndex = make(map[sqlbase.ColumnID]int)
for i, updateCol := range tu.ru.updateCols {
tu.updateColIDtoRowIndex[updateCol.ID] = i
}
}
valNeededForCol := make([]bool, len(tu.tableDesc.Columns))
for i := range valNeededForCol {
if _, ok := tu.fetchColIDtoRowIndex[tu.tableDesc.Columns[i].ID]; ok {
valNeededForCol[i] = true
}
}
err = tu.fetcher.Init(
tu.tableDesc, tu.fetchColIDtoRowIndex, &tu.tableDesc.PrimaryIndex, false, false,
tu.tableDesc.Columns, valNeededForCol)
if err != nil {
return err
}
return nil
}
// getTableSpan returns a span containing the start and end key for a table.
func (sc *SchemaChanger) getTableSpan() (sqlbase.Span, error) {
var tableDesc *sqlbase.TableDescriptor
if err := sc.db.Txn(func(txn *client.Txn) error {
var err error
tableDesc, err = getTableDescFromID(txn, sc.tableID)
return err
}); err != nil {
return sqlbase.Span{}, err
}
prefix := roachpb.Key(sqlbase.MakeIndexKeyPrefix(tableDesc.ID, tableDesc.PrimaryIndex.ID))
return sqlbase.Span{
Start: prefix,
End: prefix.PrefixEnd(),
}, nil
}
func makeBaseFKHelper(
txn *client.Txn,
otherTables tableLookupsByID,
writeIdx sqlbase.IndexDescriptor,
ref sqlbase.ForeignKeyReference,
colMap map[sqlbase.ColumnID]int, // col ids (for idx being written) to row offset.
) (baseFKHelper, error) {
b := baseFKHelper{txn: txn, writeIdx: writeIdx, searchTable: otherTables[ref.Table].table}
if b.searchTable == nil {
return b, errors.Errorf("referenced table %d not in provided table map %+v", ref.Table, otherTables)
}
b.searchPrefix = sqlbase.MakeIndexKeyPrefix(b.searchTable, ref.Index)
searchIdx, err := b.searchTable.FindIndexByID(ref.Index)
if err != nil {
return b, err
}
b.prefixLen = len(searchIdx.ColumnIDs)
if len(writeIdx.ColumnIDs) < b.prefixLen {
b.prefixLen = len(writeIdx.ColumnIDs)
}
b.searchIdx = searchIdx
ids := colIDtoRowIndexFromCols(b.searchTable.Columns)
needed := make([]bool, len(ids))
for _, i := range searchIdx.ColumnIDs {
needed[ids[i]] = true
}
isSecondary := b.searchTable.PrimaryIndex.ID != searchIdx.ID
err = b.rf.Init(b.searchTable, ids, searchIdx, false, isSecondary, b.searchTable.Columns, needed)
if err != nil {
return b, err
}
b.ids = make(map[sqlbase.ColumnID]int, len(writeIdx.ColumnIDs))
nulls := true
for i, writeColID := range writeIdx.ColumnIDs[:b.prefixLen] {
if found, ok := colMap[writeColID]; ok {
b.ids[searchIdx.ColumnIDs[i]] = found
nulls = false
}
}
if nulls {
return b, errSkipUnsedFK
}
return b, nil
}
func (sc *SchemaChanger) truncateIndexes(
lease *sqlbase.TableDescriptor_SchemaChangeLease,
dropped []sqlbase.IndexDescriptor,
) error {
for _, desc := range dropped {
// First extend the schema change lease.
l, err := sc.ExtendLease(*lease)
if err != nil {
return err
}
*lease = l
if err := sc.db.Txn(func(txn *client.Txn) error {
tableDesc, err := getTableDescFromID(txn, sc.tableID)
if err != nil {
return err
}
// Short circuit the truncation if the table has been deleted.
if tableDesc.Deleted {
return nil
}
indexPrefix := sqlbase.MakeIndexKeyPrefix(tableDesc.ID, desc.ID)
// Delete the index.
indexStartKey := roachpb.Key(indexPrefix)
indexEndKey := indexStartKey.PrefixEnd()
if log.V(2) {
log.Infof("DelRange %s - %s", indexStartKey, indexEndKey)
}
b := &client.Batch{}
b.DelRange(indexStartKey, indexEndKey, false)
if err := txn.Run(b); err != nil {
return err
}
return nil
}); err != nil {
return err
}
}
return nil
}
// initScan sets up the rowFetcher and starts a scan.
func (n *scanNode) initScan() (err error) {
if len(n.spans) == 0 {
// If no spans were specified retrieve all of the keys that start with our
// index key prefix. This isn't needed for the fetcher, but it is for
// other external users of n.spans.
start := roachpb.Key(sqlbase.MakeIndexKeyPrefix(n.desc.ID, n.index.ID))
n.spans = append(n.spans, sqlbase.Span{Start: start, End: start.PrefixEnd()})
}
limitHint := n.limitHint
if limitHint != 0 && n.limitSoft {
// Read a multiple of the limit if the limit is "soft".
limitHint *= 2
}
if err := n.fetcher.StartScan(n.p.txn, n.spans, limitHint); err != nil {
return err
}
n.scanInitialized = true
return nil
}
// encodeIndexes encodes the primary and secondary index keys. The
// secondaryIndexEntries are only valid until the next call to encodeIndexes or
// encodeSecondaryIndexes.
func (rh *rowHelper) encodeIndexes(
colIDtoRowIndex map[sqlbase.ColumnID]int, values []parser.Datum,
) (
primaryIndexKey []byte,
secondaryIndexEntries []sqlbase.IndexEntry,
err error,
) {
if rh.primaryIndexKeyPrefix == nil {
rh.primaryIndexKeyPrefix = sqlbase.MakeIndexKeyPrefix(rh.tableDesc,
rh.tableDesc.PrimaryIndex.ID)
}
primaryIndexKey, _, err = sqlbase.EncodeIndexKey(
rh.tableDesc, &rh.tableDesc.PrimaryIndex, colIDtoRowIndex, values, rh.primaryIndexKeyPrefix)
if err != nil {
return nil, nil, err
}
secondaryIndexEntries, err = rh.encodeSecondaryIndexes(colIDtoRowIndex, values)
if err != nil {
return nil, nil, err
}
return primaryIndexKey, secondaryIndexEntries, nil
}
func (td *tableDeleter) deleteAllRowsScan(ctx context.Context) error {
tablePrefix := sqlbase.MakeIndexKeyPrefix(
td.rd.helper.tableDesc, td.rd.helper.tableDesc.PrimaryIndex.ID)
span := sqlbase.Span{Start: roachpb.Key(tablePrefix), End: roachpb.Key(tablePrefix).PrefixEnd()}
valNeededForCol := make([]bool, len(td.rd.helper.tableDesc.Columns))
for _, idx := range td.rd.fetchColIDtoRowIndex {
valNeededForCol[idx] = true
}
var rf sqlbase.RowFetcher
err := rf.Init(
td.rd.helper.tableDesc, td.rd.fetchColIDtoRowIndex, &td.rd.helper.tableDesc.PrimaryIndex,
false, false, td.rd.fetchCols, valNeededForCol)
if err != nil {
return err
}
if err := rf.StartScan(td.txn, sqlbase.Spans{span}, 0); err != nil {
return err
}
for {
row, err := rf.NextRow()
if err != nil {
return err
}
if row == nil {
// Done deleting rows.
break
}
_, err = td.row(ctx, row)
if err != nil {
return err
}
}
return td.finalize(ctx)
}
func restoreTable(
ctx context.Context,
sst engine.RocksDBSstFileReader,
txn *client.Txn,
table *sqlbase.TableDescriptor,
overwrite bool,
) error {
log.Infof(ctx, "Restoring Table %q", table.Name)
tableStartKey := roachpb.Key(sqlbase.MakeIndexKeyPrefix(table, table.PrimaryIndex.ID))
tableEndKey := tableStartKey.PrefixEnd()
existingDesc, err := txn.Get(sqlbase.MakeDescMetadataKey(table.GetID()))
if err != nil {
return err
}
existingData, err := txn.Scan(tableStartKey, tableEndKey, 1)
if err != nil {
return err
}
if existingDesc.Value != nil || len(existingData) > 0 {
if overwrite {
// We're about to Put the descriptor, so don't bother deleting it.
if err := txn.DelRange(tableStartKey, tableEndKey); err != nil {
return err
}
} else {
return errors.Errorf("table %q already exists", table.Name)
}
}
tableDescKey := sqlbase.MakeDescMetadataKey(table.GetID())
if err := txn.Put(tableDescKey, sqlbase.WrapDescriptor(table)); err != nil {
return err
}
return Import(ctx, sst, txn, engine.MVCCKey{Key: tableStartKey}, engine.MVCCKey{Key: tableEndKey})
}
func (tu *tableUpserter) init(txn *client.Txn) error {
tu.txn = txn
tu.tableDesc = tu.ri.helper.tableDesc
tu.indexKeyPrefix = sqlbase.MakeIndexKeyPrefix(tu.tableDesc.ID, tu.tableDesc.PrimaryIndex.ID)
tu.updateColIDtoRowIndex = make(map[sqlbase.ColumnID]int)
for i, updateCol := range tu.ru.updateCols {
tu.updateColIDtoRowIndex[updateCol.ID] = i
}
tu.updateRow = make(parser.DTuple, len(tu.updateColIDtoRowIndex))
valNeededForCol := make([]bool, len(tu.ru.fetchCols))
for i := range valNeededForCol {
valNeededForCol[i] = true
}
err := tu.fetcher.Init(
tu.tableDesc, tu.ru.fetchColIDtoRowIndex, &tu.tableDesc.PrimaryIndex,
false, false, valNeededForCol)
if err != nil {
return err
}
return nil
}
func TestTableReader(t *testing.T) {
defer leaktest.AfterTest(t)()
_, sqlDB, kvDB, cleanup := sqlutils.SetupServer(t)
defer cleanup()
// Create a table where each row is:
//
// | a | b | sum | s |
// |-----------------------------------------------------------------|
// | rowId/10 | rowId%10 | rowId/10 + rowId%10 | IntToEnglish(rowId) |
aFn := func(row int) parser.Datum {
return parser.NewDInt(parser.DInt(row / 10))
}
bFn := func(row int) parser.Datum {
return parser.NewDInt(parser.DInt(row % 10))
}
sumFn := func(row int) parser.Datum {
return parser.NewDInt(parser.DInt(row/10 + row%10))
}
sqlutils.CreateTable(t, sqlDB, "t",
"a INT, b INT, sum INT, s STRING, PRIMARY KEY (a,b), INDEX bs (b,s)",
99,
sqlutils.ToRowFn(aFn, bFn, sumFn, sqlutils.RowEnglishFn))
td := sqlbase.GetTableDescriptor(kvDB, "test", "t")
makeIndexSpan := func(start, end int) TableReaderSpan {
var span roachpb.Span
prefix := roachpb.Key(sqlbase.MakeIndexKeyPrefix(td.ID, td.Indexes[0].ID))
span.Key = append(prefix, encoding.EncodeVarintAscending(nil, int64(start))...)
span.EndKey = append(span.EndKey, prefix...)
span.EndKey = append(span.EndKey, encoding.EncodeVarintAscending(nil, int64(end))...)
return TableReaderSpan{Span: span}
}
testCases := []struct {
spec TableReaderSpec
expected string
}{
{
spec: TableReaderSpec{
Filter: Expression{Expr: "$2 < 5 AND $1 != 3"}, // sum < 5 && b != 3
OutputColumns: []uint32{0, 1},
},
expected: "[[0 1] [0 2] [0 4] [1 0] [1 1] [1 2] [2 0] [2 1] [2 2] [3 0] [3 1] [4 0]]",
},
{
spec: TableReaderSpec{
Filter: Expression{Expr: "$2 < 5 AND $1 != 3"},
OutputColumns: []uint32{3}, // s
HardLimit: 4,
},
expected: "[['one'] ['two'] ['four'] ['one-zero']]",
},
{
spec: TableReaderSpec{
IndexIdx: 1,
Reverse: true,
Spans: []TableReaderSpan{makeIndexSpan(4, 6)},
Filter: Expression{Expr: "$0 < 3"}, // sum < 8
OutputColumns: []uint32{0, 1},
SoftLimit: 1,
},
expected: "[[2 5] [1 5] [0 5] [2 4] [1 4] [0 4]]",
},
}
for _, c := range testCases {
ts := c.spec
ts.Table = *td
txn := client.NewTxn(context.Background(), *kvDB)
out := &RowBuffer{}
tr, err := newTableReader(&ts, txn, out, &parser.EvalContext{})
if err != nil {
t.Fatal(err)
}
tr.Run(nil)
if out.err != nil {
t.Fatal(out.err)
}
if !out.closed {
t.Fatalf("output RowReceiver not closed")
}
if result := out.rows.String(); result != c.expected {
t.Errorf("invalid results: %s, expected %s'", result, c.expected)
}
}
}
// mainLoop runs the mainLoop and returns any error.
// It does not close the output.
func (jr *joinReader) mainLoop() error {
primaryKeyPrefix := sqlbase.MakeIndexKeyPrefix(&jr.desc, jr.index.ID)
var alloc sqlbase.DatumAlloc
spans := make(sqlbase.Spans, 0, joinReaderBatchSize)
for {
// TODO(radu): figure out how to send smaller batches if the source has
// a soft limit (perhaps send the batch out if we don't get a result
// within a certain amount of time).
for spans = spans[:0]; len(spans) < joinReaderBatchSize; {
row, err := jr.input.NextRow()
if err != nil {
return err
}
if row == nil {
if len(spans) == 0 {
return nil
}
break
}
key, err := jr.generateKey(row, &alloc, primaryKeyPrefix)
if err != nil {
return err
}
spans = append(spans, sqlbase.Span{
Start: key,
End: key.PrefixEnd(),
})
}
err := jr.fetcher.StartScan(jr.txn, spans, 0)
if err != nil {
return err
}
// TODO(radu): we are consuming all results from a fetch before starting
// the next batch. We could start the next batch early while we are
// outputting rows.
for {
outRow, err := jr.nextRow()
if err != nil {
return err
}
if outRow == nil {
// Done.
break
}
// Push the row to the output RowReceiver; stop if they don't need more
// rows.
if !jr.output.PushRow(outRow) {
return nil
}
}
if len(spans) != joinReaderBatchSize {
// This was the last batch.
return nil
}
}
}
func TestMakeSpans(t *testing.T) {
defer leaktest.AfterTest(t)()
testData := []struct {
expr string
columns string
expectedAsc string
expectedDesc string
}{
{`a = 1`, `a`, `/1-/2`, `/1-/0`},
{`a != 1`, `a`, `/#-`, `-/#`},
{`a > 1`, `a`, `/2-`, `-/1`},
{`a >= 1`, `a`, `/1-`, `-/0`},
{`a < 1`, `a`, `/#-/1`, `/0-/#`},
{`a <= 1`, `a`, `/#-/2`, `/1-/#`},
{`a IS NULL`, `a`, `-/#`, `/NULL-`},
{`a IS NOT NULL`, `a`, `/#-`, `-/#`},
{`a IN (1,2,3)`, `a`, `/1-/4`, `/3-/0`},
{`a IN (1,3,5)`, `a`, `/1-/2 /3-/4 /5-/6`, `/5-/4 /3-/2 /1-/0`},
{`a IN (1,2,3) AND b = 1`, `a,b`,
`/1/1-/1/2 /2/1-/2/2 /3/1-/3/2`, `/3/1-/3/0 /2/1-/2/0 /1/1-/1/0`},
{`a = 1 AND b IN (1,2,3)`, `a,b`,
`/1/1-/1/4`, `/1/3-/1/0`},
{`a = 1 AND b IN (1,3,5)`, `a,b`,
`/1/1-/1/2 /1/3-/1/4 /1/5-/1/6`, `/1/5-/1/4 /1/3-/1/2 /1/1-/1/0`},
{`a >= 1 AND b IN (1,2,3)`, `a,b`, `/1-`, `-/0`},
{`a <= 1 AND b IN (1,2,3)`, `a,b`, `/#-/2`, `/1-/#`},
{`(a, b) IN ((1, 2), (3, 4))`, `a,b`,
`/1/2-/1/3 /3/4-/3/5`, `/3/4-/3/3 /1/2-/1/1`},
{`(b, a) IN ((1, 2), (3, 4))`, `a,b`,
`/2/1-/2/2 /4/3-/4/4`, `/4/3-/4/2 /2/1-/2/0`},
{`(a, b) IN ((1, 2), (3, 4))`, `b`, `/2-/3 /4-/5`, `/4-/3 /2-/1`},
{`a = 1 AND b = 1`, `a,b`, `/1/1-/1/2`, `/1/1-/1/0`},
{`a = 1 AND b != 1`, `a,b`, `/1/#-/2`, `/1-/1/#`},
{`a = 1 AND b > 1`, `a,b`, `/1/2-/2`, `/1-/1/1`},
{`a = 1 AND b >= 1`, `a,b`, `/1/1-/2`, `/1-/1/0`},
{`a = 1 AND b < 1`, `a,b`, `/1/#-/1/1`, `/1/0-/1/#`},
{`a = 1 AND b <= 1`, `a,b`, `/1/#-/1/2`, `/1/1-/1/#`},
{`a = 1 AND b IS NULL`, `a,b`, `/1-/1/#`, `/1/NULL-/0`},
{`a = 1 AND b IS NOT NULL`, `a,b`, `/1/#-/2`, `/1-/1/#`},
{`a != 1 AND b = 1`, `a,b`, `/#-`, `-/#`},
{`a != 1 AND b != 1`, `a,b`, `/#-`, `-/#`},
{`a != 1 AND b > 1`, `a,b`, `/#-`, `-/#`},
{`a != 1 AND b >= 1`, `a,b`, `/#-`, `-/#`},
{`a != 1 AND b < 1`, `a,b`, `/#-`, `-/#`},
{`a != 1 AND b <= 1`, `a,b`, `/#-`, `-/#`},
{`a != 1 AND b IS NULL`, `a,b`, `/#-`, `-/#`},
{`a != 1 AND b IS NOT NULL`, `a,b`, `/#-`, `-/#`},
{`a > 1 AND b = 1`, `a,b`, `/2/1-`, `-/2/0`},
{`a > 1 AND b != 1`, `a,b`, `/2/#-`, `-/2/#`},
{`a > 1 AND b > 1`, `a,b`, `/2/2-`, `-/2/1`},
{`a > 1 AND b >= 1`, `a,b`, `/2/1-`, `-/2/0`},
{`a > 1 AND b < 1`, `a,b`, `/2-`, `-/1`},
{`a > 1 AND b <= 1`, `a,b`, `/2-`, `-/1`},
{`a > 1 AND b IS NULL`, `a,b`, `/2-`, `-/1`},
{`a > 1 AND b IS NOT NULL`, `a,b`, `/2/#-`, `-/2/#`},
{`a >= 1 AND b = 1`, `a,b`, `/1/1-`, `-/1/0`},
{`a >= 1 AND b != 1`, `a,b`, `/1/#-`, `-/1/#`},
{`a >= 1 AND b > 1`, `a,b`, `/1/2-`, `-/1/1`},
{`a >= 1 AND b >= 1`, `a,b`, `/1/1-`, `-/1/0`},
{`a >= 1 AND b < 1`, `a,b`, `/1-`, `-/0`},
{`a >= 1 AND b <= 1`, `a,b`, `/1-`, `-/0`},
{`a >= 1 AND b IS NULL`, `a,b`, `/1-`, `-/0`},
{`a >= 1 AND b IS NOT NULL`, `a,b`, `/1/#-`, `-/1/#`},
{`a < 1 AND b = 1`, `a,b`, `/#-/0/2`, `/0/1-/#`},
{`a < 1 AND b != 1`, `a,b`, `/#-/1`, `/0-/#`},
{`a < 1 AND b > 1`, `a,b`, `/#-/1`, `/0-/#`},
{`a < 1 AND b >= 1`, `a,b`, `/#-/1`, `/0-/#`},
{`a < 1 AND b < 1`, `a,b`, `/#-/0/1`, `/0/0-/#`},
{`a < 1 AND b <= 1`, `a,b`, `/#-/0/2`, `/0/1-/#`},
{`a < 1 AND b IS NULL`, `a,b`, `/#-/0/#`, `/0/NULL-/#`},
{`a < 1 AND b IS NOT NULL`, `a,b`, `/#-/1`, `/0-/#`},
{`a <= 1 AND b = 1`, `a,b`, `/#-/1/2`, `/1/1-/#`},
{`a <= 1 AND b != 1`, `a,b`, `/#-/2`, `/1-/#`},
{`a <= 1 AND b > 1`, `a,b`, `/#-/2`, `/1-/#`},
{`a <= 1 AND b >= 1`, `a,b`, `/#-/2`, `/1-/#`},
{`a <= 1 AND b < 1`, `a,b`, `/#-/1/1`, `/1/0-/#`},
{`a <= 1 AND b <= 1`, `a,b`, `/#-/1/2`, `/1/1-/#`},
{`a <= 1 AND b IS NULL`, `a,b`, `/#-/1/#`, `/1/NULL-/#`},
{`a <= 1 AND b IS NOT NULL`, `a,b`, `/#-/2`, `/1-/#`},
{`a IN (1) AND b = 1`, `a,b`, `/1/1-/1/2`, `/1/1-/1/0`},
{`a IN (1) AND b != 1`, `a,b`, `/1/#-/2`, `/1-/1/#`},
{`a IN (1) AND b > 1`, `a,b`, `/1/2-/2`, `/1-/1/1`},
{`a IN (1) AND b >= 1`, `a,b`, `/1/1-/2`, `/1-/1/0`},
{`a IN (1) AND b < 1`, `a,b`, `/1/#-/1/1`, `/1/0-/1/#`},
{`a IN (1) AND b <= 1`, `a,b`, `/1/#-/1/2`, `/1/1-/1/#`},
{`a IN (1) AND b IS NULL`, `a,b`, `/1-/1/#`, `/1/NULL-/0`},
{`a IN (1) AND b IS NOT NULL`, `a,b`, `/1/#-/2`, `/1-/1/#`},
{`(a, b) = (1, 2)`, `a`, `/1-/2`, `/1-/0`},
{`(a, b) = (1, 2)`, `a,b`, `/1/2-/1/3`, `/1/2-/1/1`},
{`a > 1 OR a >= 5`, `a`, `/2-`, `-/1`},
{`a < 5 OR a >= 1`, `a`, `/#-`, `-/#`},
{`a < 1 OR a >= 5`, `a`, `/#-/1 /5-`, `-/4 /0-/#`},
{`a = 1 OR a > 8`, `a`, `/1-/2 /9-`, `-/8 /1-/0`},
{`a = 8 OR a > 1`, `a`, `/2-`, `-/1`},
{`a < 1 OR a = 5 OR a > 8`, `a`, `/#-/1 /5-/6 /9-`, `-/8 /5-/4 /0-/#`},
{`a < 8 OR a = 8 OR a > 8`, `a`, `/#-`, `-/#`},
{`(a = 1 AND b = 5) OR (a = 3 AND b = 7)`, `a`, `/1-/2 /3-/4`, `/3-/2 /1-/0`},
{`(a = 1 AND b = 5) OR (a = 3 AND b = 7)`, `b`, `/5-/6 /7-/8`, `/7-/6 /5-/4`},
{`(a = 1 AND b = 5) OR (a = 3 AND b = 7)`, `a,b`,
`/1/5-/1/6 /3/7-/3/8`, `/3/7-/3/6 /1/5-/1/4`},
{`(a = 1 AND b < 5) OR (a = 3 AND b > 7)`, `a`, `/1-/2 /3-/4`, `/3-/2 /1-/0`},
{`(a = 1 AND b < 5) OR (a = 3 AND b > 7)`, `b`, `/#-/5 /8-`, `-/7 /4-/#`},
{`(a = 1 AND b < 5) OR (a = 3 AND b > 7)`, `a,b`,
`/1/#-/1/5 /3/8-/4`, `/3-/3/7 /1/4-/1/#`},
{`(a = 1 AND b > 5) OR (a = 3 AND b > 7)`, `a`, `/1-/2 /3-/4`, `/3-/2 /1-/0`},
{`(a = 1 AND b > 5) OR (a = 3 AND b > 7)`, `b`, `/6-`, `-/5`},
{`(a = 1 AND b > 5) OR (a = 3 AND b > 7)`, `a,b`,
`/1/6-/2 /3/8-/4`, `/3-/3/7 /1-/1/5`},
{`(a = 1 AND b > 5) OR (a = 3 AND b < 7)`, `a`, `/1-/2 /3-/4`, `/3-/2 /1-/0`},
{`(a = 1 AND b > 5) OR (a = 3 AND b < 7)`, `b`, `/#-`, `-/#`},
{`(a = 1 AND b > 5) OR (a = 3 AND b < 7)`, `a,b`,
`/1/6-/2 /3/#-/3/7`, `/3/6-/3/# /1-/1/5`},
{`(a < 1 AND b < 5) OR (a > 3 AND b > 7)`, `a`, `/#-/1 /4-`, `-/3 /0-/#`},
{`(a < 1 AND b < 5) OR (a > 3 AND b > 7)`, `b`, `/#-/5 /8-`, `-/7 /4-/#`},
{`(a < 1 AND b < 5) OR (a > 3 AND b > 7)`, `a,b`,
`/#-/0/5 /4/8-`, `-/4/7 /0/4-/#`},
{`(a > 3 AND b < 5) OR (a < 1 AND b > 7)`, `a`, `/#-/1 /4-`, `-/3 /0-/#`},
{`(a > 3 AND b < 5) OR (a < 1 AND b > 7)`, `b`, `/#-/5 /8-`, `-/7 /4-/#`},
{`(a > 3 AND b < 5) OR (a < 1 AND b > 7)`, `a,b`,
`/#-/1 /4-`, `-/3 /0-/#`},
{`(a > 1 AND b < 5) OR (a < 3 AND b > 7)`, `a`, `/#-`, `-/#`},
{`(a > 1 AND b < 5) OR (a < 3 AND b > 7)`, `b`, `/#-/5 /8-`, `-/7 /4-/#`},
{`(a > 1 AND b < 5) OR (a < 3 AND b > 7)`, `a,b`, `/#-`, `-/#`},
{`(a = 5) OR (a, b) IN ((1, 1), (3, 3))`, `a`, `/1-/2 /3-/4 /5-/6`, `/5-/4 /3-/2 /1-/0`},
{`(a = 5) OR (a, b) IN ((1, 1), (3, 3))`, `b`, `-`, `-`},
{`(a = 5) OR (a, b) IN ((1, 1), (3, 3))`, `a,b`,
`/1/1-/1/2 /3/3-/3/4 /5-/6`, `/5-/4 /3/3-/3/2 /1/1-/1/0`},
// When encoding an end constraint for a maximal datum, we use
// bytes.PrefixEnd() to go beyond the normal encodings of that datatype.
{fmt.Sprintf(`a = %d`, math.MaxInt64), `a`,
`/9223372036854775807-/<varint 9223372036854775808 overflows int64>`,
`/9223372036854775807-/9223372036854775806`},
{fmt.Sprintf(`a = %d`, math.MinInt64), `a`,
`/-9223372036854775808-/-9223372036854775807`,
`/-9223372036854775808-/<varint 9223372036854775808 overflows int64>`},
{`(a, b) >= (1, 4)`, `a,b`, `/1/4-`, `-/1/3`},
{`(a, b) > (1, 4)`, `a,b`, `/1/5-`, `-/1/4`},
{`(a, b) < (1, 4)`, `a,b`, `/#-/1/4`, `/1/3-/#`},
{`(a, b) <= (1, 4)`, `a,b`, `/#-/1/5`, `/1/4-/#`},
{`(a, b) = (1, 4)`, `a,b`, `/1/4-/1/5`, `/1/4-/1/3`},
{`(a, b) != (1, 4)`, `a,b`, `/#-`, `-/#`},
}
for _, d := range testData {
for _, dir := range []encoding.Direction{encoding.Ascending, encoding.Descending} {
columns := strings.Split(d.columns, ",")
dirs := make([]encoding.Direction, 0, len(columns))
for range columns {
dirs = append(dirs, dir)
}
desc, index := makeTestIndex(t, columns, dirs)
constraints, _ := makeConstraints(t, d.expr, desc, index)
spans := makeSpans(constraints, desc.ID, index)
s := sqlbase.PrettySpans(spans, 2)
var expected string
if dir == encoding.Ascending {
expected = d.expectedAsc
} else {
expected = d.expectedDesc
}
s = keys.MassagePrettyPrintedSpanForTest(s, indexToDirs(index))
if expected != s {
t.Errorf("[index direction: %d] %s: expected %s, but found %s", dir, d.expr, expected, s)
}
}
}
// Test indexes with mixed-directions (some cols Asc, some cols Desc) and other edge cases.
testData2 := []struct {
expr string
columns string
expected string
}{
{`a = 1 AND b = 5`, `a,b-,c`, `/1/5-/1/4`},
{`a = 7 AND b IN (1,2,3) AND c = false`, `a,b-,c`,
`/7/3/0-/7/3/1 /7/2/0-/7/2/1 /7/1/0-/7/1/1`},
// Test different directions for te columns inside a tuple.
{`(a,b,j) IN ((1,2,3), (4,5,6))`, `a-,b,j-`, `/4/5/6-/4/5/5 /1/2/3-/1/2/2`},
{`i = E'\xff'`, `i`, `/"\xff"-/"\xff\x00"`},
// Test that limits on bytes work correctly: when encoding a descending limit for bytes,
// we need to go outside the bytes encoding.
// "\xaa" is encoded as [bytesDescMarker, ^0xaa, <term escape sequence>]
{`i = E'\xaa'`, `i-`,
fmt.Sprintf("raw:%c%c\xff\xfe-%c%c\xff\xff",
encoding.BytesDescMarker, ^byte(0xaa), encoding.BytesDescMarker, ^byte(0xaa))},
// Ensure tuples with differing index directions aren't constrained.
// TODO(mjibson): fix this, see #6346
{`(a, b) >= (1, 4)`, `a-,b`, `-`},
{`(a, b) >= (1, 4)`, `a,b-`, `-`},
}
for _, d := range testData2 {
desc, index := makeTestIndexFromStr(t, d.columns)
constraints, _ := makeConstraints(t, d.expr, desc, index)
spans := makeSpans(constraints, desc.ID, index)
var got string
raw := false
if strings.HasPrefix(d.expected, "raw:") {
raw = true
span := spans[0]
d.expected = d.expected[4:]
// Trim the index prefix from the span.
prefix := string(sqlbase.MakeIndexKeyPrefix(desc.ID, index.ID))
got = strings.TrimPrefix(string(span.Start), prefix) + "-" +
strings.TrimPrefix(string(span.End), prefix)
} else {
got = keys.MassagePrettyPrintedSpanForTest(sqlbase.PrettySpans(spans, 2),
indexToDirs(index))
}
if d.expected != got {
if !raw {
t.Errorf("%s: expected %s, but found %s", d.expr, d.expected, got)
} else {
t.Errorf("%s: expected %# x, but found %# x", d.expr, []byte(d.expected), got)
}
}
}
}
func TestClusterFlow(t *testing.T) {
defer leaktest.AfterTest(t)()
const numRows = 100
args := base.TestClusterArgs{ReplicationMode: base.ReplicationManual}
tc := serverutils.StartTestCluster(t, 3, args)
defer tc.Stopper().Stop()
sumDigitsFn := func(row int) parser.Datum {
sum := 0
for row > 0 {
sum += row % 10
row /= 10
}
return parser.NewDInt(parser.DInt(sum))
}
sqlutils.CreateTable(t, tc.ServerConn(0), "t",
"num INT PRIMARY KEY, digitsum INT, numstr STRING, INDEX s (digitsum)",
numRows,
sqlutils.ToRowFn(sqlutils.RowIdxFn, sumDigitsFn, sqlutils.RowEnglishFn))
kvDB := tc.Server(0).KVClient().(*client.DB)
desc := sqlbase.GetTableDescriptor(kvDB, "test", "t")
makeIndexSpan := func(start, end int) TableReaderSpan {
var span roachpb.Span
prefix := roachpb.Key(sqlbase.MakeIndexKeyPrefix(desc, desc.Indexes[0].ID))
span.Key = append(prefix, encoding.EncodeVarintAscending(nil, int64(start))...)
span.EndKey = append(span.EndKey, prefix...)
span.EndKey = append(span.EndKey, encoding.EncodeVarintAscending(nil, int64(end))...)
return TableReaderSpan{Span: span}
}
// Set up table readers on three hosts feeding data into a join reader on
// the third host. This is a basic test for the distributed flow
// infrastructure, including local and remote streams.
//
// Note that the ranges won't necessarily be local to the table readers, but
// that doesn't matter for the purposes of this test.
tr1 := TableReaderSpec{
Table: *desc,
IndexIdx: 1,
OutputColumns: []uint32{0, 1},
Spans: []TableReaderSpan{makeIndexSpan(0, 8)},
}
tr2 := TableReaderSpec{
Table: *desc,
IndexIdx: 1,
OutputColumns: []uint32{0, 1},
Spans: []TableReaderSpan{makeIndexSpan(8, 12)},
}
tr3 := TableReaderSpec{
Table: *desc,
IndexIdx: 1,
OutputColumns: []uint32{0, 1},
Spans: []TableReaderSpan{makeIndexSpan(12, 100)},
}
jr := JoinReaderSpec{
Table: *desc,
OutputColumns: []uint32{2},
}
txn := client.NewTxn(context.Background(), *kvDB)
fid := FlowID{uuid.MakeV4()}
req1 := &SetupFlowRequest{Txn: txn.Proto}
req1.Flow = FlowSpec{
FlowID: fid,
Processors: []ProcessorSpec{{
Core: ProcessorCoreUnion{TableReader: &tr1},
Output: []OutputRouterSpec{{
Type: OutputRouterSpec_MIRROR,
Streams: []StreamEndpointSpec{
{Mailbox: &MailboxSpec{StreamID: 0, TargetAddr: tc.Server(2).ServingAddr()}},
},
}},
}},
}
req2 := &SetupFlowRequest{Txn: txn.Proto}
req2.Flow = FlowSpec{
FlowID: fid,
Processors: []ProcessorSpec{{
Core: ProcessorCoreUnion{TableReader: &tr2},
Output: []OutputRouterSpec{{
Type: OutputRouterSpec_MIRROR,
Streams: []StreamEndpointSpec{
{Mailbox: &MailboxSpec{StreamID: 1, TargetAddr: tc.Server(2).ServingAddr()}},
},
}},
}},
}
req3 := &SetupFlowRequest{Txn: txn.Proto}
req3.Flow = FlowSpec{
FlowID: fid,
Processors: []ProcessorSpec{
{
Core: ProcessorCoreUnion{TableReader: &tr3},
Output: []OutputRouterSpec{{
Type: OutputRouterSpec_MIRROR,
Streams: []StreamEndpointSpec{
{LocalStreamID: LocalStreamID(0)},
},
}},
},
{
Input: []InputSyncSpec{{
Type: InputSyncSpec_ORDERED,
Ordering: Ordering{Columns: []Ordering_Column{{1, Ordering_Column_ASC}}},
Streams: []StreamEndpointSpec{
{Mailbox: &MailboxSpec{StreamID: 0}},
{Mailbox: &MailboxSpec{StreamID: 1}},
{LocalStreamID: LocalStreamID(0)},
},
}},
Core: ProcessorCoreUnion{JoinReader: &jr},
Output: []OutputRouterSpec{{
Type: OutputRouterSpec_MIRROR,
Streams: []StreamEndpointSpec{{Mailbox: &MailboxSpec{SimpleResponse: true}}},
}}},
},
}
var clients []DistSQLClient
for i := 0; i < 3; i++ {
s := tc.Server(i)
conn, err := s.RPCContext().GRPCDial(s.ServingAddr())
if err != nil {
t.Fatal(err)
}
clients = append(clients, NewDistSQLClient(conn))
}
ctx := context.Background()
if log.V(1) {
log.Infof(ctx, "Setting up flow on 0")
}
if resp, err := clients[0].SetupFlow(context.Background(), req1); err != nil {
t.Fatal(err)
} else if resp.Error != nil {
t.Fatal(resp.Error)
}
if log.V(1) {
log.Infof(ctx, "Setting up flow on 1")
}
if resp, err := clients[1].SetupFlow(context.Background(), req2); err != nil {
t.Fatal(err)
} else if resp.Error != nil {
t.Fatal(resp.Error)
}
if log.V(1) {
log.Infof(ctx, "Running flow on 2")
}
stream, err := clients[2].RunSimpleFlow(context.Background(), req3)
if err != nil {
t.Fatal(err)
}
var decoder StreamDecoder
var rows sqlbase.EncDatumRows
for {
msg, err := stream.Recv()
if err != nil {
if err == io.EOF {
break
}
t.Fatal(err)
}
err = decoder.AddMessage(msg)
if err != nil {
t.Fatal(err)
}
rows = testGetDecodedRows(t, &decoder, rows)
}
if done, trailerErr := decoder.IsDone(); !done {
t.Fatal("stream not done")
} else if trailerErr != nil {
t.Fatal("error in the stream trailer:", trailerErr)
}
// The result should be all the numbers in string form, ordered by the
// digit sum (and then by number).
var results []string
for sum := 1; sum <= 50; sum++ {
for i := 1; i <= numRows; i++ {
if int(*sumDigitsFn(i).(*parser.DInt)) == sum {
results = append(results, fmt.Sprintf("['%s']", sqlutils.IntToEnglish(i)))
}
}
}
expected := strings.Join(results, " ")
expected = "[" + expected + "]"
if rowStr := rows.String(); rowStr != expected {
t.Errorf("Result: %s\n Expected: %s\n", rowStr, expected)
}
}
func makeIndexJoin(indexScan *scanNode, exactPrefix int) *indexJoinNode {
// Create a new table scan node with the primary index.
table := &scanNode{planner: indexScan.planner, txn: indexScan.txn}
table.desc = indexScan.desc
table.initDescDefaults()
table.initOrdering(0)
colIDtoRowIndex := map[sqlbase.ColumnID]int{}
for _, colID := range table.desc.PrimaryIndex.ColumnIDs {
idx, ok := indexScan.colIdxMap[colID]
if !ok {
panic(fmt.Sprintf("Unknown column %d in PrimaryIndex!", colID))
}
colIDtoRowIndex[colID] = idx
}
for _, colID := range indexScan.index.ColumnIDs {
idx, ok := indexScan.colIdxMap[colID]
if !ok {
panic(fmt.Sprintf("Unknown column %d in index!", colID))
}
colIDtoRowIndex[colID] = idx
}
for i := range indexScan.valNeededForCol {
// We transfer valNeededForCol to the table node.
table.valNeededForCol[i] = indexScan.valNeededForCol[i]
// For the index node, we set valNeededForCol for columns that are part of the index.
id := indexScan.desc.Columns[i].ID
_, found := colIDtoRowIndex[id]
indexScan.valNeededForCol[i] = found
}
if indexScan.filter != nil {
// Transfer the filter to the table node. We must first convert the
// IndexedVars associated with indexNode.
convFunc := func(expr parser.VariableExpr) (ok bool, newExpr parser.VariableExpr) {
iv := expr.(*parser.IndexedVar)
return true, table.filterVars.IndexedVar(iv.Idx)
}
table.filter = exprConvertVars(indexScan.filter, convFunc)
// Now we split the filter by extracting the part that can be evaluated using just the index
// columns.
splitFunc := func(expr parser.VariableExpr) (ok bool, newExpr parser.VariableExpr) {
colIdx := expr.(*parser.IndexedVar).Idx
if !indexScan.valNeededForCol[colIdx] {
return false, nil
}
return true, indexScan.filterVars.IndexedVar(colIdx)
}
indexScan.filter, table.filter = splitFilter(table.filter, splitFunc)
}
indexScan.initOrdering(exactPrefix)
primaryKeyPrefix := roachpb.Key(sqlbase.MakeIndexKeyPrefix(table.desc.ID, table.index.ID))
return &indexJoinNode{
index: indexScan,
table: table,
primaryKeyPrefix: primaryKeyPrefix,
colIDtoRowIndex: colIDtoRowIndex,
}
}
// makeSpansForIndexConstraints constructs the spans for an index given an
// instance of indexConstraints. The resulting spans are non-overlapping (by
// virtue of the input constraints being disjunct).
func makeSpansForIndexConstraints(
constraints indexConstraints, tableID sqlbase.ID, index *sqlbase.IndexDescriptor,
) sqlbase.Spans {
prefix := roachpb.Key(sqlbase.MakeIndexKeyPrefix(tableID, index.ID))
// We have one constraint per column, so each contributes something
// to the start and/or the end key of the span.
// But we also have (...) IN <tuple> constraints that span multiple columns.
// These constraints split each span, and that's how we can end up with
// multiple spans.
resultSpans := sqlbase.Spans{{
Start: append(roachpb.Key(nil), prefix...),
End: append(roachpb.Key(nil), prefix...),
}}
colIdx := 0
for i, c := range constraints {
// We perform special processing on the last end constraint to account for
// the exclusive nature of the scan end key.
lastEnd := (c.end != nil) &&
(i+1 == len(constraints) || constraints[i+1].end == nil)
// IN is handled separately.
if (c.start != nil && c.start.Operator == parser.In) ||
(c.end != nil && c.end.Operator == parser.In) {
resultSpans = applyInConstraint(resultSpans, c, colIdx, index, lastEnd)
} else {
dir, err := index.ColumnDirections[colIdx].ToEncodingDirection()
if err != nil {
panic(err)
}
if c.start != nil {
if dir == encoding.Ascending {
encodeStartConstraintAscending(resultSpans, c.start)
} else {
encodeStartConstraintDescending(resultSpans, c.start)
}
}
if c.end != nil {
if dir == encoding.Ascending {
encodeEndConstraintAscending(resultSpans, c.end, lastEnd)
} else {
encodeEndConstraintDescending(resultSpans, c.end, lastEnd)
}
}
}
colIdx += c.numColumns()
}
// If we had no end constraints, make it so that we scan the whole index.
if len(constraints) == 0 || constraints[0].end == nil {
for i := range resultSpans {
resultSpans[i].End = resultSpans[i].End.PrefixEnd()
}
}
// Remove any spans which are empty. This can happen for constraints such as
// "a > 1 AND a < 2" which we do not simplify to false but which is treated
// as "a >= 2 AND a < 2" for span generation.
n := 0
for _, s := range resultSpans {
if bytes.Compare(s.Start, s.End) < 0 {
resultSpans[n] = s
n++
}
}
return resultSpans[:n]
}
func TestTableReader(t *testing.T) {
defer leaktest.AfterTest(t)()
_, sqlDB, kvDB, cleanup := sqlutils.SetupServer(t)
defer cleanup()
if _, err := sqlDB.Exec(`
CREATE DATABASE test;
CREATE TABLE test.t (a INT PRIMARY KEY, b INT, c INT, d INT, INDEX bc (b, c));
INSERT INTO test.t VALUES (1, 10, 11, 12), (2, 20, 21, 22), (3, 30, 31, 32);
INSERT INTO test.t VALUES (4, 60, 61, 62), (5, 50, 51, 52), (6, 40, 41, 42);
`); err != nil {
t.Fatal(err)
}
td := sqlbase.GetTableDescriptor(kvDB, "test", "t")
ts := TableReaderSpec{
Table: *td,
IndexIdx: 0,
Reverse: false,
Spans: nil,
Filter: Expression{Expr: "$2 != 21"}, // c != 21
OutputColumns: []uint32{0, 3}, // a, d
}
txn := client.NewTxn(context.Background(), *kvDB)
out := &testingReceiver{}
tr, err := newTableReader(&ts, txn, out, parser.EvalContext{})
if err != nil {
t.Fatal(err)
}
tr.run()
if out.err != nil {
t.Fatal(out.err)
}
if !out.closed {
t.Fatalf("output rowReceiver not closed")
}
expected := "[[1 12] [3 32] [4 62] [5 52] [6 42]]"
if fmt.Sprintf("%s", out.rows) != expected {
t.Errorf("invalid results: %s, expected %s'", out.rows, expected)
}
// Read using the bc index
var span roachpb.Span
span.Key = roachpb.Key(sqlbase.MakeIndexKeyPrefix(td.ID, td.Indexes[0].ID))
span.EndKey = append(span.Key, encoding.EncodeVarintAscending(nil, 50)...)
ts = TableReaderSpec{
Table: *td,
IndexIdx: 1,
Reverse: true,
Spans: []TableReaderSpan{{Span: span}},
Filter: Expression{Expr: "$1 != 30"}, // b != 30
OutputColumns: []uint32{0, 2}, // a, c
}
out = &testingReceiver{}
tr, err = newTableReader(&ts, txn, out, parser.EvalContext{})
if err != nil {
t.Fatal(err)
}
tr.run()
if out.err != nil {
t.Fatal(out.err)
}
if !out.closed {
t.Fatalf("output rowReceiver not closed")
}
expected = "[[6 41] [2 21] [1 11]]"
if fmt.Sprintf("%s", out.rows) != expected {
t.Errorf("invalid results: %s, expected %s'", out.rows, expected)
}
}
func TestDropIndex(t *testing.T) {
defer leaktest.AfterTest(t)()
s, sqlDB, kvDB := setup(t)
defer cleanup(s, sqlDB)
if _, err := sqlDB.Exec(`
CREATE DATABASE t;
CREATE TABLE t.kv (k CHAR PRIMARY KEY, v CHAR);
CREATE INDEX foo on t.kv (v);
INSERT INTO t.kv VALUES ('c', 'e'), ('a', 'c'), ('b', 'd');
`); err != nil {
t.Fatal(err)
}
nameKey := sqlbase.MakeNameMetadataKey(keys.MaxReservedDescID+1, "kv")
gr, pErr := kvDB.Get(nameKey)
if pErr != nil {
t.Fatal(pErr)
}
if !gr.Exists() {
t.Fatalf("Name entry %q does not exist", nameKey)
}
descKey := sqlbase.MakeDescMetadataKey(sqlbase.ID(gr.ValueInt()))
desc := &sqlbase.Descriptor{}
if pErr := kvDB.GetProto(descKey, desc); pErr != nil {
t.Fatal(pErr)
}
tableDesc := desc.GetTable()
status, i, err := tableDesc.FindIndexByName("foo")
if err != nil {
t.Fatal(err)
}
if status != sqlbase.DescriptorActive {
t.Fatal("Index 'foo' is not active.")
}
indexPrefix := sqlbase.MakeIndexKeyPrefix(tableDesc.ID, tableDesc.Indexes[i].ID)
indexStartKey := roachpb.Key(indexPrefix)
indexEndKey := indexStartKey.PrefixEnd()
if kvs, err := kvDB.Scan(indexStartKey, indexEndKey, 0); err != nil {
t.Fatal(err)
} else if l := 3; len(kvs) != l {
t.Fatalf("expected %d key value pairs, but got %d", l, len(kvs))
}
if _, err := sqlDB.Exec(`DROP INDEX t.kv@foo`); err != nil {
t.Fatal(err)
}
if kvs, err := kvDB.Scan(indexStartKey, indexEndKey, 0); err != nil {
t.Fatal(err)
} else if l := 0; len(kvs) != l {
t.Fatalf("expected %d key value pairs, but got %d", l, len(kvs))
}
if err := kvDB.GetProto(descKey, desc); err != nil {
t.Fatal(err)
}
tableDesc = desc.GetTable()
if _, _, err := tableDesc.FindIndexByName("foo"); err == nil {
t.Fatalf("table descriptor still contains index after index is dropped")
}
if err != nil {
t.Fatal(err)
}
}
func (sc *SchemaChanger) truncateAndBackfillColumnsChunk(
added []sqlbase.ColumnDescriptor,
dropped []sqlbase.ColumnDescriptor,
defaultExprs []parser.TypedExpr,
evalCtx *parser.EvalContext,
sp sqlbase.Span,
) (roachpb.Key, bool, error) {
var curIndexKey roachpb.Key
done := false
err := sc.db.Txn(func(txn *client.Txn) error {
tableDesc, err := getTableDescFromID(txn, sc.tableID)
if err != nil {
return err
}
// Short circuit the backfill if the table has been deleted.
if tableDesc.Deleted() {
done = true
return nil
}
updateCols := append(added, dropped...)
fkTables := TablesNeededForFKs(*tableDesc, CheckUpdates)
for k := range fkTables {
if fkTables[k], err = getTableDescFromID(txn, k); err != nil {
return err
}
}
// TODO(dan): Tighten up the bound on the requestedCols parameter to
// makeRowUpdater.
requestedCols := make([]sqlbase.ColumnDescriptor, 0, len(tableDesc.Columns)+len(added))
requestedCols = append(requestedCols, tableDesc.Columns...)
requestedCols = append(requestedCols, added...)
ru, err := makeRowUpdater(
txn, tableDesc, fkTables, updateCols, requestedCols, rowUpdaterOnlyColumns,
)
if err != nil {
return err
}
// TODO(dan): This check is an unfortunate bleeding of the internals of
// rowUpdater. Extract the sql row to k/v mapping logic out into something
// usable here.
if !ru.isColumnOnlyUpdate() {
panic("only column data should be modified, but the rowUpdater is configured otherwise")
}
// Run a scan across the table using the primary key. Running
// the scan and applying the changes in many transactions is
// fine because the schema change is in the correct state to
// handle intermediate OLTP commands which delete and add
// values during the scan.
var rf sqlbase.RowFetcher
colIDtoRowIndex := colIDtoRowIndexFromCols(tableDesc.Columns)
valNeededForCol := make([]bool, len(tableDesc.Columns))
for i := range valNeededForCol {
_, valNeededForCol[i] = ru.fetchColIDtoRowIndex[tableDesc.Columns[i].ID]
}
err = rf.Init(tableDesc, colIDtoRowIndex, &tableDesc.PrimaryIndex, false, false,
tableDesc.Columns, valNeededForCol)
if err != nil {
return err
}
// StartScan uses 0 as a sentinal for the default limit of entries scanned.
if err := rf.StartScan(txn, sqlbase.Spans{sp}, 0); err != nil {
return err
}
indexKeyPrefix := sqlbase.MakeIndexKeyPrefix(tableDesc, tableDesc.PrimaryIndex.ID)
oldValues := make(parser.DTuple, len(ru.fetchCols))
updateValues := make(parser.DTuple, len(updateCols))
writeBatch := &client.Batch{}
var i int
for ; i < ColumnTruncateAndBackfillChunkSize; i++ {
row, err := rf.NextRow()
if err != nil {
return err
}
if row == nil {
break // Done
}
curIndexKey, _, err = sqlbase.EncodeIndexKey(
tableDesc, &tableDesc.PrimaryIndex, colIDtoRowIndex, row, indexKeyPrefix)
for j, col := range added {
if defaultExprs == nil || defaultExprs[j] == nil {
updateValues[j] = parser.DNull
} else {
updateValues[j], err = defaultExprs[j].Eval(evalCtx)
if err != nil {
return err
}
}
if !col.Nullable && updateValues[j].Compare(parser.DNull) == 0 {
return sqlbase.NewNonNullViolationError(col.Name)
}
}
for j := range dropped {
updateValues[j+len(added)] = parser.DNull
}
copy(oldValues, row)
for j := len(row); j < len(oldValues); j++ {
oldValues[j] = parser.DNull
}
if _, err := ru.updateRow(writeBatch, oldValues, updateValues); err != nil {
return err
}
}
if i < ColumnTruncateAndBackfillChunkSize {
done = true
}
if err := txn.Run(writeBatch); err != nil {
return convertBackfillError(tableDesc, writeBatch)
}
return nil
})
return curIndexKey.PrefixEnd(), done, err
}
func TestTableReader(t *testing.T) {
defer leaktest.AfterTest(t)()
ctx, _ := createTestServerContext()
server, sqlDB, kvDB := setupWithContext(t, ctx)
defer cleanup(server, sqlDB)
if _, err := sqlDB.Exec(`
CREATE DATABASE test;
CREATE TABLE test.t (a INT PRIMARY KEY, b INT, c INT, d INT, INDEX bc (b, c));
INSERT INTO test.t VALUES (1, 10, 11, 12), (2, 20, 21, 22), (3, 30, 31, 32);
INSERT INTO test.t VALUES (4, 60, 61, 62), (5, 50, 51, 52), (6, 40, 41, 42);
`); err != nil {
t.Fatal(err)
}
td := getTableDescriptor(kvDB, "test", "t")
ts := sql.TableReaderSpec{
Table: *td,
IndexIdx: 0,
Reverse: false,
Spans: nil,
Filter: sql.Expression{Expr: "$2 != 21"}, // c != 21
OutputColumns: []uint32{0, 3}, // a, d
}
txn := client.NewTxn(context.Background(), *kvDB)
tr, err := sql.NewTableReader(&ts, txn, parser.EvalContext{})
if err != nil {
t.Fatal(err)
}
pErr := tr.Run()
if pErr != nil {
t.Fatal(pErr)
}
// TODO(radu): currently the table reader just prints out stuff; when it
// will output results we will be able to verify them.
// Expected output:
// RESULT: 1 <skipped> 11 12
// RESULT: 3 <skipped> 31 32
// RESULT: 4 <skipped> 61 62
// RESULT: 5 <skipped> 51 52
// RESULT: 6 <skipped> 41 42
// Read using the bc index
var span roachpb.Span
span.Key = roachpb.Key(sqlbase.MakeIndexKeyPrefix(td.ID, td.Indexes[0].ID))
span.EndKey = append(span.Key, encoding.EncodeVarintAscending(nil, 50)...)
ts = sql.TableReaderSpec{
Table: *td,
IndexIdx: 1,
Reverse: true,
Spans: []sql.TableReaderSpan{{Span: span}},
Filter: sql.Expression{Expr: "$1 != 30"}, // b != 30
OutputColumns: []uint32{0, 1}, // a, c
}
tr, err = sql.NewTableReader(&ts, txn, parser.EvalContext{})
if err != nil {
t.Fatal(err)
}
pErr = tr.Run()
if pErr != nil {
t.Fatal(pErr)
}
// Expected output:
// RESULT: 6 40 41 <skipped>
// RESULT: 2 20 21 <skipped>
// RESULT: 1 10 11 <skipped>
}