GH-38516: [Go][Parquet] Increment the number of rows written when app…

…ending a new row group (#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: #38516 Authored-by: Tim Schaub <[email protected]> Signed-off-by: Matt Topol <[email protected]>
apache · Nov 13, 2023 · 1ff43ab · 1ff43ab
1 parent e18ace3
commit 1ff43ab
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())