mirror of
https://github.com/featurebasedb/featurebase.git
synced 2026-10-06 19:07:50 +00:00
Merge pull request #1599 from travisturner/import-key-cli
Support keys on import CLI
This commit is contained in:
commit
227e3bfc9f
9 changed files with 464 additions and 153 deletions
25
api.go
25
api.go
|
|
@ -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.holder.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.holder.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 {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -33,6 +34,7 @@ type InternalClient interface {
|
|||
Schema(ctx context.Context) ([]*IndexInfo, error)
|
||||
CreateIndex(ctx context.Context, index string, opt IndexOptions) error
|
||||
FragmentNodes(ctx context.Context, index string, shard uint64) ([]*Node, error)
|
||||
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
|
||||
|
|
@ -88,6 +90,9 @@ func (n nopInternalClient) CreateIndex(ctx context.Context, index string, opt In
|
|||
func (n nopInternalClient) FragmentNodes(ctx context.Context, index string, shard uint64) ([]*Node, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (n nopInternalClient) Nodes(ctx context.Context) ([]*Node, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (n nopInternalClient) Query(ctx context.Context, index string, queryRequest *QueryRequest) (*QueryResponse, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -51,7 +51,7 @@ 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.BoolVar(&Importer.StringKeys, "string-keys", false, "REMOVED (key type is now determined by index/field configuration): 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.")
|
||||
|
|
|
|||
168
ctl/import.go
168
ctl/import.go
|
|
@ -46,7 +46,8 @@ type ImportCommand struct { // nolint: maligned
|
|||
// CreateSchema ensures the schema exists before import
|
||||
CreateSchema bool
|
||||
|
||||
// Indicates that the payload should be treated as string keys.
|
||||
// REMOVED: Indicates that the payload should be treated as string keys.
|
||||
// TODO: remove this in a future release
|
||||
StringKeys bool `json:"StringKeys"`
|
||||
|
||||
// Filenames to import from.
|
||||
|
|
@ -79,6 +80,11 @@ func NewImportCommand(stdin io.Reader, stdout, stderr io.Writer) *ImportCommand
|
|||
func (cmd *ImportCommand) Run(ctx context.Context) error {
|
||||
logger := log.New(cmd.Stderr, "", log.LstdFlags)
|
||||
|
||||
// REMOVED: warning that --string-keys flag has been deprecated.
|
||||
if cmd.StringKeys {
|
||||
logger.Printf("REMOVED: The string-keys flag is no longer used.")
|
||||
}
|
||||
|
||||
// Validate arguments.
|
||||
// Index and field are validated early before the files are parsed.
|
||||
if cmd.Index == "" {
|
||||
|
|
@ -108,20 +114,26 @@ 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
|
||||
break
|
||||
}
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// 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 +154,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 +205,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] != "" {
|
||||
|
|
@ -224,7 +235,7 @@ func (cmd *ImportCommand) bufferBits(ctx context.Context, path string) error {
|
|||
|
||||
// If we've reached the buffer size then import bits.
|
||||
if len(a) == cmd.BufferSize {
|
||||
if err := cmd.importBits(ctx, a); err != nil {
|
||||
if err := cmd.importBits(ctx, useColumnKeys, useRowKeys, a); err != nil {
|
||||
return err
|
||||
}
|
||||
a = a[:0]
|
||||
|
|
@ -232,13 +243,22 @@ func (cmd *ImportCommand) bufferBits(ctx context.Context, path string) error {
|
|||
}
|
||||
|
||||
// If there are still bits in the buffer then flush them.
|
||||
return cmd.importBits(ctx, a)
|
||||
return cmd.importBits(ctx, useColumnKeys, useRowKeys, a)
|
||||
}
|
||||
|
||||
// importBits sends batches of bits to the server.
|
||||
func (cmd *ImportCommand) importBits(ctx context.Context, bits []pilosa.Bit) error {
|
||||
func (cmd *ImportCommand) importBits(ctx context.Context, useColumnKeys, useRowKeys bool, bits []pilosa.Bit) error {
|
||||
logger := log.New(cmd.Stderr, "", log.LstdFlags)
|
||||
|
||||
// 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 {
|
||||
return errors.Wrap(err, "importing keys")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Group bits by shard.
|
||||
logger.Printf("grouping %d bits", len(bits))
|
||||
bitsByShard := http.Bits(bits).GroupByShard()
|
||||
|
|
@ -258,100 +278,8 @@ func (cmd *ImportCommand) importBits(ctx context.Context, bits []pilosa.Bit) err
|
|||
return nil
|
||||
}
|
||||
|
||||
// bufferBitsK buffers slices of keys to be imported as a batch.
|
||||
func (cmd *ImportCommand) bufferBitsK(ctx context.Context, path string) error {
|
||||
a := make([]pilosa.Bit, 0, cmd.BufferSize)
|
||||
|
||||
var r *csv.Reader
|
||||
|
||||
if path != "-" {
|
||||
// Open file for reading.
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "opening file")
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
// Read rows as bits.
|
||||
r = csv.NewReader(f)
|
||||
} else {
|
||||
r = csv.NewReader(cmd.Stdin)
|
||||
}
|
||||
|
||||
r.FieldsPerRecord = -1
|
||||
rnum := 0
|
||||
for {
|
||||
rnum++
|
||||
|
||||
// Read CSV row.
|
||||
record, err := r.Read()
|
||||
if err == io.EOF {
|
||||
break
|
||||
} else if err != nil {
|
||||
return errors.Wrap(err, "reading")
|
||||
}
|
||||
|
||||
// Ignore blank rows.
|
||||
if record[0] == "" {
|
||||
continue
|
||||
} else if len(record) < 2 {
|
||||
return fmt.Errorf("bad column count on row %d: col=%d", rnum, len(record))
|
||||
}
|
||||
|
||||
var bit pilosa.Bit
|
||||
|
||||
// Parse row key.
|
||||
if record[0] == "" {
|
||||
return fmt.Errorf("invalid row key on row %d: %q", rnum, record[0])
|
||||
}
|
||||
bit.RowKey = record[0]
|
||||
|
||||
// Parse column key.
|
||||
if record[1] == "" {
|
||||
return fmt.Errorf("invalid column id on row %d: %q", rnum, record[1])
|
||||
}
|
||||
bit.ColumnKey = record[1]
|
||||
|
||||
// Parse time, if exists.
|
||||
if len(record) > 2 && record[2] != "" {
|
||||
t, err := time.Parse(pilosa.TimeFormat, record[2])
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid timestamp on row %d: %q", rnum, record[2])
|
||||
}
|
||||
bit.Timestamp = t.UnixNano()
|
||||
}
|
||||
|
||||
a = append(a, bit)
|
||||
|
||||
// If we've reached the buffer size then import bits.
|
||||
if len(a) == cmd.BufferSize {
|
||||
if err := cmd.importBitsK(ctx, a); err != nil {
|
||||
return err
|
||||
}
|
||||
a = a[:0]
|
||||
}
|
||||
}
|
||||
|
||||
// If there are still bitKs in the buffer then flush them.
|
||||
return cmd.importBitsK(ctx, a)
|
||||
}
|
||||
|
||||
// importBitsK sends batches of bitKs to the server.
|
||||
func (cmd *ImportCommand) importBitsK(ctx context.Context, bits []pilosa.Bit) error {
|
||||
logger := log.New(cmd.Stderr, "", log.LstdFlags)
|
||||
|
||||
// TODO: does it help to sort the rowKeys?
|
||||
|
||||
logger.Printf("importing keys: n=%d", len(bits))
|
||||
if err := cmd.client.ImportK(ctx, cmd.Index, cmd.Field, bits); err != nil {
|
||||
return errors.Wrap(err, "importing keys")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// 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 +321,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)
|
||||
|
|
|
|||
|
|
@ -101,6 +101,33 @@ func TestImportCommand_RunValue(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
// Ensure that import with keys runs.
|
||||
func TestImportCommand_RunKeys(t *testing.T) {
|
||||
buf := bytes.Buffer{}
|
||||
stdin, stdout, stderr := GetIO(buf)
|
||||
cm := NewImportCommand(stdin, stdout, stderr)
|
||||
file, err := ioutil.TempFile("", "import-key.csv")
|
||||
file.Write([]byte("foo1,bar2\nfoo3,bar4\nfoo5,bar6"))
|
||||
ctx := context.Background()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
cmd := test.MustRunCluster(t, 1)[0]
|
||||
cm.Host = cmd.API.Node().URI.HostPort()
|
||||
|
||||
http.DefaultClient.Do(MustNewHTTPRequest("POST", "http://"+cm.Host+"/index/i", strings.NewReader(`{"options":{"keys": true}}`)))
|
||||
http.DefaultClient.Do(MustNewHTTPRequest("POST", "http://"+cm.Host+"/index/i/field/f", strings.NewReader(`{"options":{"keys": true}}`)))
|
||||
|
||||
cm.Index = "i"
|
||||
cm.Field = "f"
|
||||
cm.Paths = []string{file.Name()}
|
||||
err = cm.Run(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("Import Run with keys doesn't work: %s", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestImportCommand_InvalidFile(t *testing.T) {
|
||||
cmd := test.MustRunCluster(t, 1)[0]
|
||||
|
||||
|
|
|
|||
151
http/client.go
151
http/client.go
|
|
@ -207,6 +207,37 @@ func (c *InternalClient) FragmentNodes(ctx context.Context, index string, shard
|
|||
return a, nil
|
||||
}
|
||||
|
||||
// Nodes returns a list of all nodes.
|
||||
func (c *InternalClient) Nodes(ctx context.Context) ([]*pilosa.Node, error) {
|
||||
// Execute request against the host.
|
||||
u := uriPathToURL(c.defaultURI, "/internal/nodes")
|
||||
|
||||
// Build request.
|
||||
req, err := http.NewRequest("GET", u.String(), nil)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "creating request")
|
||||
}
|
||||
|
||||
req.Header.Set("User-Agent", "pilosa/"+pilosa.Version)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
// Execute request.
|
||||
resp, err := c.httpClient.Do(req.WithContext(ctx))
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "executing request")
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
var a []*pilosa.Node
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("http: status=%d", resp.StatusCode)
|
||||
} else if err := json.NewDecoder(resp.Body).Decode(&a); err != nil {
|
||||
return nil, fmt.Errorf("json decode: %s", err)
|
||||
}
|
||||
|
||||
return a, nil
|
||||
}
|
||||
|
||||
// Query executes query against the index.
|
||||
func (c *InternalClient) Query(ctx context.Context, index string, queryRequest *pilosa.QueryRequest) (*pilosa.QueryResponse, error) {
|
||||
return c.QueryNode(ctx, c.defaultURI, index, queryRequest)
|
||||
|
|
@ -291,26 +322,42 @@ func (c *InternalClient) Import(ctx context.Context, index, field string, shard
|
|||
return nil
|
||||
}
|
||||
|
||||
func getCoordinatorNode(nodes []*pilosa.Node) *pilosa.Node {
|
||||
for _, node := range nodes {
|
||||
if node.IsCoordinator {
|
||||
return node
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ImportK bulk imports bits specified by string keys to a host.
|
||||
func (c *InternalClient) ImportK(ctx context.Context, index, field string, columns []pilosa.Bit) error {
|
||||
func (c *InternalClient) ImportK(ctx context.Context, index, field string, bits []pilosa.Bit) error {
|
||||
if index == "" {
|
||||
return pilosa.ErrIndexRequired
|
||||
} else if field == "" {
|
||||
return pilosa.ErrFieldRequired
|
||||
}
|
||||
|
||||
buf, err := c.marshalImportPayloadK(index, field, columns)
|
||||
buf, err := c.marshalImportPayload(index, field, 0, bits)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Error Creating Payload: %s", err)
|
||||
}
|
||||
|
||||
node := &pilosa.Node{
|
||||
URI: *c.defaultURI,
|
||||
// Get the coordinator node; all bits are sent to the
|
||||
// primary translate store (i.e. coordinator).
|
||||
nodes, err := c.Nodes(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("getting nodes: %s", err)
|
||||
}
|
||||
coord := getCoordinatorNode(nodes)
|
||||
if coord == nil {
|
||||
return fmt.Errorf("could not find the coordinator node")
|
||||
}
|
||||
|
||||
// Import to node.
|
||||
if err := c.importNode(ctx, node, index, field, buf); err != nil {
|
||||
return fmt.Errorf("import node: host=%s, err=%s", node.URI, err)
|
||||
if err := c.importNode(ctx, coord, index, field, buf); err != nil {
|
||||
return fmt.Errorf("import node: host=%s, err=%s", coord.URI, err)
|
||||
}
|
||||
|
||||
return nil
|
||||
|
|
@ -336,7 +383,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,27 +394,8 @@ func (c *InternalClient) marshalImportPayload(index, field string, shard uint64,
|
|||
Field: field,
|
||||
Shard: shard,
|
||||
RowIDs: rowIDs,
|
||||
ColumnIDs: columnIDs,
|
||||
Timestamps: timestamps,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal import request: %s", err)
|
||||
}
|
||||
return buf, nil
|
||||
}
|
||||
|
||||
// marshalImportPayloadK marshalls the import parameters into a protobuf byte slice.
|
||||
func (c *InternalClient) marshalImportPayloadK(index, field string, bits []pilosa.Bit) ([]byte, error) {
|
||||
// Separate row and column IDs to reduce allocations.
|
||||
rowKeys := Bits(bits).RowKeys()
|
||||
columnKeys := Bits(bits).ColumnKeys()
|
||||
timestamps := Bits(bits).Timestamps()
|
||||
|
||||
// Marshal data to protobuf.
|
||||
buf, err := c.serializer.Marshal(&pilosa.ImportRequest{
|
||||
Index: index,
|
||||
Field: field,
|
||||
RowKeys: rowKeys,
|
||||
ColumnIDs: columnIDs,
|
||||
ColumnKeys: columnKeys,
|
||||
Timestamps: timestamps,
|
||||
})
|
||||
|
|
@ -447,15 +477,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 +890,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 +924,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 +936,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 +948,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 +993,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 +1015,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))
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ import (
|
|||
gohttp "net/http"
|
||||
"reflect"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/davecgh/go-spew/spew"
|
||||
"github.com/pilosa/pilosa"
|
||||
|
|
@ -202,6 +203,181 @@ func TestClient_Import(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
// Ensure client can bulk import data.
|
||||
func TestClient_ImportKeys(t *testing.T) {
|
||||
t.Run("SingleNode", func(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, http.GetHTTPClient(nil))
|
||||
|
||||
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.Fatalf("unexpected response type %T", resp.Results[0])
|
||||
} else if !reflect.DeepEqual(pairs, []pilosa.Pair{
|
||||
{Key: "green", Count: 3},
|
||||
{Key: "blue", Count: 2},
|
||||
{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.Fatalf("unexpected response type %T", resp.Results[0])
|
||||
} else if !reflect.DeepEqual(pairs, []pilosa.Pair{
|
||||
{ID: 1, Count: 3},
|
||||
{ID: 2, Count: 2},
|
||||
{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.Fatalf("unexpected response type %T", resp.Results[0])
|
||||
} else if !reflect.DeepEqual(pairs, []pilosa.Pair{
|
||||
{Key: "green", Count: 3},
|
||||
{Key: "blue", Count: 2},
|
||||
{Key: "purple", Count: 1},
|
||||
}) {
|
||||
t.Fatalf("unexpected topn result: %v", pairs)
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("MultiNode", func(t *testing.T) {
|
||||
cluster := test.MustRunCluster(t, 2)
|
||||
cmd0 := cluster[0]
|
||||
cmd1 := cluster[1]
|
||||
host0 := cmd0.URL()
|
||||
host1 := cmd1.URL()
|
||||
|
||||
cmd0.MustCreateIndex(t, "keyed", pilosa.IndexOptions{Keys: true})
|
||||
cmd0.MustCreateField(t, "keyed", "keyedf0", pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 1000), pilosa.OptFieldKeys())
|
||||
cmd0.MustCreateField(t, "keyed", "keyedf1", pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 1000), pilosa.OptFieldKeys())
|
||||
|
||||
// Send import request.
|
||||
c0 := MustNewClient(host0, http.GetHTTPClient(nil))
|
||||
c1 := MustNewClient(host1, http.GetHTTPClient(nil))
|
||||
|
||||
// Import to node0.
|
||||
t.Run("Import node0", func(t *testing.T) {
|
||||
if err := c0.ImportK(context.Background(), "keyed", "keyedf0", []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)
|
||||
}
|
||||
cmd0.MustRecalculateCaches(t)
|
||||
resp := cmd0.MustQuery(t, &pilosa.QueryRequest{
|
||||
Index: "keyed",
|
||||
Query: "TopN(keyedf0)",
|
||||
})
|
||||
if pairs, ok := resp.Results[0].([]pilosa.Pair); !ok {
|
||||
t.Fatalf("unexpected response type %T", resp.Results[0])
|
||||
} else if !reflect.DeepEqual(pairs, []pilosa.Pair{
|
||||
{Key: "green", Count: 3},
|
||||
{Key: "blue", Count: 2},
|
||||
{Key: "purple", Count: 1},
|
||||
}) {
|
||||
t.Fatalf("unexpected topn result: %v", pairs)
|
||||
}
|
||||
})
|
||||
|
||||
// Import to node1 (ensure import is routed to coordinator for translation).
|
||||
t.Run("Import node1", func(t *testing.T) {
|
||||
if err := c1.ImportK(context.Background(), "keyed", "keyedf1", []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)
|
||||
}
|
||||
|
||||
// Wait for translation replication.
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
|
||||
cmd1.MustRecalculateCaches(t)
|
||||
resp := cmd1.MustQuery(t, &pilosa.QueryRequest{
|
||||
Index: "keyed",
|
||||
Query: "TopN(keyedf1)",
|
||||
})
|
||||
if pairs, ok := resp.Results[0].([]pilosa.Pair); !ok {
|
||||
t.Fatalf("unexpected response type %T", resp.Results[0])
|
||||
} else if !reflect.DeepEqual(pairs, []pilosa.Pair{
|
||||
{Key: "green", Count: 3},
|
||||
{Key: "blue", Count: 2},
|
||||
{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]
|
||||
|
|
|
|||
|
|
@ -232,6 +232,7 @@ func newRouter(handler *Handler) *mux.Router {
|
|||
router.HandleFunc("/internal/fragment/nodes", handler.handleGetFragmentNodes).Methods("GET").Name("GetFragmentNodes")
|
||||
router.HandleFunc("/internal/index/{index}/attr/diff", handler.handlePostIndexAttrDiff).Methods("POST")
|
||||
router.HandleFunc("/internal/index/{index}/field/{field}/attr/diff", handler.handlePostFieldAttrDiff).Methods("POST")
|
||||
router.HandleFunc("/internal/nodes", handler.handleGetNodes).Methods("GET").Name("GetNodes")
|
||||
router.HandleFunc("/internal/shards/max", handler.handleGetShardsMax).Methods("GET") // TODO: deprecate, but it's being used by the client
|
||||
router.HandleFunc("/internal/translate/data", handler.handleGetTranslateData).Methods("GET")
|
||||
|
||||
|
|
@ -1062,6 +1063,22 @@ func (h *Handler) handleGetFragmentNodes(w http.ResponseWriter, r *http.Request)
|
|||
}
|
||||
}
|
||||
|
||||
// handleGetNodes handles /internal/nodes requests.
|
||||
func (h *Handler) handleGetNodes(w http.ResponseWriter, r *http.Request) {
|
||||
if !validHeaderAcceptJSON(r.Header) {
|
||||
http.Error(w, "JSON only acceptable response", http.StatusNotAcceptable)
|
||||
return
|
||||
}
|
||||
|
||||
// Retrieve all nodes.
|
||||
nodes := h.api.Hosts(r.Context())
|
||||
|
||||
// Write to response.
|
||||
if err := json.NewEncoder(w).Encode(nodes); err != nil {
|
||||
h.logger.Printf("json write error: %s", err)
|
||||
}
|
||||
}
|
||||
|
||||
// handleGetFragmentBlockData handles GET /internal/fragment/block/data requests.
|
||||
func (h *Handler) handleGetFragmentBlockData(w http.ResponseWriter, r *http.Request) {
|
||||
buf, err := h.api.FragmentBlockData(r.Context(), r.Body)
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue