From e8505d8a53e3b58ea5ded0f7b399459ad5cbfdf4 Mon Sep 17 00:00:00 2001 From: tgruben Date: Wed, 4 Jan 2023 14:26:08 -0600 Subject: [PATCH] [FB-1831] distribute bulk insert to owning node (#2391) * distribute bulk insert to owning node --- batch/batch.go | 1 - importer.go | 28 ++++++- internal_client.go | 86 +++++++++++++------- sql3/sql_complex_test.go | 166 +++++++++++++++++++++++++++++++-------- 4 files changed, 214 insertions(+), 67 deletions(-) diff --git a/batch/batch.go b/batch/batch.go index 45bda43b8..80b06af24 100644 --- a/batch/batch.go +++ b/batch/batch.go @@ -552,7 +552,6 @@ func (b *Batch) Add(rec Row) error { // empty string is not a valid value at this point (Pilosa refuses to translate it) if val == "" { // b.rowIDs[i] = append(rowIDs, nilSentinel) - } else if rowID, ok := b.getRowTranslation(field.Name, val); ok { b.rowIDs[i] = append(rowIDs, rowID) } else { diff --git a/importer.go b/importer.go index 5753deb21..63e62e5b0 100644 --- a/importer.go +++ b/importer.go @@ -6,6 +6,8 @@ import ( "github.com/molecula/featurebase/v3/dax" "github.com/molecula/featurebase/v3/roaring" + "github.com/pkg/errors" + "golang.org/x/sync/errgroup" ) type Importer interface { @@ -32,12 +34,14 @@ var _ Importer = &onPremImporter{} // implemtation of the Importer interface does not use, and therefore they // intentionally no-op. type onPremImporter struct { - api *API + api *API + client *InternalClient } func NewOnPremImporter(api *API) *onPremImporter { return &onPremImporter{ - api: api, + api: api, + client: api.holder.executor.client, } } @@ -63,7 +67,25 @@ func (i *onPremImporter) ImportRoaringBitmap(ctx context.Context, tid dax.TableI } func (i *onPremImporter) ImportRoaringShard(ctx context.Context, tid dax.TableID, shard uint64, request *ImportRoaringShardRequest) error { - return i.api.ImportRoaringShard(ctx, string(tid), shard, request) + nodes, err := i.api.ShardNodes(ctx, string(tid), shard) + if err != nil { + return err + } + eg := errgroup.Group{} + for _, node := range nodes { + node := node + if node.ID == i.api.NodeID() { // local + eg.Go(func() error { + return i.api.ImportRoaringShard(ctx, string(tid), shard, request) + }) + } else { + eg.Go(func() error { // forward on + return i.client.ImportRoaringShard(ctx, &node.URI, string(tid), shard, true, request) + }) + } + } + err = eg.Wait() + return errors.Wrap(err, "importing") } func (i *onPremImporter) EncodeImportValues(ctx context.Context, tid dax.TableID, fld *dax.Field, shard uint64, vals []int64, ids []uint64, clear bool) (path string, data []byte, err error) { diff --git a/internal_client.go b/internal_client.go index 7dc844a60..8892f72dd 100644 --- a/internal_client.go +++ b/internal_client.go @@ -887,7 +887,6 @@ func (c *InternalClient) Import(ctx context.Context, qcx *Qcx, req *ImportReques func (c *InternalClient) ImportValue(ctx context.Context, qcx *Qcx, req *ImportValueRequest, options *ImportOptions) error { span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.Import") defer span.Finish() - if req.ColumnKeys != nil { req.Shard = ^uint64(0) } @@ -925,35 +924,7 @@ func (c *InternalClient) ImportRoaring(ctx context.Context, uri *pnet.URI, index return errors.Wrap(err, "marshal import request") } - // Generate HTTP request. - httpReq, err := http.NewRequest("POST", url, bytes.NewBuffer(data)) - if err != nil { - return errors.Wrap(err, "creating request") - } - httpReq.Header.Set("Content-Type", "application/x-protobuf") - httpReq.Header.Set("Accept", "application/x-protobuf") - httpReq.Header.Set("X-Pilosa-Row", "roaring") - httpReq.Header.Set("User-Agent", "pilosa/"+Version) - AddAuthToken(ctx, &httpReq.Header) - - // Execute request against the host. - resp, err := c.executeRequest(httpReq.WithContext(ctx)) - if err != nil { - return err - } - defer resp.Body.Close() - - dec := json.NewDecoder(resp.Body) - rbody := &ImportResponse{} - err = dec.Decode(rbody) - // Decode can return EOF when no error occurred. helpful! - if err != nil && err != io.EOF { - return errors.Wrap(err, "decoding response body") - } - if rbody.Err != "" { - return errors.Wrap(errors.New(rbody.Err), "importing roaring") - } - return nil + return c.executeProtobufRequest(ctx, url, data) } // ExportCSV bulk exports data for a single shard from a host to CSV format. @@ -2458,3 +2429,58 @@ func (c *InternalClient) getDiskUsage(ctx context.Context, index string) (DiskUs return sum, nil } + +func (c *InternalClient) executeProtobufRequest(ctx context.Context, url string, data []byte) error { + httpReq, err := http.NewRequest("POST", url, bytes.NewBuffer(data)) + if err != nil { + return errors.Wrap(err, "creating request") + } + httpReq.Header.Set("Content-Type", "application/x-protobuf") + httpReq.Header.Set("Accept", "application/x-protobuf") + httpReq.Header.Set("X-Pilosa-Row", "roaring") + httpReq.Header.Set("User-Agent", "pilosa/"+Version) + AddAuthToken(ctx, &httpReq.Header) + + // Execute request against the host. + resp, err := c.executeRequest(httpReq.WithContext(ctx)) + if err != nil { + return err + } + defer resp.Body.Close() + + dec := json.NewDecoder(resp.Body) + rbody := &ImportResponse{} + err = dec.Decode(rbody) + // Decode can return EOF when no error occurred. helpful! + if err != nil && err != io.EOF { + return errors.Wrap(err, "decoding response body") + } + if rbody.Err != "" { + return errors.Wrap(errors.New(rbody.Err), "importing roaring") + } + return nil +} + +// ImportRoaringShard(ctx, node, string(tid), shard, request +func (c *InternalClient) ImportRoaringShard(ctx context.Context, uri *pnet.URI, index string, shard uint64, remote bool, req *ImportRoaringShardRequest) error { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.ImportRoaringShard") + defer span.Finish() + + if index == "" { + return ErrIndexRequired + } + if uri == nil { + uri = c.defaultURI + } + + vals := url.Values{} + vals.Set("remote", strconv.FormatBool(remote)) + url := fmt.Sprintf("%s%s/index/%s/shard/%d/import-roaring?%s", uri, c.prefix(), index, shard, vals.Encode()) + + // Marshal data to protobuf. + data, err := c.serializer.Marshal(req) + if err != nil { + return errors.Wrap(err, "marshal import roaring shard request") + } + return c.executeProtobufRequest(ctx, url, data) +} diff --git a/sql3/sql_complex_test.go b/sql3/sql_complex_test.go index 2785db540..89fe63770 100644 --- a/sql3/sql_complex_test.go +++ b/sql3/sql_complex_test.go @@ -5,6 +5,8 @@ import ( "context" "fmt" "os" + "reflect" + "sort" "strings" "testing" "time" @@ -19,7 +21,6 @@ import ( ) func TestPlanner_Misc(t *testing.T) { - d, err := pql.ParseDecimal("12.345678") if err != nil { t.Fatal(err) @@ -302,6 +303,7 @@ func TestPlanner_Show(t *testing.T) { } }) } + func TestPlanner_CoverCreateTable(t *testing.T) { c := test.MustRunCluster(t, 1) defer c.Close() @@ -348,7 +350,7 @@ func TestPlanner_CoverCreateTable(t *testing.T) { _, _, err := sql_test.MustQueryRows(t, server, sql) if assert.Error(t, err) { assert.Equal(t, fld.expErr, err.Error()) - //sql3.SQLErrConflictingColumnConstraint.Message + // sql3.SQLErrConflictingColumnConstraint.Message } } }) @@ -522,7 +524,7 @@ func TestPlanner_CoverCreateTable(t *testing.T) { schema, err := api.Schema(ctx, false) assert.NoError(t, err) - //spew.Dump(schema) + // spew.Dump(schema) // Get the fields from the FeatureBase schema. // fbFields is a map of fieldName to FieldInfo. @@ -759,7 +761,6 @@ func TestPlanner_AlterTable(t *testing.T) { t.Fatal(diff) } }) - } func TestPlanner_DropTable(t *testing.T) { @@ -819,7 +820,8 @@ func TestPlanner_ExpressionsInSelectListParen(t *testing.T) { Set(1, b=100) Set(2, a=20) Set(2, b=200) - `}); err != nil { + `, + }); err != nil { t.Fatal(err) } @@ -898,7 +900,8 @@ func TestPlanner_ExpressionsInSelectListLiterals(t *testing.T) { Set(1, d=10.3) Set(1, ts='2022-02-22T22:22:22Z') Set(1, str='foo') - `}); err != nil { + `, + }); err != nil { t.Fatal(err) } @@ -1044,7 +1047,8 @@ func TestPlanner_ExpressionsInSelectListCase(t *testing.T) { Set(1, d=10.3) Set(1, ts='2022-02-22T22:22:22Z') Set(1, str='foo') - `}); err != nil { + `, + }); err != nil { t.Fatal(err) } @@ -1127,7 +1131,8 @@ func TestPlanner_Select(t *testing.T) { Set(1, b=100) Set(2, a=20) Set(2, b=200) - `}); err != nil { + `, + }); err != nil { t.Fatal(err) } @@ -1315,7 +1320,8 @@ func TestPlanner_SelectOrderBy(t *testing.T) { Set(1, b=100) Set(2, a=20) Set(2, b=200) - `}); err != nil { + `, + }); err != nil { t.Fatal(err) } @@ -1500,7 +1506,6 @@ func TestPlanner_BulkInsert(t *testing.T) { }) t.Run("BulkCSVBadMap", func(t *testing.T) { - _, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `bulk insert into j (_id, a, b) map (0 id, 1 int, 10 int) from x'1,10,20 2,11,21 3,12,22 @@ -1517,7 +1522,6 @@ func TestPlanner_BulkInsert(t *testing.T) { }) t.Run("BulkCSVFileDefault", func(t *testing.T) { - tmpfile, err := os.CreateTemp("", "BulkCSVFileDefault.*.csv") if err != nil { t.Fatal(err) @@ -1540,7 +1544,6 @@ func TestPlanner_BulkInsert(t *testing.T) { }) t.Run("BulkCSVFileNoColumns", func(t *testing.T) { - tmpfile, err := os.CreateTemp("", "BulkCSVFileNoColumns.*.csv") if err != nil { t.Fatal(err) @@ -1570,7 +1573,6 @@ func TestPlanner_BulkInsert(t *testing.T) { }) t.Run("BulkCSVFileRowsLimit", func(t *testing.T) { - tmpfile, err := os.CreateTemp("", "BulkCSVFileDefault.*.csv") if err != nil { t.Fatal(err) @@ -1607,7 +1609,6 @@ func TestPlanner_BulkInsert(t *testing.T) { }, columns); diff != "" { t.Fatal(diff) } - }) t.Run("BulkCSVBlobDefault", func(t *testing.T) { @@ -1618,7 +1619,6 @@ func TestPlanner_BulkInsert(t *testing.T) { }) t.Run("BulkNDJsonBlobDefault", func(t *testing.T) { - _, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `bulk insert into j (_id, a, b) map ('$._id' id, '$.a' int, '$.b' int) from x'{ "_id": 1, "a": 10, "b": 20 } { "_id": 2, "a": 10, "b": 20 } @@ -1653,7 +1653,6 @@ func TestPlanner_BulkInsert(t *testing.T) { }) t.Run("BulkNDJsonFileDefault", func(t *testing.T) { - tmpfile, err := os.CreateTemp("", "BulkNDJsonFileDefault.*.csv") if err != nil { t.Fatal(err) @@ -1677,7 +1676,6 @@ func TestPlanner_BulkInsert(t *testing.T) { }) t.Run("BulkNDJsonFileTransform", func(t *testing.T) { - tmpfile, err := os.CreateTemp("", "BulkNDJsonFileTransform.*.csv") if err != nil { t.Fatal(err) @@ -1701,7 +1699,6 @@ func TestPlanner_BulkInsert(t *testing.T) { }) t.Run("BulkNDJsonAllTypes", func(t *testing.T) { - _, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `bulk insert into alltypes (_id, id1, i1, ids1, ss1, ts1, s1, b1, d1) @@ -1721,7 +1718,6 @@ func TestPlanner_BulkInsert(t *testing.T) { }) t.Run("BulkNDJsonBadJsonPath", func(t *testing.T) { - _, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `bulk insert into alltypes (_id, id1, i1, ids1, ss1, ts1, s1, b1, d1) @@ -1741,7 +1737,6 @@ func TestPlanner_BulkInsert(t *testing.T) { }) t.Run("BulkNDJsonBadJson", func(t *testing.T) { - _, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `bulk insert into alltypes (_id, id1, i1, ids1, ss1, ts1, s1, b1, d1) @@ -1761,7 +1756,6 @@ func TestPlanner_BulkInsert(t *testing.T) { }) t.Run("BulkInsertDecimals", func(t *testing.T) { - _, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `create table iris ( _id id, sepallength decimal(2), @@ -1811,7 +1805,6 @@ func TestPlanner_BulkInsert(t *testing.T) { if err != nil { t.Fatal(err) } - }) t.Run("BulkInsertDupeColumnPlusNullsInJson", func(t *testing.T) { @@ -1839,7 +1832,6 @@ func TestPlanner_BulkInsert(t *testing.T) { }) t.Run("BulkInsertCSVStringIDSet", func(t *testing.T) { - _, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `create table greg-test ( _id STRING, id_col ID, @@ -1898,7 +1890,6 @@ func TestPlanner_BulkInsert(t *testing.T) { }) t.Run("BulkInsertAllowMissingValues", func(t *testing.T) { - _, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `create table greg-test-amv ( _id STRING, id_col ID, @@ -1956,7 +1947,6 @@ func TestPlanner_BulkInsert(t *testing.T) { } }) t.Run("BulkInsertNDJSONStringIDSet", func(t *testing.T) { - _, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `create table greg-test-01 ( _id STRING, id_col ID, @@ -2073,7 +2063,8 @@ func TestPlanner_SelectSelectSource(t *testing.T) { Set(1, b=100) Set(2, a=20) Set(2, b=200) - `}); err != nil { + `, + }); err != nil { t.Fatal(err) } @@ -2153,7 +2144,8 @@ func TestPlanner_In(t *testing.T) { Set(1, a=10) Set(2, a=20) Set(3, a=30) - `}); err != nil { + `, + }); err != nil { t.Fatal(err) } @@ -2168,15 +2160,16 @@ func TestPlanner_In(t *testing.T) { Set(3, parentid=2) Set(3, x=300) - `}); err != nil { + `, + }); err != nil { t.Fatal(err) } t.Run("Count", func(t *testing.T) { t.Skip("Need to add join conditions to get this to pass") results, columns, err := sql_test.MustQueryRows(t, c.GetNode(0).Server, fmt.Sprintf(`SELECT %j._id, %j.a, %k._id, %k.parentid, %k.x FROM %j INNER JOIN %k ON %j._id = %k.parentid`, c, c, c, c, c, c, c, c, c)) - //results, columns, err := sql_test.MustQueryRows(t, c.GetNode(0).Server, fmt.Sprintf(`SELECT COUNT(*) FROM %j INNER JOIN %k ON %j._id = %k.parentid`, c, c, c, c)) - //results, columns, err := sql_test.MustQueryRows(t, c.GetNode(0).Server, fmt.Sprintf(`SELECT a FROM %j where a = 20`, c)) // SELECT COUNT(*) FROM %j INNER JOIN %k ON %j._id = %k.parentid + // results, columns, err := sql_test.MustQueryRows(t, c.GetNode(0).Server, fmt.Sprintf(`SELECT COUNT(*) FROM %j INNER JOIN %k ON %j._id = %k.parentid`, c, c, c, c)) + // results, columns, err := sql_test.MustQueryRows(t, c.GetNode(0).Server, fmt.Sprintf(`SELECT a FROM %j where a = 20`, c)) // SELECT COUNT(*) FROM %j INNER JOIN %k ON %j._id = %k.parentid if err != nil { t.Fatal(err) } @@ -2284,7 +2277,8 @@ func TestPlanner_Distinct(t *testing.T) { Set(1, a=10) Set(2, a=20) Set(3, a=30) - `}); err != nil { + `, + }); err != nil { t.Fatal(err) } @@ -2299,7 +2293,8 @@ func TestPlanner_Distinct(t *testing.T) { Set(3, parentid=2) Set(3, x=300) - `}); err != nil { + `, + }); err != nil { t.Fatal(err) } @@ -2396,7 +2391,8 @@ func TestPlanner_SelectTop(t *testing.T) { Set(1, b=100) Set(2, b=200) Set(3, b=300) - `}); err != nil { + `, + }); err != nil { t.Fatal(err) } @@ -2454,6 +2450,7 @@ func wireQueryFieldID(name string) *pilosa.WireQueryField { BaseType: dax.BaseTypeID, } } + func wireQueryFieldBool(name string) *pilosa.WireQueryField { return &pilosa.WireQueryField{ Name: dax.FieldName(name), @@ -2461,6 +2458,7 @@ func wireQueryFieldBool(name string) *pilosa.WireQueryField { BaseType: dax.BaseTypeBool, } } + func wireQueryFieldString(name string) *pilosa.WireQueryField { return &pilosa.WireQueryField{ Name: dax.FieldName(name), @@ -2468,6 +2466,7 @@ func wireQueryFieldString(name string) *pilosa.WireQueryField { BaseType: dax.BaseTypeString, } } + func wireQueryFieldInt(name string) *pilosa.WireQueryField { return &pilosa.WireQueryField{ Name: dax.FieldName(name), @@ -2475,6 +2474,7 @@ func wireQueryFieldInt(name string) *pilosa.WireQueryField { BaseType: dax.BaseTypeInt, } } + func wireQueryFieldTimestamp(name string) *pilosa.WireQueryField { return &pilosa.WireQueryField{ Name: dax.FieldName(name), @@ -2482,6 +2482,7 @@ func wireQueryFieldTimestamp(name string) *pilosa.WireQueryField { BaseType: dax.BaseTypeTimestamp, } } + func wireQueryFieldDecimal(name string, scale int64) *pilosa.WireQueryField { return &pilosa.WireQueryField{ Name: dax.FieldName(name), @@ -2492,3 +2493,102 @@ func wireQueryFieldDecimal(name string, scale int64) *pilosa.WireQueryField { }, } } + +// This test verifies that data sent to all nodes shows up in the results +func TestPlanner_BulkInsert_FB1831(t *testing.T) { + c := test.MustRunCluster(t, 3) + defer c.Close() + + _, _, err := sql_test.MustQueryRows(t, c.GetNode(0).Server, `create table iris (_id id, sepallength decimal(2), sepalwidth decimal(2), petallength decimal(2), petalwidth decimal(2), species string cachetype ranked size 1000);`) + if err != nil { + t.Fatal(err) + } + + _, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `bulk insert + into iris (_id, sepallength, sepalwidth, petallength, petalwidth, species) + map('id' id, + 'sepalLength' DECIMAL(2), + 'sepalWidth' DECIMAL(2), + 'petalLength' DECIMAL(2), + 'petalWidth' DECIMAL(2), + 'species' STRING) + from + x'{"id": 1, "sepalLength": "5.1", "sepalWidth": "3.5", "petalLength": "1.4", "petalWidth": "0.2", "species": "setosa"} + {"id": 2, "sepalLength": "4.9", "sepalWidth": "3.0", "petalLength": "1.4", "petalWidth": "0.2", "species": "setosa"} + {"id": 3, "sepalLength": "4.7", "sepalWidth": "3.2", "petalLength": "1.3", "petalWidth": "0.2", "species": "setosa"}' + with + format 'NDJSON' + input 'STREAM';`) + if err != nil { + t.Fatal(err) + } + _, _, err = sql_test.MustQueryRows(t, c.GetNode(1).Server, `bulk insert + into iris (_id, sepallength, sepalwidth, petallength, petalwidth, species) + map('id' id, + 'sepalLength' DECIMAL(2), + 'sepalWidth' DECIMAL(2), + 'petalLength' DECIMAL(2), + 'petalWidth' DECIMAL(2), + 'species' STRING) + from + x'{"id": 4, "sepalLength": "5.1", "sepalWidth": "3.5", "petalLength": "1.4", "petalWidth": "0.2", "species": "setosa"} + {"id": 5, "sepalLength": "4.9", "sepalWidth": "3.0", "petalLength": "1.4", "petalWidth": "0.2", "species": "setosa"} + {"id": 6, "sepalLength": "4.7", "sepalWidth": "3.2", "petalLength": "1.3", "petalWidth": "0.2", "species": "setosa"}' + with + format 'NDJSON' + input 'STREAM';`) + if err != nil { + t.Fatal(err) + } + _, _, err = sql_test.MustQueryRows(t, c.GetNode(2).Server, `bulk insert + into iris (_id, sepallength, sepalwidth, petallength, petalwidth, species) + map('id' id, + 'sepalLength' DECIMAL(2), + 'sepalWidth' DECIMAL(2), + 'petalLength' DECIMAL(2), + 'petalWidth' DECIMAL(2), + 'species' STRING) + from + x'{"id": 7, "sepalLength": "5.1", "sepalWidth": "3.5", "petalLength": "1.4", "petalWidth": "0.2", "species": "setosa"} + {"id": 8, "sepalLength": "4.9", "sepalWidth": "3.0", "petalLength": "1.4", "petalWidth": "0.2", "species": "setosa"} + {"id": 9, "sepalLength": "4.7", "sepalWidth": "3.2", "petalLength": "1.3", "petalWidth": "0.2", "species": "setosa"}' + with + format 'NDJSON' + input 'STREAM';`) + if err != nil { + t.Fatal(err) + } + _, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `bulk insert + into iris (_id, sepallength, sepalwidth, petallength, petalwidth, species) + map('id' id, + 'sepalLength' DECIMAL(2), + 'sepalWidth' DECIMAL(2), + 'petalLength' DECIMAL(2), + 'petalWidth' DECIMAL(2), + 'species' STRING) + from + x'{"id": 1048577, "sepalLength": "5.1", "sepalWidth": "3.5", "petalLength": "1.4", "petalWidth": "0.2", "species": "setosa"} + {"id": 2097153, "sepalLength": "4.9", "sepalWidth": "3.0", "petalLength": "1.4", "petalWidth": "0.2", "species": "setosa"} + {"id": 3145729, "sepalLength": "4.7", "sepalWidth": "3.2", "petalLength": "1.3", "petalWidth": "0.2", "species": "setosa"}' + with + format 'NDJSON' + input 'STREAM';`) + if err != nil { + t.Fatal(err) + } + results, _, err := sql_test.MustQueryRows(t, c.GetNode(0).Server, `select _id from iris`) + if err != nil { + t.Fatal(err) + } + got := make([]int64, 0) + for i := range results { + got = append(got, results[i][0].(int64)) + } + sort.Slice(got, func(i, j int) bool { + return got[i] < got[j] + }) + expected := []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 1048577, 2097153, 3145729} + if !reflect.DeepEqual(got, expected) { + t.Fatal("Expecting to be equal") + } +}