Fix pre sort cmd (#2222)

* make ndjson pre_sort parallel, fix several bugs, test

* remove json tags (not needed), rename pre_sort -> presort
This commit is contained in:
Matthew Jaffee 2023-01-25 11:01:58 -06:00 committed by GitHub
parent 468461fbcf
commit 8f1f3c6d06
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 132 additions and 71 deletions

View file

@ -10,7 +10,7 @@ import (
func newPreSortCommand(logdest logger.Logger) *cobra.Command {
cmd := ctl.NewPreSortCommand(logdest)
ccmd := &cobra.Command{
Use: "pre_sort",
Use: "presort",
Short: "Sort records within files into files by FB partition for more efficient ingest",
Long: `
Takes all input files and writes PartitionN numbered files to a directory, where each file contains only records that will go into the partition it is named for.
@ -22,8 +22,11 @@ Takes all input files and writes PartitionN numbered files to a directory, where
flags.StringVarP(&cmd.File, "file", "", "", "Input file or directory.")
flags.StringVarP(&cmd.Table, "table", "", "", "Name of table (used to hash keys to determine partition).")
flags.StringVarP(&cmd.Type, "type", "", cmd.Type, "Input file type (csv or ndjson).")
flags.StringSliceVar(&cmd.PrimaryKeyFields, "primary-key-fields", []string{}, "Names of primary key fields. For CSV there must be a header row and these pulled from there.")
flags.StringSliceVar(&cmd.PrimaryKeyFields, "primary-key-fields", []string{}, "Names of primary key fields. For CSV there must be a header row and these are pulled from there.")
flags.IntVar(&cmd.PartitionN, "partition-n", cmd.PartitionN, "Number of partitions.")
flags.StringVarP(&cmd.OutputDir, "output-dir", "", cmd.OutputDir, "Directory name to write output to.")
flags.StringVarP(&cmd.PrimaryKeySeparator, "primary-key-separator", "", cmd.PrimaryKeySeparator, "Separator to write in between primary key fields, can be empty.")
flags.IntVar(&cmd.JobSize, "job-size", cmd.JobSize, "Number of lines to put into each job (purely a performance tuning parameter, only supported by ndjson mode).")
flags.IntVar(&cmd.NumWorkers, "num-workers", cmd.NumWorkers, "Number of parallel worker routines doing decode->hash->encode. Only supported by ndjson mode.")
return ccmd
}

View file

@ -1,6 +1,7 @@
package ctl
import (
"bufio"
"context"
"encoding/csv"
"encoding/json"
@ -10,22 +11,28 @@ import (
"io/fs"
"os"
"path/filepath"
"sync"
"github.com/featurebasedb/featurebase/v3/logger"
"github.com/pkg/errors"
"golang.org/x/sync/errgroup"
)
type PreSortCommand struct {
// Optional Index filter
File string `json:"file"`
File string
Type string `json:"type"`
Type string
Table string `json:"table"`
Table string
PrimaryKeyFields []string `json:"primary-key-fields"`
PrimaryKeyFields []string
PrimaryKeySeparator string
PartitionN int `json:"partition-n"`
PartitionN int
JobSize int
NumWorkers int
// Path to write sorted files to
OutputDir string
@ -33,17 +40,23 @@ type PreSortCommand struct {
// Standard input/output
logDest logger.Logger
outputFiles map[int]*os.File
outputLocks []sync.Mutex
outputFiles map[int]*os.File
outputEncoders map[int]*json.Encoder
outputWriters map[int]*csv.Writer
}
// NewPreSortCommand returns a new instance of PreSortCommand.
func NewPreSortCommand(logdest logger.Logger) *PreSortCommand {
return &PreSortCommand{
logDest: logdest,
Type: "ndjson",
OutputDir: "presorted_files",
PartitionN: 256,
outputFiles: make(map[int]*os.File),
logDest: logdest,
Type: "ndjson",
OutputDir: "presorted_files",
PartitionN: 256,
JobSize: 1000,
NumWorkers: 4,
PrimaryKeySeparator: "|",
outputFiles: make(map[int]*os.File),
}
}
@ -58,6 +71,7 @@ func (cmd *PreSortCommand) Run(ctx context.Context) error {
if cmd.Table == "" {
return errors.New("must set a table name")
}
cmd.outputLocks = make([]sync.Mutex, cmd.PartitionN)
filelist := []string{}
walkfunc := func(path string, d fs.DirEntry, err error) error {
if !d.IsDir() {
@ -67,7 +81,7 @@ func (cmd *PreSortCommand) Run(ctx context.Context) error {
}
err := filepath.WalkDir(cmd.File, walkfunc)
if err != nil {
return errors.Wrapf(err, "wakling %s", cmd.File)
return errors.Wrapf(err, "walking %s", cmd.File)
}
if len(filelist) > 3 {
cmd.logDest.Printf("%d files: %v...", len(filelist), filelist[:3])
@ -89,58 +103,111 @@ func (cmd *PreSortCommand) Run(ctx context.Context) error {
}
defer func() {
for _, f := range cmd.outputFiles {
_ = f.Sync()
_ = f.Close()
}
}()
linech := make(chan [][]byte, 16)
var eg *errgroup.Group
switch cmd.Type {
case "ndjson":
if err := cmd.RunNDJSON(filelist); err != nil {
return errors.Wrap(err, "running ndjson")
}
eg = cmd.setupNDJSON(linech)
case "csv":
if err := cmd.RunCSV(filelist); err != nil {
return errors.Wrap(err, "running csv")
}
// CSV ends here
return errors.Wrap(cmd.RunCSV(filelist), "running CSV command")
default:
return errors.Errorf("unsupported type %s, try ndjson or csv", cmd.Type)
}
// The rest of this is just for NDJSON... CSV doesn't currently
// support multiple workers. It needs to do the file reading
// differently in order to handle potential newlines within
// fields.
for _, fname := range filelist {
f, err := os.Open(fname)
if err != nil {
return errors.Wrapf(err, "opening %s", fname)
}
sc := bufio.NewScanner(f)
batch:
for {
lines := make([][]byte, 0, cmd.JobSize)
for i := 0; i < cmd.JobSize; i++ {
if !sc.Scan() {
linech <- lines
break batch
}
line := sc.Bytes()
// must copy line as scanner will reuse the buffer it's given us
lc := make([]byte, len(line))
copy(lc, line)
lines = append(lines, lc)
}
linech <- lines
}
if err := sc.Err(); err != nil {
return errors.Wrap(err, "scanning file")
}
}
close(linech)
if err := eg.Wait(); err != nil {
return err
}
return nil
}
func (cmd *PreSortCommand) RunNDJSON(filelist []string) error {
outputEncoders := make(map[int]*json.Encoder)
func (cmd *PreSortCommand) setupNDJSON(linech chan [][]byte) *errgroup.Group {
cmd.outputEncoders = make(map[int]*json.Encoder)
for i, f := range cmd.outputFiles {
outputEncoders[i] = json.NewEncoder(f)
outputEncoders[i].SetIndent("", "")
outputEncoders[i].SetEscapeHTML(false)
cmd.outputEncoders[i] = json.NewEncoder(f)
cmd.outputEncoders[i].SetIndent("", "")
cmd.outputEncoders[i].SetEscapeHTML(false)
}
for _, fname := range filelist {
f, err := os.Open(fname)
if err != nil {
cmd.logDest.Warnf("Could not open %s, skipping", fname)
}
dec := json.NewDecoder(f)
// start workers
eg := &errgroup.Group{}
for i := 0; i < cmd.NumWorkers; i++ {
eg.Go(func() error {
return cmd.ndjsonWorker(linech)
})
}
return eg
}
func (cmd *PreSortCommand) ndjsonWorker(linech chan [][]byte) error {
for lines := range linech {
rec := map[string]interface{}{}
for err = dec.Decode(&rec); err == nil; err = dec.Decode(&rec) {
for _, line := range lines {
err := json.Unmarshal(line, &rec)
if err != nil {
return errors.Wrapf(err, "unmarshaling '%s'", line)
}
partition, err := cmd.ndjsonPartition(rec)
if err != nil {
return err
}
err = outputEncoders[partition].Encode(rec)
if err != nil {
return errors.Wrapf(err, "writing record to partition %d", partition)
if err := cmd.outputNDJSON(partition, rec); err != nil {
return errors.Wrap(err, "outputting ndjson")
}
for k := range rec {
delete(rec, k)
}
}
if err != io.EOF && err != nil {
return errors.Wrapf(err, "error attempting to decode json from %s", fname)
}
}
return nil
}
func (cmd *PreSortCommand) outputNDJSON(partition int, record map[string]interface{}) error {
cmd.outputLocks[partition].Lock()
defer cmd.outputLocks[partition].Unlock()
return cmd.outputEncoders[partition].Encode(record)
}
func (cmd *PreSortCommand) ndjsonPartition(rec map[string]interface{}) (int, error) {
h := fnv.New64a()
_, _ = h.Write([]byte(cmd.Table))
@ -152,12 +219,23 @@ func (cmd *PreSortCommand) ndjsonPartition(rec map[string]interface{}) (int, err
}
h.Write([]byte(toString(val)))
if i < len(cmd.PrimaryKeyFields)-1 {
h.Write([]byte{'|'})
h.Write([]byte(cmd.PrimaryKeySeparator))
}
}
return int(h.Sum64() % uint64(cmd.PartitionN)), nil
}
func toString(val interface{}) string {
switch valt := val.(type) {
case string:
return valt
case float64, int, int64, uint64, float32, uint:
return fmt.Sprintf("%d", valt)
default:
return fmt.Sprintf("%v", valt)
}
}
func (cmd *PreSortCommand) RunCSV(filelist []string) error {
outputWriters := make(map[int]*csv.Writer)
for i, f := range cmd.outputFiles {
@ -209,37 +287,10 @@ func (cmd *PreSortCommand) csvPartition(header map[string]int, rec []string) (in
if !ok {
return 0, errors.Errorf("couldn't find primary key part '%s' in header: %+v", name, header)
}
val := rec[pos]
h.Write([]byte(toString(val)))
h.Write([]byte(rec[pos]))
if i < len(cmd.PrimaryKeyFields)-1 {
h.Write([]byte{'|'})
h.Write([]byte(cmd.PrimaryKeySeparator))
}
}
return int(h.Sum64() % uint64(cmd.PartitionN)), nil
}
func toString(val interface{}) string {
switch valt := val.(type) {
case string:
return valt
case float64, int, int64, uint64, float32, uint:
return fmt.Sprintf("%d", valt)
default:
return fmt.Sprintf("%v", valt)
}
}
func (cmd *PreSortCommand) partition(key string) int {
h := fnv.New64a()
_, _ = h.Write([]byte(cmd.Table))
_, _ = h.Write([]byte(key))
return int(h.Sum64() % uint64(cmd.PartitionN))
}
func partition(key string) int {
indexName := "karambit"
h := fnv.New64a()
_, _ = h.Write([]byte(indexName))
_, _ = h.Write([]byte(key))
return int(h.Sum64() % uint64(256))
}

View file

@ -22,6 +22,7 @@ func TestPreSort(t *testing.T) {
com.File = "sample.ndjson"
com.Table = "blah"
com.PartitionN = 5
com.JobSize = 2
com.PrimaryKeyFields = []string{"url"}
os.Chdir(td)
@ -64,11 +65,12 @@ func TestPreSort(t *testing.T) {
com.File = "sample.csv"
com.Table = "blah"
com.PartitionN = 5
com.JobSize = 2
com.PrimaryKeyFields = []string{"a", "b"}
err = com.Run(context.Background())
if err != nil {
t.Fatalf("running ndjson: %v", err)
t.Fatalf("running csv: %v", err)
}
entries, err = os.ReadDir(filepath.Join(td, com.OutputDir))
@ -120,11 +122,12 @@ func testFile(t *testing.T, name, dir, contents string) {
}
}
var sampleNDJSON string = `{"url":"https://www.yelp.com/search?find_desc=Desserts&find_loc=San+Jose,+CA&start=0","result":{"extractorData":{"url":"https://www.yelp.com/search?find_desc=Desserts&find_loc=San+Jose,+CA&start=0","data":[{"group":[{"Business":[{"href":"https://www.yelp.com/biz/milk-and-wood-san-jose?osq=Desserts","text":"Milk & Wood"}]},{"Business":[{"href":"https://www.yelp.com/biz/dzuis-cakes-and-desserts-san-jose?osq=Desserts","text":"Dzuis Cakes & Desserts"}]},{"Business":[{"href":"https://www.yelp.com/biz/recess-italian-ice-and-desserts-san-jose?osq=Desserts","text":"Recess Italian Ice and Desserts"}]},{"Business":[{"href":"https://www.yelp.com/biz/icicles-san-jose-7?osq=Desserts","text":"ICICLES"}]},{"Business":[{"href":"https://www.yelp.com/biz/sweet-rendezvous-san-jose?osq=Desserts","text":"Sweet Rendezvous"}]},{"Business":[{"href":"https://www.yelp.com/biz/passion-t-snacks-and-desserts-san-jose?osq=Desserts","text":"Passion-T Snacks and Desserts"}]},{"Business":[{"href":"https://www.yelp.com/biz/vampire-penguin-featuring-jastea-san-jose?osq=Desserts","text":"Vampire Penguin featuring Jastea"}]},{"Business":[{"href":"https://www.yelp.com/biz/sweet-gelato-tea-lounge-san-jose?osq=Desserts","text":"Sweet Gelato Tea Lounge"}]},{"Business":[{"href":"https://www.yelp.com/biz/matcha-love-san-jose-6?osq=Desserts","text":"Matcha Love"}]},{"Business":[{"href":"https://www.yelp.com/biz/anton-sv-p%C3%A2tisserie-san-jose-2?osq=Desserts","text":"Anton SV Pâtisserie"}]}]}]},"pageData":{"statusCode":200,"timestamp":1513286383006},"timestamp":1513286383006,"sequenceNumber":0}}
var sampleNDJSON string = `{"url":"https://www.yelp.com/search?find_desc=Desserts&find_loc=San+Jose,+CA&start=0","anything":"somedata","result":{"stuff":"blahhhhhhhhh","blufff":"hahahahah"}}
{"url":"https://www.yelp.com/search?find_desc=Desserts&find_loc=San+Jose,+CA&start=10","result":{"extractorData":{"url":"https://www.yelp.com/search?find_desc=Desserts&find_loc=San+Jose,+CA&start=10","data":[{"group":[{"Business":[{"href":"https://www.yelp.com/biz/oooh-san-jose-4?osq=Desserts","text":"Oooh"}]},{"Business":[{"href":"https://www.yelp.com/biz/hannah-san-jose?osq=Desserts","text":"Hannah"}]},{"Business":[{"href":"https://www.yelp.com/biz/chocatoo-san-jose?osq=Desserts","text":"Chocatoo"}]},{"Business":[{"href":"https://www.yelp.com/biz/nox-cookie-bar-san-jose?osq=Desserts","text":"Nox Cookie Bar"}]},{"Business":[{"href":"https://www.yelp.com/biz/sweet-fix-creamery-san-jose?osq=Desserts","text":"Sweet Fix Creamery"}]},{"Business":[{"href":"https://www.yelp.com/biz/my-milkshake-san-jose?osq=Desserts","text":"My Milkshake"}]},{"Business":[{"href":"https://www.yelp.com/biz/matcha-love-san-jose-6?osq=Desserts","text":"Matcha Love"}]},{"Business":[{"href":"https://www.yelp.com/biz/banana-cr%C3%AApe-san-jose-2?osq=Desserts","text":"Banana Crêpe"}]},{"Business":[{"href":"https://www.yelp.com/biz/marco-polo-italian-ice-cream-san-jose-4?osq=Desserts","text":"Marco Polo Italian Ice Cream"}]},{"Business":[{"href":"https://www.yelp.com/biz/blackball-desserts-san-jose-san-jose?osq=Desserts","text":"BlackBall Desserts San Jose"}]}]}]},"pageData":{"statusCode":200,"timestamp":1513286384917},"timestamp":1513286384917,"sequenceNumber":1}}
{"url":"https://www.yelp.com/search?find_desc=Desserts&find_loc=San+Jose,+CA&start=20","result":{"extractorData":{"url":"https://www.yelp.com/search?find_desc=Desserts&find_loc=San+Jose,+CA&start=20","data":[{"group":[{"Business":[{"href":"https://www.yelp.com/biz/anton-sv-p%C3%A2tisserie-san-jose-2?osq=Desserts","text":"Anton SV Pâtisserie"}]},{"Business":[{"href":"https://www.yelp.com/biz/soyful-desserts-san-jose-8?osq=Desserts","text":"Soyful Desserts"}]},{"Business":[{"href":"https://www.yelp.com/biz/cocola-bakery-san-jose?osq=Desserts","text":"Cocola Bakery"}]},{"Business":[{"href":"https://www.yelp.com/biz/charlies-cheesecake-works-san-jose?osq=Desserts","text":"Charlies Cheesecake Works"}]},{"Business":[{"href":"https://www.yelp.com/biz/jt-express-san-jose-2?osq=Desserts","text":"JT Express"}]},{"Business":[{"href":"https://www.yelp.com/biz/nox-cookie-bar-san-jose?osq=Desserts","text":"Nox Cookie Bar"}]},{"Business":[{"href":"https://www.yelp.com/biz/shuei-do-manju-shop-san-jose?osq=Desserts","text":"Shuei-Do Manju Shop"}]},{"Business":[{"href":"https://www.yelp.com/biz/churros-el-guero-san-jose?osq=Desserts","text":"Churros El Guero"}]},{"Business":[{"href":"https://www.yelp.com/biz/sno-crave-tea-house-san-jose-4?osq=Desserts","text":"Sno-Crave Tea House"}]},{"Business":[{"href":"https://www.yelp.com/biz/j-sweets-san-jose-3?osq=Desserts","text":"J.Sweets"}]}]}]},"pageData":{"statusCode":200,"timestamp":1513286395948},"timestamp":1513286395948,"sequenceNumber":2}}
{"url":"https://www.yelp.com/search?find_desc=Desserts&find_loc=San+Jose,+CA&start=30","result":{"extractorData":{"url":"https://www.yelp.com/search?find_desc=Desserts&find_loc=San+Jose,+CA&start=30","data":[{"group":[{"Business":[{"href":"https://www.yelp.com/biz/treatbot-san-jose-2?osq=Desserts","text":"Treatbot"}]},{"Business":[{"href":"https://www.yelp.com/biz/cream-san-jose?osq=Desserts","text":"CREAM"}]},{"Business":[{"href":"https://www.yelp.com/biz/my-milkshake-san-jose?osq=Desserts","text":"My Milkshake"}]},{"Business":[{"href":"https://www.yelp.com/biz/peters-bakery-san-jose?osq=Desserts","text":"Peters Bakery"}]},{"Business":[{"href":"https://www.yelp.com/biz/the-charming-kitchen-san-jose-5?osq=Desserts","text":"The Charming Kitchen"}]},{"Business":[{"href":"https://www.yelp.com/biz/sweet-fix-creamery-san-jose?osq=Desserts","text":"Sweet Fix Creamery"}]},{"Business":[{"href":"https://www.yelp.com/biz/california-mochi-santa-clara-4?osq=Desserts","text":"California Mochi"}]},{"Business":[{"href":"https://www.yelp.com/biz/raw-sugar-milpitas-2?osq=Desserts","text":"Raw Sugar"}]},{"Business":[{"href":"https://www.yelp.com/biz/chola-desserts-san-jose?osq=Desserts","text":"Chola Desserts"}]},{"Business":[{"href":"https://www.yelp.com/biz/san-jose-tofu-company-san-jose?osq=Desserts","text":"San Jose Tofu Company"}]}]}]},"pageData":{"statusCode":200,"timestamp":1513286386420},"timestamp":1513286386420,"sequenceNumber":3}}
{"url":"https://www.yelp.com/search?find_desc=Desserts&find_loc=San+Jose,+CA&start=30","result":{}}
{"url":"https://www.yelp.com/search?find_desc=Desserts&find_loc=San+Jose,+CA&start=40","result":{"extractorData":{"url":"https://www.yelp.com/search?find_desc=Desserts&find_loc=San+Jose,+CA&start=40","data":[{"group":[{"Business":[{"href":"https://www.yelp.com/biz/the-sweet-corner-san-jose?osq=Desserts","text":"The Sweet Corner"}]},{"Business":[{"href":"https://www.yelp.com/biz/marco-polo-italian-ice-cream-san-jose-4?osq=Desserts","text":"Marco Polo Italian Ice Cream"}]},{"Business":[{"href":"https://www.yelp.com/biz/honeyberry-san-jose-9?osq=Desserts","text":"Honeyberry"}]},{"Business":[{"href":"https://www.yelp.com/biz/my-ch%C3%A8-san-jose?osq=Desserts","text":"My Chè"}]},{"Business":[{"href":"https://www.yelp.com/biz/creme-paris-san-jose-3?osq=Desserts","text":"CreMe Paris"}]},{"Business":[{"href":"https://www.yelp.com/biz/snowflake-san-jose?osq=Desserts","text":"Snowflake"}]},{"Business":[{"href":"https://www.yelp.com/biz/willow-glen-creamery-san-jose-4?osq=Desserts","text":"Willow Glen Creamery"}]},{"Business":[{"href":"https://www.yelp.com/biz/vans-bakery-san-jose?osq=Desserts","text":"Vans Bakery"}]},{"Business":[{"href":"https://www.yelp.com/biz/happiness-cafe-san-jose?osq=Desserts","text":"Happiness Cafe"}]},{"Business":[{"href":"https://www.yelp.com/biz/la-original-paleteria-y-neveria-san-jose?osq=Desserts","text":"La Original Paleteria Y Neveria"}]}]}]},"pageData":{"statusCode":200,"timestamp":1513286387763},"timestamp":1513286387763,"sequenceNumber":4}}
{"url":"https://www.yelp.com/search?find_desc=Desserts&find_loc=San+Jose,+CA&start=0","result":{}}
`
var sampleCSV string = `a,b,c,d,e,f,g
@ -133,4 +136,8 @@ var sampleCSV string = `a,b,c,d,e,f,g
3,4,5,6,7,8,9
4,5,6,7,8,9,0
0,1,2,3,4,5,6
1,3,3,4,5,6,7
4,5,6,7,8,9,0
3,4,5,8,7,8,9
2,3,4,9,6,7,8
`