Support keys on import CLI.

This commit is contained in:
Ben Johnson 2018-07-13 12:46:20 +01:00 committed by Travis Turner
parent f278af6194
commit a31a08c330
No known key found for this signature in database
GPG key ID: 7F08008DFD9314C9
7 changed files with 264 additions and 34 deletions

25
api.go
View file

@ -610,11 +610,36 @@ func (api *API) Import(_ context.Context, req *ImportRequest) error {
return errors.Wrap(err, "validating api method")
}
index := api.holder.Index(req.Index)
if index == nil {
return newNotFoundError(ErrIndexNotFound)
}
field, err := api.indexField(req.Index, req.Field, req.Shard)
if err != nil {
return errors.Wrap(err, "getting field")
}
// Translate row keys.
if field.keys() {
if len(req.RowIDs) != 0 {
return errors.New("row ids cannot be used because field uses string keys")
}
if req.RowIDs, err = api.server.translateFile.TranslateRowsToUint64(index.Name(), field.Name(), req.RowKeys); err != nil {
return errors.Wrap(err, "translating rows")
}
}
// Translate column keys.
if index.Keys() {
if len(req.ColumnIDs) != 0 {
return errors.New("column ids cannot be used because index uses string keys")
}
if req.ColumnIDs, err = api.server.translateFile.TranslateColumnsToUint64(index.Name(), req.ColumnKeys); err != nil {
return errors.Wrap(err, "translating columns")
}
}
// Convert timestamps to time.Time.
timestamps := make([]*time.Time, len(req.Timestamps))
for i, ts := range req.Timestamps {

View file

@ -18,8 +18,9 @@ type Bit struct {
// FieldValue represents the value for a column within a
// range-encoded field.
type FieldValue struct {
ColumnID uint64
Value int64
ColumnID uint64
ColumnKey string
Value int64
}
// InternalClient should be implemented by any struct that enables any transport between nodes

View file

@ -51,7 +51,6 @@ omitted. If it is present then its format should be YYYY-MM-DDTHH:MM.
flags.StringVarP(&Importer.Host, "host", "", "localhost:10101", "host:port of Pilosa.")
flags.StringVarP(&Importer.Index, "index", "i", "", "Pilosa index to import into.")
flags.StringVarP(&Importer.Field, "field", "f", "", "Field to import into.")
flags.BoolVar(&Importer.StringKeys, "string-keys", false, "Treat payload as string keys.")
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.")

View file

@ -46,9 +46,6 @@ type ImportCommand struct { // nolint: maligned
// CreateSchema ensures the schema exists before import
CreateSchema bool
// Indicates that the payload should be treated as string keys.
StringKeys bool `json:"StringKeys"`
// Filenames to import from.
Paths []string `json:"paths"`
@ -108,10 +105,14 @@ func (cmd *ImportCommand) Run(ctx context.Context) error {
if err != nil {
return errors.Wrap(err, "getting schema")
}
var useColumnKeys, useRowKeys bool
for _, index := range schema {
if index.Name == cmd.Index {
useColumnKeys = index.Options.Keys
for _, field := range index.Fields {
if field.Name == cmd.Field {
useRowKeys = field.Options.Keys
fieldType = field.Options.Type
}
}
@ -121,7 +122,7 @@ func (cmd *ImportCommand) Run(ctx context.Context) error {
// Import each path and import by shard.
for _, path := range cmd.Paths {
logger.Printf("parsing: %s", path)
if err := cmd.importPath(ctx, fieldType, path); err != nil {
if err := cmd.importPath(ctx, fieldType, useColumnKeys, useRowKeys, path); err != nil {
return err
}
}
@ -142,21 +143,16 @@ func (cmd *ImportCommand) ensureSchema(ctx context.Context) error {
}
// importPath parses a path into bits and imports it to the server.
func (cmd *ImportCommand) importPath(ctx context.Context, fieldType, path string) error {
func (cmd *ImportCommand) importPath(ctx context.Context, fieldType string, useColumnKeys, useRowKeys bool, path string) error {
// If fieldType is `int`, treat the import data as values to be range-encoded.
if fieldType == pilosa.FieldTypeInt {
return cmd.bufferValues(ctx, path)
} else {
if cmd.StringKeys {
return cmd.bufferBitsK(ctx, path)
} else {
return cmd.bufferBits(ctx, path)
}
return cmd.bufferValues(ctx, useColumnKeys, path)
}
return cmd.bufferBits(ctx, useColumnKeys, useRowKeys, path)
}
// bufferBits buffers slices of bits to be imported as a batch.
func (cmd *ImportCommand) bufferBits(ctx context.Context, path string) error {
func (cmd *ImportCommand) bufferBits(ctx context.Context, useColumnKeys, useRowKeys bool, path string) error {
a := make([]pilosa.Bit, 0, cmd.BufferSize)
var r *csv.Reader
@ -198,18 +194,22 @@ func (cmd *ImportCommand) bufferBits(ctx context.Context, path string) error {
var bit pilosa.Bit
// Parse row id.
rowID, err := strconv.ParseUint(record[0], 10, 64)
if err != nil {
return fmt.Errorf("invalid row id on row %d: %q", rnum, record[0])
if useRowKeys {
bit.RowKey = record[0]
} else {
if bit.RowID, err = strconv.ParseUint(record[0], 10, 64); err != nil {
return fmt.Errorf("invalid row id on row %d: %q", rnum, record[0])
}
}
bit.RowID = rowID
// Parse column id.
columnID, err := strconv.ParseUint(record[1], 10, 64)
if err != nil {
return fmt.Errorf("invalid column id on row %d: %q", rnum, record[1])
if useColumnKeys {
bit.ColumnKey = record[1]
} else {
if bit.ColumnID, err = strconv.ParseUint(record[1], 10, 64); err != nil {
return fmt.Errorf("invalid column id on row %d: %q", rnum, record[1])
}
}
bit.ColumnID = columnID
// Parse time, if exists.
if len(record) > 2 && record[2] != "" {
@ -351,7 +351,7 @@ func (cmd *ImportCommand) importBitsK(ctx context.Context, bits []pilosa.Bit) er
}
// bufferValues buffers slices of FieldValues to be imported as a batch.
func (cmd *ImportCommand) bufferValues(ctx context.Context, path string) error {
func (cmd *ImportCommand) bufferValues(ctx context.Context, useColumnKeys bool, path string) error {
a := make([]pilosa.FieldValue, 0, cmd.BufferSize)
var r *csv.Reader
@ -393,11 +393,13 @@ func (cmd *ImportCommand) bufferValues(ctx context.Context, path string) error {
var val pilosa.FieldValue
// Parse column id.
columnID, err := strconv.ParseUint(record[0], 10, 64)
if err != nil {
return fmt.Errorf("invalid column id on row %d: %q", rnum, record[0])
if useColumnKeys {
val.ColumnKey = record[0]
} else {
if val.ColumnID, err = strconv.ParseUint(record[0], 10, 64); err != nil {
return fmt.Errorf("invalid column id on row %d: %q", rnum, record[0])
}
}
val.ColumnID = columnID
// Parse FieldValue.
value, err := strconv.ParseInt(record[1], 10, 64)

View file

@ -336,7 +336,9 @@ func (c *InternalClient) EnsureField(ctx context.Context, indexName string, fiel
func (c *InternalClient) marshalImportPayload(index, field string, shard uint64, bits []pilosa.Bit) ([]byte, error) {
// Separate row and column IDs to reduce allocations.
rowIDs := Bits(bits).RowIDs()
rowKeys := Bits(bits).RowKeys()
columnIDs := Bits(bits).ColumnIDs()
columnKeys := Bits(bits).ColumnKeys()
timestamps := Bits(bits).Timestamps()
// Marshal data to protobuf.
@ -345,7 +347,9 @@ func (c *InternalClient) marshalImportPayload(index, field string, shard uint64,
Field: field,
Shard: shard,
RowIDs: rowIDs,
RowKeys: rowKeys,
ColumnIDs: columnIDs,
ColumnKeys: columnKeys,
Timestamps: timestamps,
})
if err != nil {
@ -447,15 +451,17 @@ func (c *InternalClient) ImportValue(ctx context.Context, index, field string, s
func (c *InternalClient) marshalImportValuePayload(index, field string, shard uint64, vals []pilosa.FieldValue) ([]byte, error) {
// Separate row and column IDs to reduce allocations.
columnIDs := FieldValues(vals).ColumnIDs()
columnKeys := FieldValues(vals).ColumnKeys()
values := FieldValues(vals).Values()
// Marshal data to protobuf.
buf, err := c.serializer.Marshal(&pilosa.ImportValueRequest{
Index: index,
Field: field,
Shard: shard,
ColumnIDs: columnIDs,
Values: values,
Index: index,
Field: field,
Shard: shard,
ColumnIDs: columnIDs,
ColumnKeys: columnKeys,
Values: values,
})
if err != nil {
return nil, fmt.Errorf("marshal import request: %s", err)
@ -858,8 +864,31 @@ func (p Bits) Less(i, j int) bool {
return p[i].RowID < p[j].RowID
}
// HasRowKeys returns true if any values use a row key.
func (p Bits) HasRowKeys() bool {
for i := range p {
if p[i].RowKey != "" {
return true
}
}
return false
}
// HasColumnKeys returns true if any values use a column key.
func (p Bits) HasColumnKeys() bool {
for i := range p {
if p[i].ColumnKey != "" {
return true
}
}
return false
}
// RowIDs returns a slice of all the row IDs.
func (p Bits) RowIDs() []uint64 {
if p.HasRowKeys() {
return nil
}
other := make([]uint64, len(p))
for i := range p {
other[i] = p[i].RowID
@ -869,6 +898,9 @@ func (p Bits) RowIDs() []uint64 {
// ColumnIDs returns a slice of all the column IDs.
func (p Bits) ColumnIDs() []uint64 {
if p.HasColumnKeys() {
return nil
}
other := make([]uint64, len(p))
for i := range p {
other[i] = p[i].ColumnID
@ -878,6 +910,9 @@ func (p Bits) ColumnIDs() []uint64 {
// RowKeys returns a slice of all the row keys.
func (p Bits) RowKeys() []string {
if !p.HasRowKeys() {
return nil
}
other := make([]string, len(p))
for i := range p {
other[i] = p[i].RowKey
@ -887,6 +922,9 @@ func (p Bits) RowKeys() []string {
// ColumnKeys returns a slice of all the column keys.
func (p Bits) ColumnKeys() []string {
if !p.HasColumnKeys() {
return nil
}
other := make([]string, len(p))
for i := range p {
other[i] = p[i].ColumnKey
@ -929,8 +967,21 @@ func (p FieldValues) Less(i, j int) bool {
return p[i].ColumnID < p[j].ColumnID
}
// HasColumnKeys returns true if any values use a column key.
func (p FieldValues) HasColumnKeys() bool {
for i := range p {
if p[i].ColumnKey != "" {
return true
}
}
return false
}
// ColumnIDs returns a slice of all the column IDs.
func (p FieldValues) ColumnIDs() []uint64 {
if p.HasColumnKeys() {
return nil
}
other := make([]uint64, len(p))
for i := range p {
other[i] = p[i].ColumnID
@ -938,6 +989,18 @@ func (p FieldValues) ColumnIDs() []uint64 {
return other
}
// ColumnKeys returns a slice of all the column keys.
func (p FieldValues) ColumnKeys() []string {
if !p.HasColumnKeys() {
return nil
}
other := make([]string, len(p))
for i := range p {
other[i] = p[i].ColumnKey
}
return other
}
// Values returns a slice of all the values.
func (p FieldValues) Values() []int64 {
other := make([]int64, len(p))

View file

@ -202,6 +202,104 @@ func TestClient_Import(t *testing.T) {
}
}
// Ensure client can bulk import data.
func TestClient_ImportKeys(t *testing.T) {
cmd := test.MustRunCluster(t, 1)[0]
host := cmd.URL()
cmd.MustCreateIndex(t, "keyed", pilosa.IndexOptions{Keys: true})
cmd.MustCreateIndex(t, "unkeyed", pilosa.IndexOptions{Keys: false})
cmd.MustCreateField(t, "keyed", "keyedf", pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 1000), pilosa.OptFieldKeys())
cmd.MustCreateField(t, "keyed", "unkeyedf", pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 1000))
cmd.MustCreateField(t, "unkeyed", "keyedf", pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 1000), pilosa.OptFieldKeys())
// Send import request.
c := MustNewClient(host, defaultClient)
t.Run("Import keyed,keyed", func(t *testing.T) {
if err := c.Import(context.Background(), "keyed", "keyedf", 0, []pilosa.Bit{
{RowKey: "green", ColumnKey: "eve"},
{RowKey: "green", ColumnKey: "alice"},
{RowKey: "green", ColumnKey: "bob"},
{RowKey: "blue", ColumnKey: "eve"},
{RowKey: "blue", ColumnKey: "alice"},
{RowKey: "purple", ColumnKey: "eve"},
}); err != nil {
t.Fatal(err)
}
cmd.MustRecalculateCaches(t)
resp := cmd.MustQuery(t, &pilosa.QueryRequest{
Index: "keyed",
Query: "TopN(keyedf)",
})
if pairs, ok := resp.Results[0].([]pilosa.Pair); !ok {
t.Fatal("unexpected response type %T", resp.Results[0])
} else if !reflect.DeepEqual(pairs, []pilosa.Pair{
pilosa.Pair{Key: "green", Count: 3},
pilosa.Pair{Key: "blue", Count: 2},
pilosa.Pair{Key: "purple", Count: 1},
}) {
t.Fatalf("unexpected topn result: %v", pairs)
}
})
t.Run("Import keyed,unkeyedf", func(t *testing.T) {
if err := c.Import(context.Background(), "keyed", "unkeyedf", 0, []pilosa.Bit{
{RowID: 1, ColumnKey: "eve"},
{RowID: 1, ColumnKey: "alice"},
{RowID: 1, ColumnKey: "bob"},
{RowID: 2, ColumnKey: "eve"},
{RowID: 2, ColumnKey: "alice"},
{RowID: 3, ColumnKey: "eve"},
}); err != nil {
t.Fatal(err)
}
cmd.MustRecalculateCaches(t)
resp := cmd.MustQuery(t, &pilosa.QueryRequest{
Index: "keyed",
Query: "TopN(unkeyedf)",
})
if pairs, ok := resp.Results[0].([]pilosa.Pair); !ok {
t.Fatal("unexpected response type %T", resp.Results[0])
} else if !reflect.DeepEqual(pairs, []pilosa.Pair{
pilosa.Pair{ID: 1, Count: 3},
pilosa.Pair{ID: 2, Count: 2},
pilosa.Pair{ID: 3, Count: 1},
}) {
t.Fatalf("unexpected topn result: %v", pairs)
}
})
t.Run("Import unkeyed,keyed", func(t *testing.T) {
if err := c.Import(context.Background(), "unkeyed", "keyedf", 0, []pilosa.Bit{
{RowKey: "green", ColumnID: 1},
{RowKey: "green", ColumnID: 2},
{RowKey: "green", ColumnID: 3},
{RowKey: "blue", ColumnID: 1},
{RowKey: "blue", ColumnID: 2},
{RowKey: "purple", ColumnID: 1},
}); err != nil {
t.Fatal(err)
}
cmd.MustRecalculateCaches(t)
resp := cmd.MustQuery(t, &pilosa.QueryRequest{
Index: "unkeyed",
Query: "TopN(keyedf)",
})
if pairs, ok := resp.Results[0].([]pilosa.Pair); !ok {
t.Fatal("unexpected response type %T", resp.Results[0])
} else if !reflect.DeepEqual(pairs, []pilosa.Pair{
pilosa.Pair{Key: "green", Count: 3},
pilosa.Pair{Key: "blue", Count: 2},
pilosa.Pair{Key: "purple", Count: 1},
}) {
t.Fatalf("unexpected topn result: %v", pairs)
}
})
}
// Ensure client can bulk import value data.
func TestClient_ImportValue(t *testing.T) {
cmd := test.MustRunCluster(t, 1)[0]

View file

@ -16,6 +16,7 @@ package test
import (
"bytes"
"context"
"fmt"
"io"
"io/ioutil"
@ -27,6 +28,7 @@ import (
"testing"
"time"
"github.com/pilosa/pilosa"
"github.com/pilosa/pilosa/http"
"github.com/pilosa/pilosa/server"
"github.com/pkg/errors"
@ -125,6 +127,45 @@ func (m *Command) Reopen() error {
return m.Start()
}
// MustCreateIndex uses this command's API to create an index and fails the test
// if there is an error.
func (m *Command) MustCreateIndex(t *testing.T, name string, opts pilosa.IndexOptions) *pilosa.Index {
idx, err := m.API.CreateIndex(context.Background(), name, opts)
if err != nil {
t.Fatalf("creating index: %v with options: %v, err: %v", name, opts, err)
}
return idx
}
// MustCreateField uses this command's API to create the field. The index must
// already exist - it fails the test if there is an error.
func (m *Command) MustCreateField(t *testing.T, index, field string, opts ...pilosa.FieldOption) *pilosa.Field {
f, err := m.API.CreateField(context.Background(), index, field, opts...)
if err != nil {
t.Fatalf("creating field: %s in index: %s err: %v", field, index, err)
}
return f
}
// MustQuery uses this command's API to execute the given query request, failing
// if Query returns a non-nil error, otherwise returning the QueryResponse.
func (m *Command) MustQuery(t *testing.T, req *pilosa.QueryRequest) pilosa.QueryResponse {
resp, err := m.API.Query(context.Background(), req)
if err != nil {
t.Fatalf("making query: %v, err: %v", req, err)
}
return resp
}
// MustRecalculateCaches calls RecalculateCaches on the command's API, and fails
// if there is an error.
func (m *Command) MustRecalculateCaches(t *testing.T) {
err := m.API.RecalculateCaches(context.Background())
if err != nil {
t.Fatalf("recalcluating caches: %v", err)
}
}
// URL returns the base URL string for accessing the running program.
func (m *Command) URL() string { return m.API.Node().URI.String() }
@ -146,6 +187,7 @@ func (m *Command) Query(index, rawQuery, query string) (string, error) {
return resp.Body, nil
}
// RecalculateCaches is deprecated. Use MustRecalculateCaches.
func (m *Command) RecalculateCaches() error {
resp := MustDo("POST", fmt.Sprintf("%s/recalculate-caches", m.URL()), "")
if resp.StatusCode != 204 {