apacheGH-38516: [Go][Parquet] Increment the number of rows written wh…

…en appending a new row group (apache#38517) ### Rationale for this change This makes it so the `NumRows` method on the `file.Writer` reports the total number of rows written across multiple row groups. ### Are these changes tested? A regression test is added that asserts that the total number of rows written matches expectations. * Closes: apache#38516 Authored-by: Tim Schaub <[email protected]> Signed-off-by: Matt Topol <[email protected]>
dgreiss · Feb 17, 2024 · e46408a · e46408a
1 parent 82b02c7
commit e46408a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/go/parquet/file/file_writer.go b/go/parquet/file/file_writer.go
@@ -121,6 +121,7 @@ func (fw *Writer) AppendRowGroup() SerialRowGroupWriter {
 
 func (fw *Writer) appendRowGroup(buffered bool) *rowGroupWriter {
 	if fw.rowGroupWriter != nil {
+		fw.nrows += fw.rowGroupWriter.nrows
 		fw.rowGroupWriter.Close()
 	}
 	fw.rowGroups++

diff --git a/go/parquet/file/file_writer_test.go b/go/parquet/file/file_writer_test.go
@@ -97,6 +97,8 @@ func (t *SerializeTestSuite) fileSerializeTest(codec compress.Compression, expec
 	writer.Close()
 
 	nrows := t.numRowGroups * t.rowsPerRG
+	t.EqualValues(nrows, writer.NumRows())
+
 	reader, err := file.NewParquetReader(bytes.NewReader(sink.Bytes()))
 	t.NoError(err)
 	t.Equal(t.numCols, reader.MetaData().Schema.NumColumns())