diff --git a/cmd/dataframe-csv-loader.go b/cmd/dataframe-csv-loader.go index a0f677ca5..fd0679b1d 100644 --- a/cmd/dataframe-csv-loader.go +++ b/cmd/dataframe-csv-loader.go @@ -24,6 +24,7 @@ func newDataframeCsvLoaderCommand(logdest logger.Logger) *cobra.Command { flags.StringVar(&cmd.AuthToken, "auth-token", "", "Authentication token") flags.StringVar(&cmd.Index, "index", "", "Destination Index. ") flags.IntVar(&cmd.MaxCapacity, "buffer", 0, "Maximum size of of the line buffer defaults to go bufio default ") + flags.IntVar(&cmd.BatchSize, "batch-size", 1048576, "Maximum number of records to send in a single batch ") ctl.SetTLSConfig( flags, "", &cmd.TLS.CertificatePath, diff --git a/ctl/dataframe-csv-loader.go b/ctl/dataframe-csv-loader.go index 70cc45e3e..2a5f893dc 100644 --- a/ctl/dataframe-csv-loader.go +++ b/ctl/dataframe-csv-loader.go @@ -50,6 +50,9 @@ type DataframeCsvLoaderCommand struct { // max line length of csv file MaxCapacity int + // Batch Size + BatchSize int + // Host:port on which to listen for pprof. Pprof string `json:"pprof"` @@ -234,7 +237,9 @@ func (cmd *DataframeCsvLoaderCommand) Run(ctx context.Context) (err error) { fileScanner.Scan() // skip the header cmd.Logger().Infof("Build the dataframe input package in memory") id := uint64(0) + recordCounter := 0 for fileScanner.Scan() { + recordCounter++ records := strings.Split(fileScanner.Text(), ",") if cmd.needTranslation { id = lookup[records[0]] @@ -278,13 +283,23 @@ func (cmd *DataframeCsvLoaderCommand) Run(ctx context.Context) (err error) { } } } + if recordCounter > cmd.BatchSize { + cmd.Logger().Infof("sending package to featurebase") + err = sharder.Store(arrowSchema, cmd.client) + if err != nil { + return err + } + sharder.Reset() + recordCounter = 0 + } } - cmd.Logger().Infof("sending package to featurebase") - err = sharder.Store(arrowSchema, cmd.client) - if err != nil { - return err + if recordCounter > 0 { + err = sharder.Store(arrowSchema, cmd.client) + if err != nil { + return err + } } - return err + return nil } type pair struct { @@ -378,6 +393,10 @@ type Sharder struct { log logger.Logger } +func (s *Sharder) Reset() { + s.shards = make(map[uint64]*ShardDiff) +} + func (s *Sharder) GetShard(shard uint64) (*ShardDiff, error) { f, ok := s.shards[shard] if ok {