add a batch size to limit upload payloads (#2275)

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,

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 {