From caf8e067129b8c7fdcbd44c6b1c937dcc67cbc74 Mon Sep 17 00:00:00 2001 From: Travis Turner Date: Mon, 22 Oct 2018 13:35:32 -0500 Subject: [PATCH] add `clear` functional option for imports --- api.go | 36 +++++- client.go | 8 +- cmd/import.go | 1 + cmd/import_test.go | 11 ++ ctl/import.go | 7 +- ctl/import_test.go | 67 +++++++--- field.go | 22 +++- fragment.go | 26 ++-- fragment_internal_test.go | 250 +++++++++++++++++++++++++++++++++++--- http/client.go | 46 +++++-- http/handler.go | 10 +- 11 files changed, 413 insertions(+), 71 deletions(-) diff --git a/api.go b/api.go index 0b27d19a8..bb900db4e 100644 --- a/api.go +++ b/api.go @@ -670,12 +670,36 @@ func (api *API) FieldAttrDiff(_ context.Context, indexName string, fieldName str return attrs, nil } +// ImportOptions holds the options for the API.Import method. +type ImportOptions struct { + Clear bool +} + +// ImportOption is a functional option type for API.Import +type ImportOption func(*ImportOptions) error + +func OptImportOptionsClear(c bool) ImportOption { + return func(o *ImportOptions) error { + o.Clear = c + return nil + } +} + // Import bulk imports data into a particular index,field,shard. -func (api *API) Import(_ context.Context, req *ImportRequest) error { +func (api *API) Import(_ context.Context, req *ImportRequest, opts ...ImportOption) error { if err := api.validate(apiImport); err != nil { return errors.Wrap(err, "validating api method") } + // Set up import options. + options := &ImportOptions{} + for _, opt := range opts { + err := opt(options) + if err != nil { + return errors.Wrap(err, "applying option") + } + } + index := api.holder.Index(req.Index) if index == nil { return newNotFoundError(ErrIndexNotFound) @@ -717,13 +741,15 @@ func (api *API) Import(_ context.Context, req *ImportRequest) error { } // Import columnIDs into existence field. - if err := importExistenceColumns(index, req.ColumnIDs); err != nil { - api.server.logger.Printf("import existence error: index=%s, field=%s, shard=%d, columns=%d, err=%s", req.Index, req.Field, req.Shard, len(req.ColumnIDs), err) - return errors.Wrap(err, "importing existence columns") + if !options.Clear { + if err := importExistenceColumns(index, req.ColumnIDs); err != nil { + api.server.logger.Printf("import existence error: index=%s, field=%s, shard=%d, columns=%d, err=%s", req.Index, req.Field, req.Shard, len(req.ColumnIDs), err) + return errors.Wrap(err, "importing existence columns") + } } // Import into fragment. - err = field.Import(req.RowIDs, req.ColumnIDs, timestamps) + err = field.Import(req.RowIDs, req.ColumnIDs, timestamps, opts...) if err != nil { api.server.logger.Printf("import error: index=%s, field=%s, shard=%d, columns=%d, err=%s", req.Index, req.Field, req.Shard, len(req.ColumnIDs), err) } diff --git a/client.go b/client.go index a2ca0bf96..889544557 100644 --- a/client.go +++ b/client.go @@ -37,8 +37,8 @@ type InternalClient interface { Nodes(ctx context.Context) ([]*Node, error) Query(ctx context.Context, index string, queryRequest *QueryRequest) (*QueryResponse, error) QueryNode(ctx context.Context, uri *URI, index string, queryRequest *QueryRequest) (*QueryResponse, error) - Import(ctx context.Context, index, field string, shard uint64, bits []Bit) error - ImportK(ctx context.Context, index, field string, bits []Bit) error + Import(ctx context.Context, index, field string, shard uint64, bits []Bit, opts ...ImportOption) error + ImportK(ctx context.Context, index, field string, bits []Bit, opts ...ImportOption) error EnsureIndex(ctx context.Context, name string, options IndexOptions) error EnsureField(ctx context.Context, indexName string, fieldName string) error EnsureFieldWithOptions(ctx context.Context, index, field string, opt FieldOptions) error @@ -103,10 +103,10 @@ func (n nopInternalClient) Query(ctx context.Context, index string, queryRequest func (n nopInternalClient) QueryNode(ctx context.Context, uri *URI, index string, queryRequest *QueryRequest) (*QueryResponse, error) { return nil, nil } -func (n nopInternalClient) Import(ctx context.Context, index, field string, shard uint64, bits []Bit) error { +func (n nopInternalClient) Import(ctx context.Context, index, field string, shard uint64, bits []Bit, opts ...ImportOption) error { return nil } -func (n nopInternalClient) ImportK(ctx context.Context, index, field string, bits []Bit) error { +func (n nopInternalClient) ImportK(ctx context.Context, index, field string, bits []Bit, opts ...ImportOption) error { return nil } func (n nopInternalClient) ImportRoaring(ctx context.Context, uri *URI, index, field string, shard uint64, remote bool, data []byte) error { diff --git a/cmd/import.go b/cmd/import.go index 7bfca1a84..878aad286 100644 --- a/cmd/import.go +++ b/cmd/import.go @@ -63,6 +63,7 @@ omitted. If it is present then its format should be YYYY-MM-DDTHH:MM. flags.IntVarP(&Importer.BufferSize, "buffer-size", "s", 10000000, "Number of bits to buffer/sort before importing.") flags.BoolVarP(&Importer.Sort, "sort", "", false, "Enables sorting before import.") flags.BoolVarP(&Importer.CreateSchema, "create", "e", false, "Create the schema if it does not exist before import.") + flags.BoolVarP(&Importer.Clear, "clear", "", false, "Clear the data provided in the import.") ctl.SetTLSConfig(flags, &Importer.TLS.CertificatePath, &Importer.TLS.CertificateKeyPath, &Importer.TLS.SkipVerify) return importCmd diff --git a/cmd/import_test.go b/cmd/import_test.go index ba015dc04..7b566ce65 100644 --- a/cmd/import_test.go +++ b/cmd/import_test.go @@ -95,6 +95,17 @@ field = "f1" return v.Error() }, }, + { + args: []string{"import", "--index", "i1", "--field", "f1", "--clear", "true"}, + env: map[string]string{}, + validation: func() error { + v := validator{} + v.Check(cmd.Importer.Index, "i1") + v.Check(cmd.Importer.Field, "f1") + v.Check(cmd.Importer.Clear, true) + return v.Error() + }, + }, } executeDry(t, tests) } diff --git a/ctl/import.go b/ctl/import.go index d2694a64f..8f0f4a4e5 100644 --- a/ctl/import.go +++ b/ctl/import.go @@ -49,6 +49,9 @@ type ImportCommand struct { // nolint: maligned // CreateSchema ensures the schema exists before import CreateSchema bool + // Clear clears the import data as opposed to setting it. + Clear bool + // Filenames to import from. Paths []string `json:"paths"` @@ -255,7 +258,7 @@ func (cmd *ImportCommand) importBits(ctx context.Context, useColumnKeys, useRowK // If keys are used, all bits are sent to the primary translate store (i.e. coordinator). if useColumnKeys || useRowKeys { logger.Printf("importing keys: n=%d", len(bits)) - if err := cmd.client.ImportK(ctx, cmd.Index, cmd.Field, bits); err != nil { + if err := cmd.client.ImportK(ctx, cmd.Index, cmd.Field, bits, pilosa.OptImportOptionsClear(cmd.Clear)); err != nil { return errors.Wrap(err, "importing keys") } return nil @@ -272,7 +275,7 @@ func (cmd *ImportCommand) importBits(ctx context.Context, useColumnKeys, useRowK } logger.Printf("importing shard: %d, n=%d", shard, len(chunk)) - if err := cmd.client.Import(ctx, cmd.Index, cmd.Field, shard, chunk); err != nil { + if err := cmd.client.Import(ctx, cmd.Index, cmd.Field, shard, chunk, pilosa.OptImportOptionsClear(cmd.Clear)); err != nil { return errors.Wrap(err, "importing") } } diff --git a/ctl/import_test.go b/ctl/import_test.go index 2913005fb..fdf73720a 100644 --- a/ctl/import_test.go +++ b/ctl/import_test.go @@ -50,28 +50,55 @@ func TestImportCommand_Validation(t *testing.T) { } } -func TestImportCommand_Run(t *testing.T) { - buf := bytes.Buffer{} - stdin, stdout, stderr := GetIO(buf) - cm := NewImportCommand(stdin, stdout, stderr) - file, err := ioutil.TempFile("", "import.csv") - file.Write([]byte("1,2\n3,4\n5,6")) - ctx := context.Background() - if err != nil { - t.Fatal(err) - } +func TestImportCommand_Basic(t *testing.T) { + t.Run("set", func(t *testing.T) { + buf := bytes.Buffer{} + stdin, stdout, stderr := GetIO(buf) + cm := NewImportCommand(stdin, stdout, stderr) + file, err := ioutil.TempFile("", "import.csv") + file.Write([]byte("1,2\n3,4\n5,6")) + ctx := context.Background() + if err != nil { + t.Fatal(err) + } - cmd := test.MustRunCluster(t, 1)[0] - cm.Host = cmd.API.Node().URI.HostPort() + cmd := test.MustRunCluster(t, 1)[0] + cm.Host = cmd.API.Node().URI.HostPort() - cm.Index = "i" - cm.Field = "f" - cm.CreateSchema = true - cm.Paths = []string{file.Name()} - err = cm.Run(ctx) - if err != nil { - t.Fatalf("Import Run doesn't work: %s", err) - } + cm.Index = "i" + cm.Field = "f" + cm.CreateSchema = true + cm.Paths = []string{file.Name()} + err = cm.Run(ctx) + if err != nil { + t.Fatalf("Import Run doesn't work: %s", err) + } + }) + + t.Run("clear", func(t *testing.T) { + buf := bytes.Buffer{} + stdin, stdout, stderr := GetIO(buf) + cm := NewImportCommand(stdin, stdout, stderr) + file, err := ioutil.TempFile("", "import.csv") + file.Write([]byte("1,2\n3,4\n5,6")) + ctx := context.Background() + if err != nil { + t.Fatal(err) + } + + cmd := test.MustRunCluster(t, 1)[0] + cm.Host = cmd.API.Node().URI.HostPort() + + cm.Index = "i" + cm.Field = "f" + cm.CreateSchema = true + cm.Clear = true + cm.Paths = []string{file.Name()} + err = cm.Run(ctx) + if err != nil { + t.Fatalf("Import Run clear doesn't work: %s", err) + } + }) } // Ensure that the ImportValue path runs. diff --git a/field.go b/field.go index 63bb08b63..1a880a266 100644 --- a/field.go +++ b/field.go @@ -1047,11 +1047,25 @@ func (f *Field) Range(name string, op pql.Token, predicate int64) (*Row, error) } // Import bulk imports data. -func (f *Field) Import(rowIDs, columnIDs []uint64, timestamps []*time.Time) error { +func (f *Field) Import(rowIDs, columnIDs []uint64, timestamps []*time.Time, opts ...ImportOption) error { + + // Set up import options. + options := &ImportOptions{} + for _, opt := range opts { + err := opt(options) + if err != nil { + return errors.Wrap(err, "applying option") + } + } + // Determine quantum if timestamps are set. q := f.TimeQuantum() - if hasTime(timestamps) && q == "" { - return errors.New("time quantum not set in field") + if hasTime(timestamps) { + if q == "" { + return errors.New("time quantum not set in field") + } else if options.Clear { + return errors.New("import clear is not supported with timestamps") + } } fieldType := f.Type() @@ -1103,7 +1117,7 @@ func (f *Field) Import(rowIDs, columnIDs []uint64, timestamps []*time.Time) erro return errors.Wrap(err, "creating view") } - if err := frag.bulkImport(data.RowIDs, data.ColumnIDs); err != nil { + if err := frag.bulkImport(data.RowIDs, data.ColumnIDs, options); err != nil { return err } } diff --git a/fragment.go b/fragment.go index 63c9cce37..94e45ab1e 100644 --- a/fragment.go +++ b/fragment.go @@ -1418,20 +1418,20 @@ func (f *fragment) mergeBlock(id int, data []pairSet) (sets, clears []pairSet, e // bulkImport bulk imports a set of bits and then snapshots the storage. // The cache is updated to reflect the new data. -func (f *fragment) bulkImport(rowIDs, columnIDs []uint64) error { +func (f *fragment) bulkImport(rowIDs, columnIDs []uint64, options *ImportOptions) error { // Verify that there are an equal number of row ids and column ids. if len(rowIDs) != len(columnIDs) { return fmt.Errorf("mismatch of row/column len: %d != %d", len(rowIDs), len(columnIDs)) } - if f.mutexVector != nil { - return f.bulkImportMutex(rowIDs, columnIDs) + if f.mutexVector != nil && !options.Clear { + return f.bulkImportMutex(rowIDs, columnIDs, options) } - return f.bulkImportStandard(rowIDs, columnIDs) + return f.bulkImportStandard(rowIDs, columnIDs, options) } // bulkImportStandard performs a bulk import on a standard fragment. -func (f *fragment) bulkImportStandard(rowIDs, columnIDs []uint64) error { +func (f *fragment) bulkImportStandard(rowIDs, columnIDs []uint64, options *ImportOptions) error { // Create a temporary bitmap which will be populated by rowIDs and columnIDs // and then merged into the existing fragment's bitmap. localBitmap := roaring.NewBitmap() @@ -1480,10 +1480,18 @@ func (f *fragment) bulkImportStandard(rowIDs, columnIDs []uint64) error { // Merge localBitmap into fragment's existing data. var results *roaring.Bitmap - if f.storage.Count() > 0 { - results = f.storage.Union(localBitmap) + if options.Clear { + if f.storage.Count() > 0 { + results = f.storage.Difference(localBitmap) + } else { + results = roaring.NewBitmap() + } } else { - results = localBitmap + if f.storage.Count() > 0 { + results = f.storage.Union(localBitmap) + } else { + results = localBitmap + } } // Update cache counts for all affected rows. @@ -1500,7 +1508,7 @@ func (f *fragment) bulkImportStandard(rowIDs, columnIDs []uint64) error { // mutex restrictions. Because the mutex requirements must be checked // against storage, this method must acquire a write lock on the fragment // during the entire process, and it handles every bit independently. -func (f *fragment) bulkImportMutex(rowIDs, columnIDs []uint64) error { +func (f *fragment) bulkImportMutex(rowIDs, columnIDs []uint64, options *ImportOptions) error { f.mu.Lock() defer f.mu.Unlock() diff --git a/fragment_internal_test.go b/fragment_internal_test.go index be0db3b01..67ad70f45 100644 --- a/fragment_internal_test.go +++ b/fragment_internal_test.go @@ -1253,12 +1253,15 @@ func TestFragment_SetMutex(t *testing.T) { } } -// Ensure a fragment can import mutually exclusive values. -func TestFragment_ImportMutex(t *testing.T) { +// Ensure a fragment can import into set fields. +func TestFragment_ImportSet(t *testing.T) { tests := []struct { - rowIDs []uint64 - colIDs []uint64 - exp map[uint64][]uint64 + setRowIDs []uint64 + setColIDs []uint64 + setExp map[uint64][]uint64 + clearRowIDs []uint64 + clearColIDs []uint64 + clearExp map[uint64][]uint64 }{ { []uint64{1, 1, 1, 1}, @@ -1266,6 +1269,129 @@ func TestFragment_ImportMutex(t *testing.T) { map[uint64][]uint64{ 1: {0, 1, 2, 3}, }, + []uint64{}, + []uint64{}, + map[uint64][]uint64{ + 1: {0, 1, 2, 3}, + }, + }, + { + []uint64{1, 1, 1, 1, 2, 2, 2, 2}, + []uint64{0, 1, 2, 3, 0, 1, 2, 3}, + map[uint64][]uint64{ + 1: {0, 1, 2, 3}, + 2: {0, 1, 2, 3}, + }, + []uint64{1, 1, 2}, + []uint64{1, 2, 3}, + map[uint64][]uint64{ + 1: {0, 3}, + 2: {0, 1, 2}, + }, + }, + { + []uint64{1, 1, 1, 1, 2}, + []uint64{0, 1, 2, 3, 1}, + map[uint64][]uint64{ + 1: {0, 1, 2, 3}, + 2: {1}, + }, + []uint64{1, 1, 1, 1}, + []uint64{0, 1, 2, 3}, + map[uint64][]uint64{ + 1: {}, + 2: {1}, + }, + }, + { + []uint64{1, 1, 1, 1, 2, 2, 1}, + []uint64{0, 1, 2, 3, 1, 8, 1}, + map[uint64][]uint64{ + 1: {0, 1, 2, 3}, + 2: {1, 8}, + }, + []uint64{1, 1}, + []uint64{0, 0}, + map[uint64][]uint64{ + 1: {1, 2, 3}, + 2: {1, 8}, + }, + }, + { + []uint64{1, 2, 3}, + []uint64{8, 8, 8}, + map[uint64][]uint64{ + 1: {8}, + 2: {8}, + 3: {8}, + }, + []uint64{1, 2, 3}, + []uint64{9, 9, 9}, + map[uint64][]uint64{ + 1: {8}, + 2: {8}, + 3: {8}, + }, + }, + } + + for i, test := range tests { + t.Run(fmt.Sprintf("importset%d", i), func(t *testing.T) { + f := mustOpenFragment("i", "f", viewStandard, 0, "") + defer f.Close() + + // Set import. + err := f.bulkImport(test.setRowIDs, test.setColIDs, &ImportOptions{}) + if err != nil { + t.Fatalf("bulk importing ids: %v", err) + } + + // Check for expected results. + for k, v := range test.setExp { + cols := f.row(k).Columns() + if !reflect.DeepEqual(cols, v) { + t.Fatalf("expected: %v, but got: %v", v, cols) + } + } + + // Clear import. + err = f.bulkImport(test.clearRowIDs, test.clearColIDs, &ImportOptions{Clear: true}) + if err != nil { + t.Fatalf("bulk clearing ids: %v", err) + } + + // Check for expected results. + for k, v := range test.clearExp { + cols := f.row(k).Columns() + if !reflect.DeepEqual(cols, v) { + t.Fatalf("expected: %v, but got: %v", v, cols) + } + } + }) + } +} + +// Ensure a fragment can import mutually exclusive values. +func TestFragment_ImportMutex(t *testing.T) { + tests := []struct { + setRowIDs []uint64 + setColIDs []uint64 + setExp map[uint64][]uint64 + clearRowIDs []uint64 + clearColIDs []uint64 + clearExp map[uint64][]uint64 + }{ + { + []uint64{1, 1, 1, 1}, + []uint64{0, 1, 2, 3}, + map[uint64][]uint64{ + 1: {0, 1, 2, 3}, + }, + []uint64{}, + []uint64{}, + map[uint64][]uint64{ + 1: {0, 1, 2, 3}, + }, }, { []uint64{1, 1, 1, 1, 2, 2, 2, 2}, @@ -1274,6 +1400,12 @@ func TestFragment_ImportMutex(t *testing.T) { 1: {}, 2: {0, 1, 2, 3}, }, + []uint64{1, 1, 2}, + []uint64{1, 2, 3}, + map[uint64][]uint64{ + 1: {}, + 2: {0, 1, 2}, + }, }, { []uint64{1, 1, 1, 1, 2}, @@ -1282,6 +1414,12 @@ func TestFragment_ImportMutex(t *testing.T) { 1: {0, 2, 3}, 2: {1}, }, + []uint64{1, 1, 1, 1}, + []uint64{0, 1, 2, 3}, + map[uint64][]uint64{ + 1: {}, + 2: {1}, + }, }, { []uint64{1, 1, 1, 1, 2, 2, 1}, @@ -1290,6 +1428,12 @@ func TestFragment_ImportMutex(t *testing.T) { 1: {0, 1, 2, 3}, 2: {8}, }, + []uint64{1, 1}, + []uint64{0, 0}, + map[uint64][]uint64{ + 1: {1, 2, 3}, + 2: {8}, + }, }, { []uint64{1, 2, 3}, @@ -1299,6 +1443,13 @@ func TestFragment_ImportMutex(t *testing.T) { 2: {}, 3: {8}, }, + []uint64{1, 2, 3}, + []uint64{9, 9, 9}, + map[uint64][]uint64{ + 1: {}, + 2: {}, + 3: {8}, + }, }, } @@ -1307,13 +1458,28 @@ func TestFragment_ImportMutex(t *testing.T) { f := mustOpenMutexFragment("i", "f", viewStandard, 0, "") defer f.Close() - err := f.bulkImport(test.rowIDs, test.colIDs) + // Set import. + err := f.bulkImport(test.setRowIDs, test.setColIDs, &ImportOptions{}) if err != nil { t.Fatalf("bulk importing ids: %v", err) } // Check for expected results. - for k, v := range test.exp { + for k, v := range test.setExp { + cols := f.row(k).Columns() + if !reflect.DeepEqual(cols, v) { + t.Fatalf("expected: %v, but got: %v", v, cols) + } + } + + // Clear import. + err = f.bulkImport(test.clearRowIDs, test.clearColIDs, &ImportOptions{Clear: true}) + if err != nil { + t.Fatalf("bulk clearing ids: %v", err) + } + + // Check for expected results. + for k, v := range test.clearExp { cols := f.row(k).Columns() if !reflect.DeepEqual(cols, v) { t.Fatalf("expected: %v, but got: %v", v, cols) @@ -1326,9 +1492,12 @@ func TestFragment_ImportMutex(t *testing.T) { // Ensure a fragment can import bool values. func TestFragment_ImportBool(t *testing.T) { tests := []struct { - rowIDs []uint64 - colIDs []uint64 - exp map[uint64][]uint64 + setRowIDs []uint64 + setColIDs []uint64 + setExp map[uint64][]uint64 + clearRowIDs []uint64 + clearColIDs []uint64 + clearExp map[uint64][]uint64 }{ { []uint64{1, 1, 1, 1}, @@ -1336,6 +1505,11 @@ func TestFragment_ImportBool(t *testing.T) { map[uint64][]uint64{ 1: {0, 1, 2, 3}, }, + []uint64{}, + []uint64{}, + map[uint64][]uint64{ + 1: {0, 1, 2, 3}, + }, }, { []uint64{0, 0, 0, 0, 1, 1, 1, 1}, @@ -1344,6 +1518,13 @@ func TestFragment_ImportBool(t *testing.T) { 0: {}, 1: {0, 1, 2, 3}, }, + []uint64{1, 1, 2}, + []uint64{1, 2, 3}, + map[uint64][]uint64{ + 0: {}, + 1: {0, 3}, + 2: {}, + }, }, { []uint64{0, 0, 0, 0, 1}, @@ -1352,6 +1533,12 @@ func TestFragment_ImportBool(t *testing.T) { 0: {0, 2, 3}, 1: {1}, }, + []uint64{1, 1, 1, 1}, + []uint64{0, 1, 2, 3}, + map[uint64][]uint64{ + 0: {0, 2, 3}, + 1: {}, + }, }, { []uint64{1, 1, 1, 1, 0, 0, 1}, @@ -1360,6 +1547,12 @@ func TestFragment_ImportBool(t *testing.T) { 0: {8}, 1: {0, 1, 2, 3}, }, + []uint64{1, 1}, + []uint64{0, 0}, + map[uint64][]uint64{ + 0: {8}, + 1: {1, 2, 3}, + }, }, { []uint64{0, 1, 2}, @@ -1369,6 +1562,13 @@ func TestFragment_ImportBool(t *testing.T) { 1: {}, // This isn't {8} because fragment doesn't validate bool values. 2: {8}, }, + []uint64{1, 2, 3}, + []uint64{9, 9, 9}, + map[uint64][]uint64{ + 0: {}, + 1: {}, + 2: {8}, + }, }, } @@ -1377,13 +1577,28 @@ func TestFragment_ImportBool(t *testing.T) { f := mustOpenBoolFragment("i", "f", viewStandard, 0, "") defer f.Close() - err := f.bulkImport(test.rowIDs, test.colIDs) + // Set import. + err := f.bulkImport(test.setRowIDs, test.setColIDs, &ImportOptions{}) if err != nil { t.Fatalf("bulk importing ids: %v", err) } // Check for expected results. - for k, v := range test.exp { + for k, v := range test.setExp { + cols := f.row(k).Columns() + if !reflect.DeepEqual(cols, v) { + t.Fatalf("expected: %v, but got: %v", v, cols) + } + } + + // Clear import. + err = f.bulkImport(test.clearRowIDs, test.clearColIDs, &ImportOptions{Clear: true}) + if err != nil { + t.Fatalf("bulk importing ids: %v", err) + } + + // Check for expected results. + for k, v := range test.clearExp { cols := f.row(k).Columns() if !reflect.DeepEqual(cols, v) { t.Fatalf("expected: %v, but got: %v", v, cols) @@ -1427,6 +1642,7 @@ func BenchmarkFragment_FullSnapshot(b *testing.B) { rows := make([]uint64, sz) cols := make([]uint64, sz) + options := &ImportOptions{} max := 0 for row := 0; row < 100; row++ { val := 1 @@ -1437,7 +1653,7 @@ func BenchmarkFragment_FullSnapshot(b *testing.B) { val += 2 i++ } - if err := f.bulkImport(rows, cols); err != nil { + if err := f.bulkImport(rows, cols, options); err != nil { b.Fatalf("Error Building Sample: %s", err) } if row > max { @@ -1477,8 +1693,9 @@ func BenchmarkFragment_Import(b *testing.B) { } b.ResetTimer() b.ReportAllocs() + options := &ImportOptions{} for i := 0; i < b.N; i++ { - if err := f.bulkImport(rows, cols); err != nil { + if err := f.bulkImport(rows, cols, options); err != nil { b.Fatalf("Error Building Sample: %s", err) } } @@ -1692,7 +1909,8 @@ func TestFragment_RoaringImportTopN(t *testing.T) { f := mustOpenFragment("i", "f", viewStandard, 0, CacheTypeRanked) defer f.Close() - err := f.bulkImport(test.rowIDs, test.colIDs) + options := &ImportOptions{} + err := f.bulkImport(test.rowIDs, test.colIDs, options) if err != nil { t.Fatalf("bulk importing ids: %v", err) } @@ -1705,7 +1923,7 @@ func TestFragment_RoaringImportTopN(t *testing.T) { t.Fatalf("post bulk import:\n exp: %v\n got: %v\n", expPairs, pairs) } - err = f.bulkImport(test.rowIDs2, test.colIDs2) + err = f.bulkImport(test.rowIDs2, test.colIDs2, options) if err != nil { t.Fatalf("bulk importing ids: %v", err) } diff --git a/http/client.go b/http/client.go index 922ac6234..39f574bb9 100644 --- a/http/client.go +++ b/http/client.go @@ -294,13 +294,22 @@ func (c *InternalClient) QueryNode(ctx context.Context, uri *pilosa.URI, index s } // Import bulk imports bits for a single shard to a host. -func (c *InternalClient) Import(ctx context.Context, index, field string, shard uint64, bits []pilosa.Bit) error { +func (c *InternalClient) Import(ctx context.Context, index, field string, shard uint64, bits []pilosa.Bit, opts ...pilosa.ImportOption) error { if index == "" { return pilosa.ErrIndexRequired } else if field == "" { return pilosa.ErrFieldRequired } + // Set up import options. + options := &pilosa.ImportOptions{} + for _, opt := range opts { + err := opt(options) + if err != nil { + return errors.Wrap(err, "applying option") + } + } + buf, err := c.marshalImportPayload(index, field, shard, bits) if err != nil { return fmt.Errorf("Error Creating Payload: %s", err) @@ -314,7 +323,7 @@ func (c *InternalClient) Import(ctx context.Context, index, field string, shard // Import to each node. for _, node := range nodes { - if err := c.importNode(ctx, node, index, field, buf); err != nil { + if err := c.importNode(ctx, node, index, field, buf, options); err != nil { return fmt.Errorf("import node: host=%s, err=%s", node.URI, err) } } @@ -332,13 +341,22 @@ func getCoordinatorNode(nodes []*pilosa.Node) *pilosa.Node { } // ImportK bulk imports bits specified by string keys to a host. -func (c *InternalClient) ImportK(ctx context.Context, index, field string, bits []pilosa.Bit) error { +func (c *InternalClient) ImportK(ctx context.Context, index, field string, bits []pilosa.Bit, opts ...pilosa.ImportOption) error { if index == "" { return pilosa.ErrIndexRequired } else if field == "" { return pilosa.ErrFieldRequired } + // Set up import options. + options := &pilosa.ImportOptions{} + for _, opt := range opts { + err := opt(options) + if err != nil { + return errors.Wrap(err, "applying option") + } + } + buf, err := c.marshalImportPayload(index, field, 0, bits) if err != nil { return fmt.Errorf("Error Creating Payload: %s", err) @@ -356,7 +374,7 @@ func (c *InternalClient) ImportK(ctx context.Context, index, field string, bits } // Import to node. - if err := c.importNode(ctx, coord, index, field, buf); err != nil { + if err := c.importNode(ctx, coord, index, field, buf, options); err != nil { return fmt.Errorf("import node: host=%s, err=%s", coord.URI, err) } @@ -410,11 +428,17 @@ func (c *InternalClient) marshalImportPayload(index, field string, shard uint64, } // importNode sends a pre-marshaled import request to a node. -func (c *InternalClient) importNode(ctx context.Context, node *pilosa.Node, index, field string, buf []byte) error { +func (c *InternalClient) importNode(ctx context.Context, node *pilosa.Node, index, field string, buf []byte, opts *pilosa.ImportOptions) error { // Create URL & HTTP request. path := fmt.Sprintf("/index/%s/field/%s/import", index, field) u := nodePathToURL(node, path) - req, err := http.NewRequest("POST", u.String(), bytes.NewReader(buf)) + + url := u.String() + if opts.Clear { + url += "?clear=true" + } + + req, err := http.NewRequest("POST", url, bytes.NewReader(buf)) if err != nil { return errors.Wrap(err, "creating request") } @@ -467,9 +491,12 @@ func (c *InternalClient) ImportValue(ctx context.Context, index, field string, s return fmt.Errorf("shard nodes: %s", err) } + // Set up import options. + options := &pilosa.ImportOptions{} + // Import to each node. for _, node := range nodes { - if err := c.importNode(ctx, node, index, field, buf); err != nil { + if err := c.importNode(ctx, node, index, field, buf, options); err != nil { return fmt.Errorf("import node: host=%s, err=%s", node.URI, err) } } @@ -495,8 +522,11 @@ func (c *InternalClient) ImportValueK(ctx context.Context, index, field string, return fmt.Errorf("could not find the coordinator node") } + // Set up import options. + options := &pilosa.ImportOptions{} + // Import to node. - if err := c.importNode(ctx, coord, index, field, buf); err != nil { + if err := c.importNode(ctx, coord, index, field, buf, options); err != nil { return fmt.Errorf("import node: host=%s, err=%s", coord.URI, err) } diff --git a/http/handler.go b/http/handler.go index 8b239ec2d..bfa3d5e4a 100644 --- a/http/handler.go +++ b/http/handler.go @@ -181,8 +181,8 @@ func (h *Handler) populateValidators() { h.validators["DeleteIndex"] = queryValidationSpecRequired() h.validators["PostField"] = queryValidationSpecRequired() h.validators["DeleteField"] = queryValidationSpecRequired() - h.validators["PostImport"] = queryValidationSpecRequired() - h.validators["PostImportRoaring"] = queryValidationSpecRequired().Optional("remote") + h.validators["PostImport"] = queryValidationSpecRequired().Optional("clear") + h.validators["PostImportRoaring"] = queryValidationSpecRequired().Optional("remote", "clear") h.validators["PostQuery"] = queryValidationSpecRequired().Optional("shards", "columnAttrs", "excludeRowAttrs", "excludeColumns") h.validators["GetInfo"] = queryValidationSpecRequired() h.validators["RecalculateCaches"] = queryValidationSpecRequired() @@ -994,6 +994,10 @@ func (h *Handler) handlePostImport(w http.ResponseWriter, r *http.Request) { indexName := mux.Vars(r)["index"] fieldName := mux.Vars(r)["field"] + // If the clear flag is true, treat the import as clear bits. + q := r.URL.Query() + doClear := q.Get("clear") == "true" + // Get index and field type to determine how to handle the // import data. field, err := h.api.Field(r.Context(), indexName, fieldName) @@ -1044,7 +1048,7 @@ func (h *Handler) handlePostImport(w http.ResponseWriter, r *http.Request) { return } - if err := h.api.Import(r.Context(), req); err != nil { + if err := h.api.Import(r.Context(), req, pilosa.OptImportOptionsClear(doClear)); err != nil { switch errors.Cause(err) { case pilosa.ErrClusterDoesNotOwnShard: http.Error(w, err.Error(), http.StatusPreconditionFailed)