change Writer.flush parameter to bool

brimdata · Dec 8, 2022 · a7a90e7 · a7a90e7
1 parent 7bb8b53
commit a7a90e7
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/zio/arrowio/writer.go b/zio/arrowio/writer.go
@@ -47,7 +47,7 @@ func NewWriter(w io.WriteCloser) *Writer {
 func (w *Writer) Close() error {
 	var err error
 	if w.writer != nil {
-		err = w.flush(1)
+		err = w.flush(true)
 		w.builder.Release()
 		if err2 := w.writer.Close(); err == nil {
 			err = err2
@@ -88,11 +88,11 @@ func (w *Writer) Write(val *zed.Value) error {
 		}
 		w.buildArrowValue(builder, recType.Columns[i].Type, b)
 	}
-	return w.flush(recordBatchSize)
+	return w.flush(false)
 }
 
-func (w *Writer) flush(min int) error {
-	if w.builder.Field(0).Len() < min {
+func (w *Writer) flush(partialBatch bool) error {
+	if n := w.builder.Field(0).Len(); n >= recordBatchSize || partialBatch && n > 0 {
 		return nil
 	}
 	rec := w.builder.NewRecord()