diff --git a/cmd/pre_sort.go b/cmd/presort.go similarity index 69% rename from cmd/pre_sort.go rename to cmd/presort.go index ae834837a..1bd4aba21 100644 --- a/cmd/pre_sort.go +++ b/cmd/presort.go @@ -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 } diff --git a/ctl/pre_sort.go b/ctl/presort.go similarity index 60% rename from ctl/pre_sort.go rename to ctl/presort.go index 164bde1b2..c90e41b2b 100644 --- a/ctl/pre_sort.go +++ b/ctl/presort.go @@ -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)) -} diff --git a/ctl/pre_sort_test.go b/ctl/presort_test.go similarity index 71% rename from ctl/pre_sort_test.go rename to ctl/presort_test.go index ff070dc4c..ae38e36a2 100644 --- a/ctl/pre_sort_test.go +++ b/ctl/presort_test.go @@ -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":"Dzui’s 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":"Charlie’s 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":"Cre’Me 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":"Van’s 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 `