mirror of
https://github.com/featurebasedb/featurebase.git
synced 2026-09-06 00:25:55 +00:00
Merge pull request #186 from jaffee/170-import-benchmark
170 import benchmark
This commit is contained in:
tgruben
GitHub
commit
010a71263b
5 changed files with 541 additions and 188 deletions
175
bench/import.go
Normal file
175
bench/import.go
Normal file
|
|
@ -0,0 +1,175 @@
|
|||
package bench
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"math/rand"
|
||||
|
||||
"github.com/pilosa/pilosa/pilosactl"
|
||||
)
|
||||
|
||||
func NewImport(stdin io.Reader, stdout, stderr io.Writer) *Import {
|
||||
return &Import{
|
||||
ImportCommand: pilosactl.NewImportCommand(stdin, stdout, stderr),
|
||||
}
|
||||
}
|
||||
|
||||
// Import sets bits with increasing profile id and bitmap id.
|
||||
type Import struct {
|
||||
BaseBitmapID int64
|
||||
MaxBitmapID int64
|
||||
BaseProfileID int64
|
||||
MaxProfileID int64
|
||||
RandomBitmapOrder bool
|
||||
MinBitsPerMap int64
|
||||
MaxBitsPerMap int64
|
||||
AgentControls string
|
||||
Seed int64
|
||||
numbits int
|
||||
|
||||
*pilosactl.ImportCommand
|
||||
}
|
||||
|
||||
func (b *Import) Usage() string {
|
||||
return `
|
||||
import generates an import file and imports using pilosa's bulk import interface
|
||||
|
||||
Usage: import [arguments]
|
||||
|
||||
The following arguments are available:
|
||||
|
||||
-base-bitmap-id int
|
||||
bits being set will all be greater than this
|
||||
|
||||
-maximum-bitmap-id int
|
||||
bits being set will all be less than this
|
||||
|
||||
-base-profile-id int
|
||||
profile id num to start from
|
||||
|
||||
-max-profile-id int
|
||||
maximum profile id to generate
|
||||
|
||||
-random-bitmap-order
|
||||
if this option is set, the import file will not be sorted by bitmap id
|
||||
|
||||
-min-bits-per-map int
|
||||
minimum number of bits set per bitmap
|
||||
|
||||
-max-bits-per-map int
|
||||
maximum number of bits set per bitmap
|
||||
|
||||
-agent-controls string
|
||||
can be 'height', 'width', or empty (TODO or square?)- increasing
|
||||
number of agents modulates bitmap id range, profile id range,
|
||||
or just sets more bits in the same range.
|
||||
|
||||
-seed int
|
||||
seed for RNG
|
||||
|
||||
-db string
|
||||
pilosa db to use
|
||||
|
||||
-frame string
|
||||
frame to import into
|
||||
`[1:]
|
||||
}
|
||||
|
||||
func (b *Import) ConsumeFlags(args []string) ([]string, error) {
|
||||
fs := flag.NewFlagSet("Import", flag.ContinueOnError)
|
||||
fs.SetOutput(ioutil.Discard)
|
||||
fs.Int64Var(&b.BaseBitmapID, "base-bitmap-id", 0, "")
|
||||
fs.Int64Var(&b.MaxBitmapID, "max-bitmap-id", 1000, "")
|
||||
fs.Int64Var(&b.BaseProfileID, "base-profile-id", 0, "")
|
||||
fs.Int64Var(&b.MaxProfileID, "max-profile-id", 1000, "")
|
||||
fs.BoolVar(&b.RandomBitmapOrder, "random-bitmap-order", false, "")
|
||||
fs.Int64Var(&b.MinBitsPerMap, "min-bits-per-map", 0, "")
|
||||
fs.Int64Var(&b.MaxBitsPerMap, "max-bits-per-map", 10, "")
|
||||
fs.StringVar(&b.AgentControls, "agent-controls", "", "")
|
||||
fs.Int64Var(&b.Seed, "seed", 0, "")
|
||||
fs.StringVar(&b.Database, "db", "benchdb", "")
|
||||
fs.StringVar(&b.Frame, "frame", "testframe", "")
|
||||
fs.IntVar(&b.BufferSize, "buffer-size", 10000000, "")
|
||||
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return fs.Args(), nil
|
||||
}
|
||||
|
||||
func (b *Import) Init(hosts []string, agentNum int) error {
|
||||
if len(hosts) == 0 {
|
||||
return fmt.Errorf("Need at least one host")
|
||||
}
|
||||
b.Host = hosts[0]
|
||||
// generate csv data
|
||||
baseBitmapID, maxBitmapID, baseProfileID, maxProfileID := b.BaseBitmapID, b.MaxBitmapID, b.BaseProfileID, b.MaxProfileID
|
||||
switch b.AgentControls {
|
||||
case "height":
|
||||
numBitmapIDs := (b.MaxBitmapID - b.BaseBitmapID)
|
||||
baseBitmapID = b.BaseBitmapID + (numBitmapIDs * int64(agentNum))
|
||||
maxBitmapID = baseBitmapID + numBitmapIDs
|
||||
case "width":
|
||||
numProfileIDs := (b.MaxProfileID - b.BaseProfileID)
|
||||
baseProfileID = b.BaseProfileID + (numProfileIDs * int64(agentNum))
|
||||
maxProfileID = baseProfileID + numProfileIDs
|
||||
case "":
|
||||
break
|
||||
default:
|
||||
return fmt.Errorf("agent-controls: '%v' is not supported", b.AgentControls)
|
||||
}
|
||||
f, err := ioutil.TempFile("", "")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// set b.Paths)
|
||||
num := GenerateImportCSV(f, baseBitmapID, maxBitmapID, baseProfileID, maxProfileID,
|
||||
b.MinBitsPerMap, b.MaxBitsPerMap, b.Seed+int64(agentNum), b.RandomBitmapOrder)
|
||||
b.numbits = num
|
||||
// set b.Paths
|
||||
f.Close()
|
||||
b.Paths = []string{f.Name()}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Run runs the Import benchmark
|
||||
func (b *Import) Run(agentNum int) map[string]interface{} {
|
||||
results := make(map[string]interface{})
|
||||
err := b.ImportCommand.Run(context.TODO())
|
||||
if err != nil {
|
||||
results["error"] = err.Error()
|
||||
}
|
||||
results["numbits"] = b.numbits
|
||||
results["config"] = *b
|
||||
return results
|
||||
}
|
||||
|
||||
func GenerateImportCSV(w io.Writer, baseBitmapID, maxBitmapID, baseProfileID, maxProfileID, minBitsPerMap, maxBitsPerMap, seed int64, randomOrder bool) int {
|
||||
src := rand.NewSource(seed)
|
||||
rng := rand.New(src)
|
||||
|
||||
var bitmapIDs []int
|
||||
if randomOrder {
|
||||
bitmapIDs = rng.Perm(int(maxBitmapID - baseBitmapID))
|
||||
}
|
||||
numrows := 0
|
||||
for i := baseBitmapID; i < maxBitmapID; i++ {
|
||||
var bitmapID int64
|
||||
if randomOrder {
|
||||
bitmapID = int64(bitmapIDs[i-baseBitmapID])
|
||||
} else {
|
||||
bitmapID = int64(i)
|
||||
}
|
||||
|
||||
numBitsToSet := rng.Int63n(maxBitsPerMap-minBitsPerMap) + minBitsPerMap
|
||||
for j := int64(0); j < numBitsToSet; j++ {
|
||||
profileID := rng.Int63n(maxProfileID-baseProfileID) + baseProfileID
|
||||
fmt.Fprintf(w, "%d,%d\n", bitmapID, profileID)
|
||||
numrows += 1
|
||||
}
|
||||
}
|
||||
return numrows
|
||||
}
|
||||
140
bench/import_test.go
Normal file
140
bench/import_test.go
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
package bench_test
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"log"
|
||||
"testing"
|
||||
|
||||
"io/ioutil"
|
||||
|
||||
"os"
|
||||
|
||||
"github.com/pilosa/pilosa/bench"
|
||||
)
|
||||
|
||||
func TestImportInit(t *testing.T) {
|
||||
imp := bench.Import{
|
||||
BaseBitmapID: 0,
|
||||
MaxBitmapID: 10,
|
||||
BaseProfileID: 0,
|
||||
MaxProfileID: 10,
|
||||
RandomBitmapOrder: false,
|
||||
MinBitsPerMap: 2,
|
||||
MaxBitsPerMap: 3,
|
||||
AgentControls: "width",
|
||||
Seed: 0,
|
||||
}
|
||||
|
||||
imp.Init([]string{"blah"}, 2)
|
||||
f, err := os.Open(imp.Paths[0])
|
||||
if err != nil {
|
||||
t.Fatalf("Couldn't open file: %v, err: %v", imp.Paths[0], err)
|
||||
}
|
||||
bytes, err := ioutil.ReadAll(f)
|
||||
if err != nil {
|
||||
t.Fatalf("error reading file: %v", err)
|
||||
}
|
||||
|
||||
expected := `
|
||||
0,21
|
||||
0,22
|
||||
1,22
|
||||
1,20
|
||||
2,22
|
||||
2,26
|
||||
3,21
|
||||
3,23
|
||||
4,21
|
||||
4,22
|
||||
5,20
|
||||
5,28
|
||||
6,23
|
||||
6,27
|
||||
7,20
|
||||
7,20
|
||||
8,29
|
||||
8,23
|
||||
9,29
|
||||
9,23
|
||||
`[1:]
|
||||
|
||||
if string(bytes) != expected {
|
||||
t.Fatalf("unexpected result: %v", string(bytes))
|
||||
}
|
||||
|
||||
log.Println(imp)
|
||||
}
|
||||
|
||||
func TestGenerateImportCSVNonRand(t *testing.T) {
|
||||
b := bytes.NewBuffer(make([]byte, 0))
|
||||
|
||||
bench.GenerateImportCSV(b, 0, 10, 20, 30, 2, 3, 2, false)
|
||||
|
||||
bytes, err := ioutil.ReadAll(b)
|
||||
if err != nil {
|
||||
t.Fatalf("Error reading buffer: %v", err)
|
||||
}
|
||||
|
||||
expected := `
|
||||
0,21
|
||||
0,22
|
||||
1,22
|
||||
1,20
|
||||
2,22
|
||||
2,26
|
||||
3,21
|
||||
3,23
|
||||
4,21
|
||||
4,22
|
||||
5,20
|
||||
5,28
|
||||
6,23
|
||||
6,27
|
||||
7,20
|
||||
7,20
|
||||
8,29
|
||||
8,23
|
||||
9,29
|
||||
9,23
|
||||
`[1:]
|
||||
|
||||
if string(bytes) != expected {
|
||||
t.Fatalf("unexpected value for generated csv: \n%v", string(bytes))
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateImportCSVRand(t *testing.T) {
|
||||
b := bytes.NewBuffer(make([]byte, 0))
|
||||
|
||||
bench.GenerateImportCSV(b, 0, 10, 21, 29, 1, 4, 0, true)
|
||||
|
||||
bytes, err := ioutil.ReadAll(b)
|
||||
if err != nil {
|
||||
t.Fatalf("Error reading buffer: %v", err)
|
||||
}
|
||||
|
||||
expected := `
|
||||
8,25
|
||||
2,23
|
||||
3,22
|
||||
3,28
|
||||
3,28
|
||||
0,25
|
||||
0,23
|
||||
5,26
|
||||
5,23
|
||||
7,21
|
||||
7,23
|
||||
1,25
|
||||
6,23
|
||||
6,27
|
||||
6,25
|
||||
9,26
|
||||
4,22
|
||||
4,24
|
||||
`[1:]
|
||||
|
||||
if string(bytes) != expected {
|
||||
t.Fatalf("unexpected value for generated csv: \n%v", string(bytes))
|
||||
}
|
||||
}
|
||||
15
cmd/pilosactl/import.json
Normal file
15
cmd/pilosactl/import.json
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
{
|
||||
"PilosaHosts": ["localhost:19327"],
|
||||
"CreatorArgs": ["-type", "local", "-serverN", "1", "-replicaN", "1"],
|
||||
"Agents": { "Type": "local" },
|
||||
"Benchmarks": [
|
||||
{
|
||||
"Num": 1,
|
||||
"Args": ["import", "-max-bitmap-id", "100000", "-max-profile-id", "10000", "-max-bits-per-map", "100", "-seed", "0", "-agent-controls", "width"]
|
||||
},
|
||||
{
|
||||
"Num": 1,
|
||||
"Args": ["import", "-max-bitmap-id", "100000", "-max-profile-id", "10000", "-max-bits-per-map", "100", "-seed", "0", "-agent-controls", "width", "-random-bitmap-order", "-db", "randoload"]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -24,9 +24,11 @@ import (
|
|||
"unsafe"
|
||||
|
||||
"encoding/json"
|
||||
|
||||
"github.com/pilosa/pilosa"
|
||||
"github.com/pilosa/pilosa/bench"
|
||||
"github.com/pilosa/pilosa/creator"
|
||||
"github.com/pilosa/pilosa/pilosactl"
|
||||
"github.com/pilosa/pilosa/roaring"
|
||||
)
|
||||
|
||||
|
|
@ -122,7 +124,7 @@ func (m *Main) ParseFlags(args []string) error {
|
|||
case "config":
|
||||
m.Cmd = NewConfigCommand(m.Stdin, m.Stdout, m.Stderr)
|
||||
case "import":
|
||||
m.Cmd = NewImportCommand(m.Stdin, m.Stdout, m.Stderr)
|
||||
m.Cmd = pilosactl.NewImportCommand(m.Stdin, m.Stdout, m.Stderr)
|
||||
case "export":
|
||||
m.Cmd = NewExportCommand(m.Stdin, m.Stdout, m.Stderr)
|
||||
case "sort":
|
||||
|
|
@ -220,190 +222,6 @@ path = ""
|
|||
return nil
|
||||
}
|
||||
|
||||
// ImportCommand represents a command for bulk importing data.
|
||||
type ImportCommand struct {
|
||||
// Destination host and port.
|
||||
Host string
|
||||
|
||||
// Name of the database & frame to import into.
|
||||
Database string
|
||||
Frame string
|
||||
|
||||
// Filenames to import from.
|
||||
Paths []string
|
||||
|
||||
// Size of buffer used to chunk import.
|
||||
BufferSize int
|
||||
|
||||
// Reusable client.
|
||||
Client *pilosa.Client
|
||||
|
||||
// Standard input/output
|
||||
Stdin io.Reader
|
||||
Stdout io.Writer
|
||||
Stderr io.Writer
|
||||
}
|
||||
|
||||
// NewImportCommand returns a new instance of ImportCommand.
|
||||
func NewImportCommand(stdin io.Reader, stdout, stderr io.Writer) *ImportCommand {
|
||||
return &ImportCommand{
|
||||
Stdin: stdin,
|
||||
Stdout: stdout,
|
||||
Stderr: stderr,
|
||||
|
||||
BufferSize: 10000000,
|
||||
}
|
||||
}
|
||||
|
||||
// ParseFlags parses command line flags from args.
|
||||
func (cmd *ImportCommand) ParseFlags(args []string) error {
|
||||
fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError)
|
||||
fs.SetOutput(ioutil.Discard)
|
||||
fs.StringVar(&cmd.Host, "host", "localhost:15000", "host:port")
|
||||
fs.StringVar(&cmd.Database, "d", "", "database")
|
||||
fs.StringVar(&cmd.Frame, "f", "", "frame")
|
||||
fs.IntVar(&cmd.BufferSize, "buffer-size", cmd.BufferSize, "buffer size")
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Extract the import paths.
|
||||
cmd.Paths = fs.Args()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Usage returns the usage message to be printed.
|
||||
func (cmd *ImportCommand) Usage() string {
|
||||
return strings.TrimSpace(`
|
||||
usage: pilosactl import -host HOST -d database -f frame paths
|
||||
|
||||
Bulk imports one or more CSV files to a host's database and frame. The bits
|
||||
of the CSV file are grouped by slice for the most efficient import.
|
||||
|
||||
The format of the CSV file is:
|
||||
|
||||
BITMAPID,PROFILEID
|
||||
|
||||
The file should contain no headers.
|
||||
`)
|
||||
}
|
||||
|
||||
// Run executes the main program execution.
|
||||
func (cmd *ImportCommand) Run(ctx context.Context) error {
|
||||
logger := log.New(cmd.Stderr, "", log.LstdFlags)
|
||||
|
||||
// Validate arguments.
|
||||
// Database and frame are validated early before the files are parsed.
|
||||
if cmd.Database == "" {
|
||||
return pilosa.ErrDatabaseRequired
|
||||
} else if cmd.Frame == "" {
|
||||
return pilosa.ErrFrameRequired
|
||||
} else if len(cmd.Paths) == 0 {
|
||||
return ErrPathRequired
|
||||
}
|
||||
|
||||
// Create a client to the server.
|
||||
client, err := pilosa.NewClient(cmd.Host)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cmd.Client = client
|
||||
|
||||
// Import each path and import by slice.
|
||||
for _, path := range cmd.Paths {
|
||||
// Parse path into bits.
|
||||
logger.Printf("parsing: %s", path)
|
||||
if err := cmd.importPath(ctx, path); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// importPath parses a path into bits and imports it to the server.
|
||||
func (cmd *ImportCommand) importPath(ctx context.Context, path string) error {
|
||||
a := make([]pilosa.Bit, 0, cmd.BufferSize)
|
||||
|
||||
// Open file for reading.
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
// Read rows as bits.
|
||||
r := csv.NewReader(f)
|
||||
rnum := 0
|
||||
for {
|
||||
rnum++
|
||||
|
||||
// Read CSV row.
|
||||
record, err := r.Read()
|
||||
if err == io.EOF {
|
||||
break
|
||||
} else if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 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))
|
||||
}
|
||||
|
||||
// Parse bitmap id.
|
||||
bitmapID, err := strconv.ParseUint(record[0], 10, 64)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid bitmap id on row %d: %q", rnum, record[0])
|
||||
}
|
||||
|
||||
// Parse bitmap id.
|
||||
profileID, err := strconv.ParseUint(record[1], 10, 64)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid profile id on row %d: %q", rnum, record[1])
|
||||
}
|
||||
|
||||
a = append(a, pilosa.Bit{BitmapID: bitmapID, ProfileID: profileID})
|
||||
|
||||
// If we've reached the buffer size then import bits.
|
||||
if len(a) == cmd.BufferSize {
|
||||
if err := cmd.importBits(ctx, a); err != nil {
|
||||
return err
|
||||
}
|
||||
a = a[:0]
|
||||
}
|
||||
}
|
||||
|
||||
// If there are still bits in the buffer then flush them.
|
||||
if err := cmd.importBits(ctx, a); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// importPath parses a path into bits and imports it to the server.
|
||||
func (cmd *ImportCommand) importBits(ctx context.Context, bits []pilosa.Bit) error {
|
||||
logger := log.New(cmd.Stderr, "", log.LstdFlags)
|
||||
|
||||
// Group bits by slice.
|
||||
logger.Printf("grouping %d bits", len(bits))
|
||||
bitsBySlice := pilosa.Bits(bits).GroupBySlice()
|
||||
|
||||
// Parse path into bits.
|
||||
for slice, bits := range bitsBySlice {
|
||||
logger.Printf("importing slice: %d, n=%d", slice, len(bits))
|
||||
if err := cmd.Client.Import(ctx, cmd.Database, cmd.Frame, slice, bits); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ExportCommand represents a command for bulk exporting data from a server.
|
||||
type ExportCommand struct {
|
||||
// Remote host and port.
|
||||
|
|
@ -1347,6 +1165,8 @@ func (cmd *BagentCommand) ParseFlags(args []string) error {
|
|||
bm = &bench.MultiDBSetBits{}
|
||||
case "random-query":
|
||||
bm = &bench.RandomQuery{}
|
||||
case "import":
|
||||
bm = bench.NewImport(cmd.Stdin, cmd.Stdout, cmd.Stderr)
|
||||
default:
|
||||
return fmt.Errorf("Unknown benchmark cmd: %v", remArgs[0])
|
||||
}
|
||||
|
|
@ -1387,9 +1207,7 @@ The following arguments are available:
|
|||
random-set-bits
|
||||
multi-db-set-bits
|
||||
random-query
|
||||
|
||||
|
||||
|
||||
import
|
||||
`)
|
||||
}
|
||||
|
||||
|
|
|
|||
205
pilosactl/import.go
Normal file
205
pilosactl/import.go
Normal file
|
|
@ -0,0 +1,205 @@
|
|||
package pilosactl
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/csv"
|
||||
"errors"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"log"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/pilosa/pilosa"
|
||||
)
|
||||
|
||||
// ImportCommand represents a command for bulk importing data.
|
||||
type ImportCommand struct {
|
||||
// Destination host and port.
|
||||
Host string
|
||||
|
||||
// Name of the database & frame to import into.
|
||||
Database string
|
||||
Frame string
|
||||
|
||||
// Filenames to import from.
|
||||
Paths []string
|
||||
|
||||
// Size of buffer used to chunk import.
|
||||
BufferSize int
|
||||
|
||||
// Reusable client.
|
||||
Client *pilosa.Client
|
||||
|
||||
// Standard input/output
|
||||
Stdin io.Reader
|
||||
Stdout io.Writer
|
||||
Stderr io.Writer
|
||||
}
|
||||
|
||||
// NewImportCommand returns a new instance of ImportCommand.
|
||||
func NewImportCommand(stdin io.Reader, stdout, stderr io.Writer) *ImportCommand {
|
||||
return &ImportCommand{
|
||||
Stdin: stdin,
|
||||
Stdout: stdout,
|
||||
Stderr: stderr,
|
||||
|
||||
BufferSize: 10000000,
|
||||
}
|
||||
}
|
||||
|
||||
func (cmd *ImportCommand) String() string {
|
||||
return fmt.Sprint(*cmd)
|
||||
}
|
||||
|
||||
// ParseFlags parses command line flags from args.
|
||||
func (cmd *ImportCommand) ParseFlags(args []string) error {
|
||||
fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError)
|
||||
fs.SetOutput(ioutil.Discard)
|
||||
fs.StringVar(&cmd.Host, "host", "localhost:15000", "host:port")
|
||||
fs.StringVar(&cmd.Database, "d", "", "database")
|
||||
fs.StringVar(&cmd.Frame, "f", "", "frame")
|
||||
fs.IntVar(&cmd.BufferSize, "buffer-size", cmd.BufferSize, "buffer size")
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Extract the import paths.
|
||||
cmd.Paths = fs.Args()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Usage returns the usage message to be printed.
|
||||
func (cmd *ImportCommand) Usage() string {
|
||||
return strings.TrimSpace(`
|
||||
usage: pilosactl import -host HOST -d database -f frame paths
|
||||
|
||||
Bulk imports one or more CSV files to a host's database and frame. The bits
|
||||
of the CSV file are grouped by slice for the most efficient import.
|
||||
|
||||
The format of the CSV file is:
|
||||
|
||||
BITMAPID,PROFILEID
|
||||
|
||||
The file should contain no headers.
|
||||
`)
|
||||
}
|
||||
|
||||
// Run executes the main program execution.
|
||||
func (cmd *ImportCommand) Run(ctx context.Context) error {
|
||||
logger := log.New(cmd.Stderr, "", log.LstdFlags)
|
||||
|
||||
// Validate arguments.
|
||||
// Database and frame are validated early before the files are parsed.
|
||||
if cmd.Database == "" {
|
||||
return pilosa.ErrDatabaseRequired
|
||||
} else if cmd.Frame == "" {
|
||||
return pilosa.ErrFrameRequired
|
||||
} else if len(cmd.Paths) == 0 {
|
||||
return errors.New("path required")
|
||||
}
|
||||
|
||||
// Create a client to the server.
|
||||
client, err := pilosa.NewClient(cmd.Host)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cmd.Client = client
|
||||
|
||||
// Import each path and import by slice.
|
||||
for _, path := range cmd.Paths {
|
||||
// Parse path into bits.
|
||||
logger.Printf("parsing: %s", path)
|
||||
if err := cmd.importPath(ctx, path); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// importPath parses a path into bits and imports it to the server.
|
||||
func (cmd *ImportCommand) importPath(ctx context.Context, path string) error {
|
||||
a := make([]pilosa.Bit, 0, cmd.BufferSize)
|
||||
|
||||
// Open file for reading.
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
// Read rows as bits.
|
||||
r := csv.NewReader(f)
|
||||
rnum := 0
|
||||
for {
|
||||
rnum++
|
||||
|
||||
// Read CSV row.
|
||||
record, err := r.Read()
|
||||
if err == io.EOF {
|
||||
break
|
||||
} else if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 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))
|
||||
}
|
||||
|
||||
// Parse bitmap id.
|
||||
bitmapID, err := strconv.ParseUint(record[0], 10, 64)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid bitmap id on row %d: %q", rnum, record[0])
|
||||
}
|
||||
|
||||
// Parse bitmap id.
|
||||
profileID, err := strconv.ParseUint(record[1], 10, 64)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid profile id on row %d: %q", rnum, record[1])
|
||||
}
|
||||
|
||||
a = append(a, pilosa.Bit{BitmapID: bitmapID, ProfileID: profileID})
|
||||
|
||||
// If we've reached the buffer size then import bits.
|
||||
if len(a) == cmd.BufferSize {
|
||||
if err := cmd.importBits(ctx, a); err != nil {
|
||||
return err
|
||||
}
|
||||
a = a[:0]
|
||||
}
|
||||
}
|
||||
|
||||
// If there are still bits in the buffer then flush them.
|
||||
if err := cmd.importBits(ctx, a); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// importPath parses a path into bits and imports it to the server.
|
||||
func (cmd *ImportCommand) importBits(ctx context.Context, bits []pilosa.Bit) error {
|
||||
logger := log.New(cmd.Stderr, "", log.LstdFlags)
|
||||
|
||||
// Group bits by slice.
|
||||
logger.Printf("grouping %d bits", len(bits))
|
||||
bitsBySlice := pilosa.Bits(bits).GroupBySlice()
|
||||
|
||||
// Parse path into bits.
|
||||
for slice, bits := range bitsBySlice {
|
||||
logger.Printf("importing slice: %d, n=%d", slice, len(bits))
|
||||
if err := cmd.Client.Import(ctx, cmd.Database, cmd.Frame, slice, bits); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue