mirror of
https://github.com/featurebasedb/featurebase.git
synced 2026-10-06 19:07:50 +00:00
Merge pull request #1784 from molecula/recordbatch-ingestapi
add new RecordBatch implementation which uses ingest API
This commit is contained in:
Matthew Jaffee
GitHub
commit
4488ab1062
4 changed files with 459 additions and 5 deletions
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
145
client/ingest_api_batch.go
Normal file
145
client/ingest_api_batch.go
Normal file
|
|
@ -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 }
|
||||
305
client/ingest_api_batch_test.go
Normal file
305
client/ingest_api_batch_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue