From 1ff43ab5ee13de5c3130acf10c7aa8eb9680baab Mon Sep 17 00:00:00 2001 From: Tim Schaub Date: Mon, 13 Nov 2023 10:01:01 -0700 Subject: [PATCH] GH-38516: [Go][Parquet] Increment the number of rows written when appending 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 Signed-off-by: Matt Topol --- go/parquet/file/file_writer.go | 1 + go/parquet/file/file_writer_test.go | 2 ++ 2 files changed, 3 insertions(+) diff --git a/go/parquet/file/file_writer.go b/go/parquet/file/file_writer.go index 48eb8f44d167b..1d7f7840dac50 100644 --- 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 index beceff9732584..f32e403a8d534 100644 --- 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())