diff --git a/client/batch.go b/client/batch.go index 0b2ae36bb..c71b3c5de 100644 --- a/client/batch.go +++ b/client/batch.go @@ -342,6 +342,10 @@ func (qt *QuantizedTime) SetHour(hour string) { copy(qt.ymdh[8:10], hour) } +func (qt *QuantizedTime) Time() (time.Time, error) { + return time.Parse("2006010215", string(qt.ymdh[:])) +} + // Reset sets the time to the zero value which generates no time views. func (qt *QuantizedTime) Reset() { for i := range qt.ymdh { diff --git a/client/client.go b/client/client.go index 4076069c2..d27ddcfc0 100644 --- a/client/client.go +++ b/client/client.go @@ -782,7 +782,7 @@ func (c *Client) readSchema() ([]SchemaIndex, error) { func (c *Client) IngestSchema(reqBody map[string]interface{}) (body []byte, err error) { data, err := json.Marshal(reqBody) if err != nil { - return data, errors.Wrap(err, " error building Schema body to Ingest") + return data, errors.Wrap(err, "error building Schema body to Ingest") } return c.IngestRequest("/internal/schema", data) } @@ -790,19 +790,19 @@ func (c *Client) IngestSchema(reqBody map[string]interface{}) (body []byte, err func (c *Client) IngestData(index string, reqBody []map[string]interface{}) (body []byte, err error) { data, err := json.Marshal(reqBody) if err != nil { - return data, errors.Wrap(err, " error building request body to Ingest") + return data, errors.Wrap(err, "error building request body to Ingest") } return c.IngestRequest("/internal/ingest/"+index, data) } -func (c *Client) IngestRequest(Uri string, data []byte) (body []byte, err error) { +func (c *Client) IngestRequest(uri string, data []byte) (body []byte, err error) { var header = make(map[string]string) header["Content-Type"] = "application/json" header["Accept"] = "application/json" header["User-Agent"] = "pilosa/" + pilosa.Version - _, body, err = c.HTTPRequest("POST", Uri, data, header) + status, body, err := c.HTTPRequest("POST", uri, data, header) if err != nil { - return nil, errors.Wrap(err, "requesting "+Uri) + return nil, errors.Wrapf(err, "requesting %s status: %d", uri, status) } return body, err } diff --git a/client/ingest_api_batch.go b/client/ingest_api_batch.go new file mode 100644 index 000000000..2f6d35d5e --- /dev/null +++ b/client/ingest_api_batch.go @@ -0,0 +1,145 @@ +package client + +import ( + "time" + + "github.com/molecula/featurebase/v2/logger" + "github.com/pkg/errors" +) + +// NewIngestAPIBatch creates an alternate implementation of +// RecordBatch which exists to aid in testing the new Ingest API and +// is likely far slower than the Batch. +func NewIngestAPIBatch(client *Client, size int, logger logger.Logger, fields []*Field) *ingestAPIBatch { + if len(fields) == 0 { + return nil + } + + return &ingestAPIBatch{ + client: client, + log: logger, + fields: fields, + keyed: fields[0].index.Opts().Keys(), + index: fields[0].index.Name(), + batchSize: size, + + recordsK: make(map[string]map[string]interface{}), + records: make(map[uint64]map[string]interface{}), + } +} + +type ingestAPIBatch struct { + client *Client + log logger.Logger + batchSize int + + fields []*Field + keyed bool + index string + + // map[recordKey][fieldName]value + recordsK map[string]map[string]interface{} + records map[uint64]map[string]interface{} +} + +func (b *ingestAPIBatch) Add(row Row) error { + if len(row.Clears) > 0 { + return errors.New("ingest api batch does not support clears") + } + values := make(map[string]interface{}) + for i, val := range row.Values { + field := b.fields[i] + // val can be string, uint64, int64, []string, []uint64, nil + // TODO timestamp field might need special handling + // TODO check that the Row.Clears field is only used for packed bools, and then issue a warning/error (in IDK) if the ingest API mode is used in conjunction w/ packed bools. + if val == nil { + continue + } + zero := QuantizedTime{} + if field.Options().Type() == FieldTypeTime && row.Time != zero { + timeq, err := row.Time.Time() + if err != nil { + return errors.Wrap(err, "parsing row time") + } + values[field.Name()] = map[string]interface{}{"time": timeq.Format(time.RFC3339), "values": val} + } else { + values[field.Name()] = val + } + } + + if b.keyed { + switch rowID := row.ID.(type) { + case string: + b.recordsK[rowID] = values + case []byte: + b.recordsK[string(rowID)] = values + default: + return errors.Errorf("unsupported rowID %v of type %[1]T, must be string, or []byte for keyed index", rowID) + } + if len(b.recordsK) >= b.batchSize { + return ErrBatchNowFull + } + } else { + rowID, ok := row.ID.(uint64) + if !ok { + return errors.Errorf("unsupported rowID %v of type %[1]T, must be uint64 for unkeyed index", row.ID) + } + b.records[rowID] = values + if len(b.records) >= b.batchSize { + return ErrBatchNowFull + } + } + return nil +} + +func (b *ingestAPIBatch) Import() error { + if b.keyed { + return b.importKeyed() + } + return b.importUnkeyed() +} + +func (b *ingestAPIBatch) importKeyed() error { + req := []map[string]interface{}{ + { + "action": "set", + "records": b.recordsK, + }, + } + bod, err := b.client.IngestData(b.index, req) + if err != nil { + return errors.Wrapf(err, "importKeyed, body: %s", bod) + } + + for k := range b.recordsK { + delete(b.recordsK, k) + } + return nil +} + +func (b *ingestAPIBatch) importUnkeyed() error { + req := []map[string]interface{}{ + { + "action": "set", + "records": b.records, + }, + } + bod, err := b.client.IngestData(b.index, req) + if err != nil { + return errors.Wrapf(err, "importKeyed, body: %s", bod) + } + + for v := range b.records { + delete(b.records, v) + } + return nil +} + +func (b *ingestAPIBatch) Len() int { + if b.keyed { + return len(b.recordsK) + } + return len(b.records) +} + +func (b *ingestAPIBatch) Flush() error { return nil } diff --git a/client/ingest_api_batch_test.go b/client/ingest_api_batch_test.go new file mode 100644 index 000000000..df94ae33d --- /dev/null +++ b/client/ingest_api_batch_test.go @@ -0,0 +1,305 @@ +package client + +import ( + "strings" + "testing" + "time" + + "github.com/molecula/featurebase/v2/logger" + "github.com/molecula/featurebase/v2/test" +) + +func TestIngestAPIBatchAdd(t *testing.T) { + t.Run("unkeyed", func(t *testing.T) { + batch := NewIngestAPIBatch(nil, 10, logger.NopLogger, []*Field{ + { + name: "a", + index: &Index{name: "idxname", options: &IndexOptions{}}, + options: &FieldOptions{ + fieldType: FieldTypeSet, + }, + }, + { + name: "b", + index: &Index{name: "idxname", options: &IndexOptions{}}, + options: &FieldOptions{ + fieldType: FieldTypeSet, + keys: true, + }, + }, + { + name: "c", + index: &Index{name: "idxname", options: &IndexOptions{}}, + options: &FieldOptions{ + fieldType: FieldTypeTime, + keys: true, + }, + }, + }) + qt := QuantizedTime{} + qt.Set(time.Date(2007, time.January, 1, 15, 0, 0, 0, time.UTC)) + err := batch.Add(Row{ + ID: uint64(1), + Values: []interface{}{uint64(2), "bkey", "ckey"}, + Time: qt, + }) + if err != nil { + t.Fatalf("adding row to batch: %v", err) + } + + if batch.records[1]["a"] != uint64(2) { + t.Fatalf("unexpected batch.records: %+v", batch.records) + } + if batch.records[1]["b"] != "bkey" { + t.Fatalf("unexpected batch.records: %+v", batch.records) + } + if batch.records[1]["c"].(map[string]interface{})["time"] != "2007-01-01T15:00:00Z" { + t.Fatalf("unexpected batch.records: %+v", batch.records) + } + if batch.records[1]["c"].(map[string]interface{})["values"] != "ckey" { + t.Fatalf("unexpected batch.records: %+v", batch.records) + } + + }) + + t.Run("keyed", func(t *testing.T) { + batch := NewIngestAPIBatch(nil, 10, logger.NopLogger, []*Field{ + { + name: "a", + index: &Index{name: "idxname", options: &IndexOptions{keys: true}}, + options: &FieldOptions{ + fieldType: FieldTypeSet, + }, + }, + { + name: "b", + index: &Index{name: "idxname", options: &IndexOptions{keys: true}}, + options: &FieldOptions{ + fieldType: FieldTypeSet, + keys: true, + }, + }, + { + name: "c", + index: &Index{name: "idxname", options: &IndexOptions{keys: true}}, + options: &FieldOptions{ + fieldType: FieldTypeTime, + keys: true, + }, + }, + }) + qt := QuantizedTime{} + qt.Set(time.Date(2007, time.January, 1, 15, 0, 0, 0, time.UTC)) + err := batch.Add(Row{ + ID: "1", + Values: []interface{}{uint64(2), "bkey", "ckey"}, + Time: qt, + }) + + checkResult := func(batch *ingestAPIBatch, id string, err error) { + if err != nil { + t.Fatalf("adding row to batch: %v", err) + } + + if batch.recordsK[id]["a"] != uint64(2) { + t.Fatalf("unexpected batch.records: %+v", batch.recordsK) + } + if batch.recordsK[id]["b"] != "bkey" { + t.Fatalf("unexpected batch.records: %+v", batch.recordsK) + } + if batch.recordsK[id]["c"].(map[string]interface{})["time"] != "2007-01-01T15:00:00Z" { + t.Fatalf("unexpected batch.records: %+v", batch.recordsK) + } + if batch.recordsK[id]["c"].(map[string]interface{})["values"] != "ckey" { + t.Fatalf("unexpected batch.records: %+v", batch.recordsK) + } + } + checkResult(batch, "1", err) + + // test wrong type row ID + if err := batch.Add(Row{ID: 64.5}); !strings.Contains(err.Error(), "unsupported rowID") { + t.Fatalf("unexpected error w/ floating point rowID: %v", err) + } + + // test that byte slice ID works same as string + err = batch.Add(Row{ + ID: []byte("2"), + Values: []interface{}{uint64(2), "bkey", "ckey"}, + Time: qt, + }) + checkResult(batch, "2", err) + + }) +} + +func TestIngestAPIBatch(t *testing.T) { + c := test.MustRunCluster(t, 3) + defer c.Close() + + urls := make([]string, len(c.Nodes)) + for i, n := range c.Nodes { + urls[i] = n.URL() + } + + // Create a new client for the cluster + cli, err := newClientFromAddresses(urls, &ClientOptions{}) + if err != nil { + t.Fatalf("getting new client: %v", err) + } + defer cli.Close() + + cli.IngestSchema(map[string]interface{}{ + "index-name": "test-1", + "index-action": "create", + "primary-key-type": "uint", + "field-action": "create", + "fields": []map[string]interface{}{ + { + "field-name": "astr", + "field-type": "string", + "field-options": map[string]interface{}{}, + }, + { + "field-name": "bint", + "field-type": "int", + "field-options": map[string]interface{}{}, + }, + { + "field-name": "cid", + "field-type": "id", + "field-options": map[string]interface{}{}, + }, + { + "field-name": "dtimestamp", + "field-type": "timestamp", + "field-options": map[string]interface{}{ + "unit": "s", + }, + }, + { + "field-name": "etime", + "field-type": "string", + "field-options": map[string]interface{}{ + "time-quantum": "YMD", + }, + }, + { + "field-name": "fdecimal", + "field-type": "decimal", + "field-options": map[string]interface{}{ + "scale": 3, + }, + }, + { + "field-name": "gbool", + "field-type": "bool", + "field-options": map[string]interface{}{}, + }, + }, + }) + + schema, err := cli.Schema() + if err != nil { + t.Fatalf("getting schema: %v", err) + } + index := schema.Index("test-1") + defer cli.DeleteIndex(index) + + batch := NewIngestAPIBatch(cli, 10, logger.NopLogger, []*Field{ + { + name: "astr", + index: &Index{name: "test-1", options: &IndexOptions{}}, + options: &FieldOptions{fieldType: FieldTypeSet, keys: true}, + }, + { + name: "bint", + options: &FieldOptions{fieldType: FieldTypeInt}, + }, + { + name: "cid", + options: &FieldOptions{fieldType: FieldTypeSet, keys: false}, + }, + { + name: "dtimestamp", + options: &FieldOptions{fieldType: FieldTypeTimestamp}, + }, + { + name: "etime", + options: &FieldOptions{fieldType: FieldTypeTime, keys: true, timeQuantum: TimeQuantumYearMonthDay}, + }, + { + name: "fdecimal", + options: &FieldOptions{fieldType: FieldTypeDecimal, scale: 3}, + }, + { + name: "gbool", + options: &FieldOptions{fieldType: FieldTypeBool}, + }, + }) + + qt0 := &QuantizedTime{} + qt0.Set(time.Date(2010, time.January, 1, 0, 0, 0, 0, time.UTC)) + if err := batch.Add(Row{ + ID: uint64(7), + Values: []interface{}{"a", -2, 9, 1287367623, "e", 1.2345, true}, + Time: *qt0, + }); err != nil { + t.Fatalf("adding row: %v", err) + } + + // test nil value case + if err := batch.Add(Row{ + ID: uint64(8), + Values: []interface{}{nil, nil, nil, nil, nil, nil, nil}, + Time: QuantizedTime{}, + }); err != nil { + t.Fatalf("error adding all nil batch which should affect nothing: %v", err) + } + + if err := batch.Import(); err != nil { + t.Fatalf("importing row: %v", err) + } + + if resp, err := cli.Query(NewPQLBaseQuery("Row(astr=a)", &Index{name: "test-1", options: &IndexOptions{}}, nil)); err != nil { + t.Fatalf("querying: %v", err) + } else if len(resp.Result().Row().Columns) != 1 || resp.Result().Row().Columns[0] != uint64(7) { + t.Fatalf("unexpected Row(asr=a) result: %+v", resp.Result().Row().Columns) + } + + if resp, err := cli.Query(NewPQLBaseQuery("Row(bint==-2)", &Index{name: "test-1", options: &IndexOptions{}}, nil)); err != nil { + t.Fatalf("querying: %v", err) + } else if len(resp.Result().Row().Columns) != 1 || resp.Result().Row().Columns[0] != uint64(7) { + t.Fatalf("unexpected Row(asr=a) result: %+v", resp.Result().Row().Columns) + } + + if resp, err := cli.Query(NewPQLBaseQuery("Row(cid=9)", &Index{name: "test-1", options: &IndexOptions{}}, nil)); err != nil { + t.Fatalf("querying: %v", err) + } else if len(resp.Result().Row().Columns) != 1 || resp.Result().Row().Columns[0] != uint64(7) { + t.Fatalf("unexpected Row(asr=a) result: %+v", resp.Result().Row().Columns) + } + + if resp, err := cli.Query(NewPQLBaseQuery("Row(dtimestamp=='2010-10-18T02:07:03Z')", &Index{name: "test-1", options: &IndexOptions{}}, nil)); err != nil { + t.Fatalf("querying: %v", err) + } else if len(resp.Result().Row().Columns) != 1 || resp.Result().Row().Columns[0] != uint64(7) { + t.Fatalf("unexpected Row(asr=a) result: %+v", resp.Result().Row().Columns) + } + + if resp, err := cli.Query(NewPQLBaseQuery("Row(etime=e, from='2010-01-01', to='2010-01-02')", &Index{name: "test-1", options: &IndexOptions{}}, nil)); err != nil { + t.Fatalf("querying: %v", err) + } else if len(resp.Result().Row().Columns) != 1 || resp.Result().Row().Columns[0] != uint64(7) { + t.Fatalf("unexpected Row(asr=a) result: %+v", resp.Result().Row().Columns) + } + + if resp, err := cli.Query(NewPQLBaseQuery("Row(fdecimal==1.234)", &Index{name: "test-1", options: &IndexOptions{}}, nil)); err != nil { + t.Fatalf("querying: %v", err) + } else if len(resp.Result().Row().Columns) != 1 || resp.Result().Row().Columns[0] != uint64(7) { + t.Fatalf("unexpected Row(asr=a) result: %+v", resp.Result().Row().Columns) + } + + if resp, err := cli.Query(NewPQLBaseQuery("Row(gbool=true)", &Index{name: "test-1", options: &IndexOptions{}}, nil)); err != nil { + t.Fatalf("querying: %v", err) + } else if len(resp.Result().Row().Columns) != 1 || resp.Result().Row().Columns[0] != uint64(7) { + t.Fatalf("unexpected Row(asr=a) result: %+v", resp.Result().Row().Columns) + } + +}