diff --git a/api.go b/api.go index a81f1db12..df1265433 100644 --- a/api.go +++ b/api.go @@ -256,10 +256,20 @@ func (api *API) Field(_ context.Context, indexName, fieldName string) (*Field, e // (shard*ShardWidth)+(i%ShardWidth). That is to say that "data" represents all // of the rows in this shard of this field concatenated together in one long // bitmap. -func (api *API) ImportRoaring(ctx context.Context, indexName, fieldName string, shard uint64, remote bool, data []byte) (err error) { +func (api *API) ImportRoaring(ctx context.Context, indexName, fieldName string, shard uint64, remote bool, data []byte, opts ...ImportOption) (err error) { if err = api.validate(apiField); 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") + } + } + nodes := api.cluster.shardNodes(indexName, shard) var eg errgroup.Group @@ -280,14 +290,14 @@ func (api *API) ImportRoaring(ctx context.Context, indexName, fieldName string, d2 := make([]byte, len(data)) copy(d2, data) eg.Go(func() error { - return field.importRoaring(d2, shard) + return field.importRoaring(d2, shard, options.Clear) }) go func(node *Node) { }(node) } else if !remote { // if remote == true we don't forward to other nodes // forward it on eg.Go(func() error { - return api.server.defaultClient.ImportRoaring(ctx, &node.URI, indexName, fieldName, shard, true, data) + return api.server.defaultClient.ImportRoaring(ctx, &node.URI, indexName, fieldName, shard, true, data, opts...) }) } } diff --git a/client.go b/client.go index e6a02aa96..f6c49d14b 100644 --- a/client.go +++ b/client.go @@ -53,7 +53,7 @@ type InternalClient interface { RowAttrDiff(ctx context.Context, uri *URI, index, field string, blks []AttrBlock) (map[uint64]map[string]interface{}, error) SendMessage(ctx context.Context, uri *URI, msg []byte) error RetrieveShardFromURI(ctx context.Context, index, field string, shard uint64, uri URI) (io.ReadCloser, error) - ImportRoaring(ctx context.Context, uri *URI, index, field string, shard uint64, remote bool, data []byte) error + ImportRoaring(ctx context.Context, uri *URI, index, field string, shard uint64, remote bool, data []byte, opts ...ImportOption) error } //=============== @@ -109,7 +109,7 @@ func (n nopInternalClient) Import(ctx context.Context, index, field string, shar 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 { +func (n nopInternalClient) ImportRoaring(ctx context.Context, uri *URI, index, field string, shard uint64, remote bool, data []byte, opts ...ImportOption) error { return nil } func (n nopInternalClient) EnsureIndex(ctx context.Context, name string, options IndexOptions) error { diff --git a/field.go b/field.go index 3422e0e19..bc656fa54 100644 --- a/field.go +++ b/field.go @@ -1182,7 +1182,7 @@ func (f *Field) importValue(columnIDs []uint64, values []int64, options *ImportO return nil } -func (f *Field) importRoaring(data []byte, shard uint64) error { +func (f *Field) importRoaring(data []byte, shard uint64, clear bool) error { viewName := viewStandard view, err := f.createViewIfNotExists(viewName) @@ -1195,7 +1195,7 @@ func (f *Field) importRoaring(data []byte, shard uint64) error { return errors.Wrap(err, "creating fragment") } - if err := frag.importRoaring(data); err != nil { + if err := frag.importRoaring(data, clear); err != nil { return err } diff --git a/fragment.go b/fragment.go index d45757762..10e5e3920 100644 --- a/fragment.go +++ b/fragment.go @@ -1651,7 +1651,7 @@ func (f *fragment) importValue(columnIDs, values []uint64, bitDepth uint, clear // importRoaring imports from the official roaring data format defined at // https://github.com/RoaringBitmap/RoaringFormatSpec or from pilosa's version // of the roaring format. The cache is updated to reflect the new data. -func (f *fragment) importRoaring(data []byte) error { +func (f *fragment) importRoaring(data []byte, clear bool) error { f.mu.Lock() defer f.mu.Unlock() bm := roaring.NewBitmap() @@ -1679,8 +1679,12 @@ func (f *fragment) importRoaring(data []byte) error { lastRow = vRow } - if f.storage.Count() > 0 { - bm = f.storage.Union(bm) + if clear { + bm = f.storage.Difference(bm) + } else { + if f.storage.Count() > 0 { + bm = f.storage.Union(bm) + } } for _, rowID := range rowSet { diff --git a/fragment_internal_test.go b/fragment_internal_test.go index 5097fb680..ac080b5b7 100644 --- a/fragment_internal_test.go +++ b/fragment_internal_test.go @@ -1900,7 +1900,7 @@ func TestFragment_RoaringImport(t *testing.T) { if err != nil { t.Fatalf("writing to buffer: %v", err) } - f.importRoaring(buf.Bytes()) + f.importRoaring(buf.Bytes(), false) exp := calcExpected(test[:num+1]...) for row, expCols := range exp { cols := f.row(uint64(row)).Columns() @@ -1972,7 +1972,7 @@ func TestFragment_RoaringImportTopN(t *testing.T) { if err != nil { t.Fatalf("writing to buffer: %v", err) } - f.importRoaring(buf.Bytes()) + f.importRoaring(buf.Bytes(), false) rows, cols := toRowsCols(test.roaring) expPairs = calcTop(append(test.rowIDs, rows...), append(test.colIDs, cols...)) pairs, err = f.top(topOptions{}) diff --git a/http/client.go b/http/client.go index d546c5b9c..17c0b2d63 100644 --- a/http/client.go +++ b/http/client.go @@ -569,7 +569,7 @@ func (c *InternalClient) marshalImportValuePayload(index, field string, shard ui // ImportRoaring does fast import of raw bits in roaring format (pilosa or // official format, see API.ImportRoaring). -func (c *InternalClient) ImportRoaring(ctx context.Context, uri *pilosa.URI, index, field string, shard uint64, remote bool, data []byte) error { +func (c *InternalClient) ImportRoaring(ctx context.Context, uri *pilosa.URI, index, field string, shard uint64, remote bool, data []byte, opts ...pilosa.ImportOption) error { if index == "" { return pilosa.ErrIndexRequired } else if field == "" { @@ -579,7 +579,19 @@ func (c *InternalClient) ImportRoaring(ctx context.Context, uri *pilosa.URI, ind uri = c.defaultURI } + // Set up import options. + options := &pilosa.ImportOptions{} + for _, opt := range opts { + err := opt(options) + if err != nil { + return errors.Wrap(err, "applying option") + } + } + url := fmt.Sprintf("%s/index/%s/field/%s/import-roaring/%d?remote=%v", uri, index, field, shard, remote) + if options.Clear { + url += "&clear=true" + } // Generate HTTP request. req, err := http.NewRequest("POST", url, bytes.NewBuffer(data))