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)
}
}
}
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 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 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>
}