From e5edf3e432e2ca71904e73e1ffaeb3ce2d9e2692 Mon Sep 17 00:00:00 2001 From: jaffee Date: Mon, 5 Dec 2016 17:37:33 -0600 Subject: [PATCH] add import benchmark --- bench/import.go | 154 ++++++++++++++++++++++++++++++++ bench/import_test.go | 84 ++++++++++++++++++ cmd/pilosactl/main.go | 187 +-------------------------------------- pilosactl/import.go | 201 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 441 insertions(+), 185 deletions(-) create mode 100644 bench/import.go create mode 100644 bench/import_test.go create mode 100644 pilosactl/import.go diff --git a/bench/import.go b/bench/import.go new file mode 100644 index 000000000..b9f5304e2 --- /dev/null +++ b/bench/import.go @@ -0,0 +1,154 @@ +package bench + +import ( + "context" + "flag" + "fmt" + "io" + "io/ioutil" + "math/rand" + + "github.com/pilosa/pilosa/pilosactl" +) + +// 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 + + 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 { + var err error + b.Client, err = firstHostClient(hosts) + if err != nil { + return err + } + // generate csv data + baseBitmapID, maxBitmapID, baseProfileID, maxProfileID := b.BaseBitmapID, b.MaxBitmapID, b.BaseProfileID, b.MaxProfileID + if b.AgentControls == "height" { + numBitmapIDs := (b.MaxBitmapID - b.BaseBitmapID) + baseBitmapID = b.BaseBitmapID + (numBitmapIDs * int64(agentNum)) + maxBitmapID = baseBitmapID + numBitmapIDs + } + if b.AgentControls == "height" { + numProfileIDs := (b.MaxProfileID - b.BaseProfileID) + baseProfileID = b.BaseProfileID + (numProfileIDs * int64(agentNum)) + maxProfileID = baseProfileID + numProfileIDs + } + f, err := ioutil.TempFile("", "") + if err != nil { + return err + } + GenerateImportCSV(f, baseBitmapID, maxBitmapID, baseProfileID, maxProfileID, + b.MinBitsPerMap, b.MaxBitsPerMap, b.Seed+int64(agentNum), b.RandomBitmapOrder) + // set b.Paths + 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{}) + b.ImportCommand.Run(context.TODO()) + return results +} + +func GenerateImportCSV(w io.Writer, baseBitmapID, maxBitmapID, baseProfileID, maxProfileID, minBitsPerMap, maxBitsPerMap, seed int64, randomOrder bool) { + src := rand.NewSource(seed) + rng := rand.New(src) + + var bitmapIDs []int + if randomOrder { + bitmapIDs = rng.Perm(int(maxBitmapID - baseBitmapID)) + } + 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) + } + } +} diff --git a/bench/import_test.go b/bench/import_test.go new file mode 100644 index 000000000..879842e84 --- /dev/null +++ b/bench/import_test.go @@ -0,0 +1,84 @@ +package bench_test + +import ( + "bytes" + "testing" + + "io/ioutil" + + "github.com/pilosa/pilosa/bench" +) + +func TestGenerateImportCSVNonRand(t *testing.T) { + b := bytes.NewBuffer(make([]byte, 0)) + + bench.GenerateImportCSV(b, 0, 10, 21, 29, 2, 3, 0, false) + + bytes, err := ioutil.ReadAll(b) + if err != nil { + t.Fatalf("Error reading buffer: %v", err) + } + + expected := ` +0,21 +0,24 +1,23 +1,28 +2,21 +2,22 +3,21 +3,25 +4,23 +4,21 +5,28 +5,28 +6,25 +6,23 +7,26 +7,23 +8,21 +8,23 +9,25 +9,24 +`[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)) + } +} diff --git a/cmd/pilosactl/main.go b/cmd/pilosactl/main.go index 95e2ee6a8..53028bad1 100644 --- a/cmd/pilosactl/main.go +++ b/cmd/pilosactl/main.go @@ -27,6 +27,7 @@ import ( "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 +123,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 +221,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. diff --git a/pilosactl/import.go b/pilosactl/import.go new file mode 100644 index 000000000..1d9358cce --- /dev/null +++ b/pilosactl/import.go @@ -0,0 +1,201 @@ +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, + } +} + +// 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 +}