apacheGH-39921: ColumnWriter not reset TotalCompressedBytes after Flush

mapleFU · Feb 2, 2024 · 69cdfff · 69cdfff
1 parent 129a529
commit 69cdfff
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 2 deletions.
diff --git a/go/parquet/file/column_writer.go b/go/parquet/file/column_writer.go
@@ -397,7 +397,6 @@ func (w *columnWriter) FlushBufferedDataPages() (err error) {
 		}
 	}
 	w.pages = w.pages[:0]
-	w.totalCompressedBytes = 0
 	return
 }
 
@@ -542,7 +541,9 @@ func (w *columnWriter) Close() (err error) {
 	if !w.closed {
 		w.closed = true
 		if w.hasDict && !w.fallbackToNonDict {
-			w.WriteDictionaryPage()
+			if err = w.WriteDictionaryPage(); err != nil {
+				return err
+			}
 		}
 
 		if err = w.FlushBufferedDataPages(); err != nil {

diff --git a/go/parquet/file/column_writer_test.go b/go/parquet/file/column_writer_test.go
@@ -426,6 +426,26 @@ func (p *PrimitiveWriterTestSuite) testDictionaryFallbackEncoding(version parque
 	}
 }
 
+func (p *PrimitiveWriterTestSuite) testDictionaryFallbackAndCompressedSize(version parquet.Version) {
+	p.GenerateData(SmallSize)
+	props := parquet.DefaultColumnProperties()
+	props.DictionaryEnabled = true
+
+	if version == parquet.V1_0 {
+		props.Encoding = parquet.Encodings.PlainDict
+	} else {
+		props.Encoding = parquet.Encodings.RLEDict
+	}
+
+	writer := p.buildWriter(SmallSize, props, parquet.WithVersion(version))
+	p.WriteBatchValues(writer, nil, nil)
+	writer.FallbackToPlain()
+	p.NotEqual(0, writer.TotalCompressedBytes())
+	writer.Close()
+	p.NotEqual(0, writer.TotalCompressedBytes())
+	p.NotEqual(0, writer.TotalBytesWritten())
+}
+
 func (p *PrimitiveWriterTestSuite) TestRequiredPlain() {
 	p.testRequiredWithEncoding(parquet.Encodings.Plain)
 }
@@ -575,6 +595,14 @@ func (p *PrimitiveWriterTestSuite) TestDictionaryFallbackEncodingV2() {
 	p.testDictionaryFallbackEncoding(parquet.V2_LATEST)
 }
 
+func (p *PrimitiveWriterTestSuite) TestDictionaryFallbackStatsV1() {
+	p.testDictionaryFallbackAndCompressedSize(parquet.V1_0)
+}
+
+func (p *PrimitiveWriterTestSuite) TestDictionaryFallbackStatsV2() {
+	p.testDictionaryFallbackAndCompressedSize(parquet.V2_LATEST)
+}
+
 func (p *PrimitiveWriterTestSuite) TestOptionalNullValueChunk() {
 	// test case for NULL values
 	p.SetupSchema(parquet.Repetitions.Optional, 1)