mirror of
https://github.com/featurebasedb/featurebase.git
synced 2026-09-07 00:55:55 +00:00
move sort to subcommand
This commit is contained in:
Matt Jaffee
parent
362539fe2d
commit
312a9a5334
3 changed files with 183 additions and 164 deletions
|
|
@ -1,20 +1,15 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/csv"
|
||||
"errors"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"log"
|
||||
"math/rand"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
"text/tabwriter"
|
||||
|
|
@ -95,7 +90,6 @@ Usage:
|
|||
|
||||
The commands are:
|
||||
|
||||
sort sorts a data file for optimal import speed
|
||||
backup backs up a frame to an archive file
|
||||
restore restores a frame from an archive file
|
||||
inspect inspects fragment data files
|
||||
|
|
@ -122,8 +116,6 @@ func (m *Main) ParseFlags(args []string) error {
|
|||
fmt.Fprintln(m.Stderr, m.Usage())
|
||||
fmt.Fprintln(m.Stderr, "")
|
||||
return flag.ErrHelp
|
||||
case "sort":
|
||||
m.Cmd = NewSortCommand(m.Stdin, m.Stdout, m.Stderr)
|
||||
case "backup":
|
||||
m.Cmd = NewBackupCommand(m.Stdin, m.Stdout, m.Stderr)
|
||||
case "restore":
|
||||
|
|
@ -157,120 +149,6 @@ type Command interface {
|
|||
Run(context.Context) error
|
||||
}
|
||||
|
||||
// SortCommand represents a command for sorting import data.
|
||||
type SortCommand struct {
|
||||
// Filename to sort
|
||||
Path string
|
||||
|
||||
// Standard input/output
|
||||
Stdin io.Reader
|
||||
Stdout io.Writer
|
||||
Stderr io.Writer
|
||||
}
|
||||
|
||||
// NewSortCommand returns a new instance of SortCommand.
|
||||
func NewSortCommand(stdin io.Reader, stdout, stderr io.Writer) *SortCommand {
|
||||
return &SortCommand{
|
||||
Stdin: stdin,
|
||||
Stdout: stdout,
|
||||
Stderr: stderr,
|
||||
}
|
||||
}
|
||||
|
||||
// ParseFlags parses command line flags from args.
|
||||
func (cmd *SortCommand) ParseFlags(args []string) error {
|
||||
fs := flag.NewFlagSet("pilosactl", flag.ContinueOnError)
|
||||
fs.SetOutput(ioutil.Discard)
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Extract the data path.
|
||||
if fs.NArg() == 0 {
|
||||
return errors.New("path required")
|
||||
} else if fs.NArg() > 1 {
|
||||
return errors.New("only one path allowed")
|
||||
}
|
||||
cmd.Path = fs.Arg(0)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Usage returns the usage message to be printed.
|
||||
func (cmd *SortCommand) Usage() string {
|
||||
return strings.TrimSpace(`
|
||||
usage: pilosactl sort PATH
|
||||
|
||||
Sorts the import data at PATH into the optimal sort order for importing.
|
||||
|
||||
The format of the CSV file is:
|
||||
|
||||
BITMAPID,PROFILEID
|
||||
|
||||
The file should contain no headers.
|
||||
`)
|
||||
}
|
||||
|
||||
// Run executes the main program execution.
|
||||
func (cmd *SortCommand) Run(ctx context.Context) error {
|
||||
// Open file for reading.
|
||||
f, err := os.Open(cmd.Path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
// Read rows as bits.
|
||||
r := csv.NewReader(f)
|
||||
r.FieldsPerRecord = -1
|
||||
a := make([]pilosa.Bit, 0, 1000000)
|
||||
for {
|
||||
bitmapID, profileID, timestamp, err := readCSVRow(r)
|
||||
if err == io.EOF {
|
||||
break
|
||||
} else if err == errBlank {
|
||||
continue
|
||||
} else if err != nil {
|
||||
return err
|
||||
}
|
||||
a = append(a, pilosa.Bit{BitmapID: bitmapID, ProfileID: profileID, Timestamp: timestamp})
|
||||
}
|
||||
|
||||
// Sort bits by position.
|
||||
sort.Sort(pilosa.BitsByPos(a))
|
||||
|
||||
// Rewrite to STDOUT.
|
||||
w := bufio.NewWriter(cmd.Stdout)
|
||||
buf := make([]byte, 0, 1024)
|
||||
for _, bit := range a {
|
||||
// Write CSV to buffer.
|
||||
buf = buf[:0]
|
||||
buf = strconv.AppendUint(buf, bit.BitmapID, 10)
|
||||
|
||||
buf = append(buf, ',')
|
||||
buf = strconv.AppendUint(buf, bit.ProfileID, 10)
|
||||
|
||||
if bit.Timestamp != 0 {
|
||||
buf = append(buf, ',')
|
||||
buf = append(buf, time.Unix(0, bit.Timestamp).UTC().Format(pilosa.TimeFormat)...)
|
||||
}
|
||||
|
||||
buf = append(buf, '\n')
|
||||
|
||||
// Write to output.
|
||||
if _, err := w.Write(buf); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Ensure buffer is flushed before exiting.
|
||||
if err := w.Flush(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// BackupCommand represents a command for backing up a frame.
|
||||
type BackupCommand struct {
|
||||
// Destination host and port.
|
||||
|
|
@ -808,45 +686,3 @@ func (cmd *BenchCommand) runSetBit(ctx context.Context, client *pilosa.Client) e
|
|||
|
||||
return nil
|
||||
}
|
||||
|
||||
// readCSVRow reads a bitmap/profile pair from a CSV row.
|
||||
func readCSVRow(r *csv.Reader) (bitmapID, profileID uint64, timestamp int64, err error) {
|
||||
// Read CSV row.
|
||||
record, err := r.Read()
|
||||
if err != nil {
|
||||
return 0, 0, 0, err
|
||||
}
|
||||
|
||||
// Ignore blank rows.
|
||||
if record[0] == "" {
|
||||
return 0, 0, 0, errBlank
|
||||
} else if len(record) < 2 {
|
||||
return 0, 0, 0, fmt.Errorf("bad column count: %d", len(record))
|
||||
}
|
||||
|
||||
// Parse bitmap id.
|
||||
bitmapID, err = strconv.ParseUint(record[0], 10, 64)
|
||||
if err != nil {
|
||||
return 0, 0, 0, fmt.Errorf("invalid bitmap id: %q", record[0])
|
||||
}
|
||||
|
||||
// Parse bitmap id.
|
||||
profileID, err = strconv.ParseUint(record[1], 10, 64)
|
||||
if err != nil {
|
||||
return 0, 0, 0, fmt.Errorf("invalid profile id: %q", record[1])
|
||||
}
|
||||
|
||||
// Parse timestamp, if available.
|
||||
if len(record) > 2 && record[2] != "" {
|
||||
t, err := time.Parse(pilosa.TimeFormat, record[2])
|
||||
if err != nil {
|
||||
return 0, 0, 0, fmt.Errorf("invalid timestamp: %q", record[2])
|
||||
}
|
||||
timestamp = t.UnixNano()
|
||||
}
|
||||
|
||||
return bitmapID, profileID, timestamp, nil
|
||||
}
|
||||
|
||||
// errBlank indicates a blank row in a CSV file.
|
||||
var errBlank = errors.New("blank row")
|
||||
|
|
|
|||
45
cmd/sort.go
Normal file
45
cmd/sort.go
Normal file
|
|
@ -0,0 +1,45 @@
|
|||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
"github.com/pilosa/pilosa/ctl"
|
||||
)
|
||||
|
||||
var sorter = ctl.NewSortCommand(os.Stdin, os.Stdout, os.Stderr)
|
||||
|
||||
var sortCmd = &cobra.Command{
|
||||
Use: "sort <path>",
|
||||
Short: "sort - sort import data for optimal import performance",
|
||||
Long: `
|
||||
Sorts the import data at PATH into the optimal sort order for importing.
|
||||
|
||||
The format of the CSV file is:
|
||||
|
||||
BITMAPID,PROFILEID
|
||||
|
||||
The file should contain no headers.
|
||||
`,
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
fmt.Println(cmd.Flags())
|
||||
if len(args) == 0 {
|
||||
fmt.Println("path required")
|
||||
return
|
||||
} else if len(args) > 1 {
|
||||
fmt.Println("only one path supported")
|
||||
return
|
||||
}
|
||||
sorter.Path = args[0]
|
||||
if err := sorter.Run(context.Background()); err != nil {
|
||||
fmt.Println(err)
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
func init() {
|
||||
RootCmd.AddCommand(sortCmd)
|
||||
}
|
||||
138
ctl/sort.go
Normal file
138
ctl/sort.go
Normal file
|
|
@ -0,0 +1,138 @@
|
|||
package ctl
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/csv"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"sort"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/pilosa/pilosa"
|
||||
)
|
||||
|
||||
// SortCommand represents a command for sorting import data.
|
||||
type SortCommand struct {
|
||||
// Filename to sort
|
||||
Path string
|
||||
|
||||
// Standard input/output
|
||||
Stdin io.Reader
|
||||
Stdout io.Writer
|
||||
Stderr io.Writer
|
||||
}
|
||||
|
||||
// NewSortCommand returns a new instance of SortCommand.
|
||||
func NewSortCommand(stdin io.Reader, stdout, stderr io.Writer) *SortCommand {
|
||||
return &SortCommand{
|
||||
Stdin: stdin,
|
||||
Stdout: stdout,
|
||||
Stderr: stderr,
|
||||
}
|
||||
}
|
||||
|
||||
// Run executes the main program execution.
|
||||
func (cmd *SortCommand) Run(ctx context.Context) error {
|
||||
// Open file for reading.
|
||||
f, err := os.Open(cmd.Path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
// Read rows as bits.
|
||||
r := csv.NewReader(f)
|
||||
r.FieldsPerRecord = -1
|
||||
a := make([]pilosa.Bit, 0, 1000000)
|
||||
for {
|
||||
bitmapID, profileID, timestamp, err := readCSVRow(r)
|
||||
if err == io.EOF {
|
||||
break
|
||||
} else if err == errBlank {
|
||||
continue
|
||||
} else if err != nil {
|
||||
return err
|
||||
}
|
||||
a = append(a, pilosa.Bit{BitmapID: bitmapID, ProfileID: profileID, Timestamp: timestamp})
|
||||
}
|
||||
|
||||
// Sort bits by position.
|
||||
sort.Sort(pilosa.BitsByPos(a))
|
||||
|
||||
// Rewrite to STDOUT.
|
||||
w := bufio.NewWriter(cmd.Stdout)
|
||||
buf := make([]byte, 0, 1024)
|
||||
for _, bit := range a {
|
||||
// Write CSV to buffer.
|
||||
buf = buf[:0]
|
||||
buf = strconv.AppendUint(buf, bit.BitmapID, 10)
|
||||
|
||||
buf = append(buf, ',')
|
||||
buf = strconv.AppendUint(buf, bit.ProfileID, 10)
|
||||
|
||||
if bit.Timestamp != 0 {
|
||||
buf = append(buf, ',')
|
||||
buf = append(buf, time.Unix(0, bit.Timestamp).UTC().Format(pilosa.TimeFormat)...)
|
||||
}
|
||||
|
||||
buf = append(buf, '\n')
|
||||
|
||||
// Write to output.
|
||||
if _, err := w.Write(buf); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Ensure buffer is flushed before exiting.
|
||||
if err := w.Flush(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// readCSVRow reads a bitmap/profile pair from a CSV row.
|
||||
func readCSVRow(r *csv.Reader) (bitmapID, profileID uint64, timestamp int64, err error) {
|
||||
// Read CSV row.
|
||||
record, err := r.Read()
|
||||
if err != nil {
|
||||
return 0, 0, 0, err
|
||||
}
|
||||
|
||||
// Ignore blank rows.
|
||||
if record[0] == "" {
|
||||
return 0, 0, 0, errBlank
|
||||
} else if len(record) < 2 {
|
||||
return 0, 0, 0, fmt.Errorf("bad column count: %d", len(record))
|
||||
}
|
||||
|
||||
// Parse bitmap id.
|
||||
bitmapID, err = strconv.ParseUint(record[0], 10, 64)
|
||||
if err != nil {
|
||||
return 0, 0, 0, fmt.Errorf("invalid bitmap id: %q", record[0])
|
||||
}
|
||||
|
||||
// Parse bitmap id.
|
||||
profileID, err = strconv.ParseUint(record[1], 10, 64)
|
||||
if err != nil {
|
||||
return 0, 0, 0, fmt.Errorf("invalid profile id: %q", record[1])
|
||||
}
|
||||
|
||||
// Parse timestamp, if available.
|
||||
if len(record) > 2 && record[2] != "" {
|
||||
t, err := time.Parse(pilosa.TimeFormat, record[2])
|
||||
if err != nil {
|
||||
return 0, 0, 0, fmt.Errorf("invalid timestamp: %q", record[2])
|
||||
}
|
||||
timestamp = t.UnixNano()
|
||||
}
|
||||
|
||||
return bitmapID, profileID, timestamp, nil
|
||||
}
|
||||
|
||||
// errBlank indicates a blank row in a CSV file.
|
||||
var errBlank = errors.New("blank row")
|
||||
Loading…
Add table
Reference in a new issue