[apache#35775] parquet.FileWriter: persist changes to KeyValueMetadat…

…a on close
xxgreg · May 26, 2023 · 8ef8820 · 8ef8820
1 parent 2d32efe
commit 8ef8820
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 10 deletions.
diff --git a/go/parquet/file/file_writer.go b/go/parquet/file/file_writer.go
@@ -170,7 +170,7 @@ func (fw *Writer) Close() error {
 		fileEncryptProps := fw.props.FileEncryptionProperties()
 		if fileEncryptProps == nil { // non encrypted file
 			var err error
-			if fw.FileMetadata, err = fw.metadata.Finish(); err != nil {
+			if fw.FileMetadata, err = fw.metadata.Finish(fw.KeyValueMetadata); err != nil {
 				return err
 			}
 
@@ -186,7 +186,7 @@ func (fw *Writer) Close() error {
 func (fw *Writer) closeEncryptedFile(props *parquet.FileEncryptionProperties) (err error) {
 	// encrypted file with encrypted footer
 	if props.EncryptedFooter() {
-		fw.FileMetadata, err = fw.metadata.Finish()
+		fw.FileMetadata, err = fw.metadata.Finish(fw.KeyValueMetadata)
 		if err != nil {
 			return
 		}
@@ -214,7 +214,7 @@ func (fw *Writer) closeEncryptedFile(props *parquet.FileEncryptionProperties) (e
 			return err
 		}
 	} else {
-		if fw.FileMetadata, err = fw.metadata.Finish(); err != nil {
+		if fw.FileMetadata, err = fw.metadata.Finish(fw.KeyValueMetadata); err != nil {
 			return
 		}
 		footerSigningEncryptor := fw.fileEncryptor.GetFooterSigningEncryptor()

diff --git a/go/parquet/metadata/file.go b/go/parquet/metadata/file.go
@@ -45,7 +45,6 @@ type FileMetaDataBuilder struct {
 	schema         *schema.Schema
 	rowGroups      []*format.RowGroup
 	currentRgBldr  *RowGroupMetaDataBuilder
-	kvmeta         KeyValueMetadata
 	cryptoMetadata *format.FileCryptoMetaData
 }
 
@@ -60,7 +59,6 @@ func NewFileMetadataBuilder(schema *schema.Schema, props *parquet.WriterProperti
 		metadata:       format.NewFileMetaData(),
 		props:          props,
 		schema:         schema,
-		kvmeta:         kvmeta,
 		cryptoMetadata: crypto,
 	}
 }
@@ -98,7 +96,7 @@ func (f *FileMetaDataBuilder) AppendRowGroup() *RowGroupMetaDataBuilder {
 // Finish will finalize the metadata of the number of rows, row groups,
 // version etc. This will clear out this filemetadatabuilder so it can
 // be re-used
-func (f *FileMetaDataBuilder) Finish() (*FileMetaData, error) {
+func (f *FileMetaDataBuilder) Finish(kvmeta KeyValueMetadata) (*FileMetaData, error) {
 	totalRows := int64(0)
 	for _, rg := range f.rowGroups {
 		totalRows += rg.NumRows
@@ -145,7 +143,7 @@ func (f *FileMetaDataBuilder) Finish() (*FileMetaData, error) {
 	}
 
 	f.metadata.Schema = schema.ToThrift(f.schema.Root())
-	f.metadata.KeyValueMetadata = f.kvmeta
+	f.metadata.KeyValueMetadata = kvmeta
 
 	out := &FileMetaData{
 		FileMetaData: f.metadata,

diff --git a/go/parquet/metadata/metadata_test.go b/go/parquet/metadata/metadata_test.go
@@ -63,7 +63,7 @@ func generateTableMetaData(schema *schema.Schema, props *parquet.WriterPropertie
 	rg2Builder.SetNumRows(nrows / 2)
 	rg2Builder.Finish(1024, -1)
 
-	return fbuilder.Finish()
+	return fbuilder.Finish(metadata.KeyValueMetadata{})
 }
 
 func assertStatsSet(t *testing.T, m *metadata.ColumnChunkMetaData) {
@@ -247,7 +247,7 @@ func TestV1VersionMetadata(t *testing.T) {
 	schema := schema.NewSchema(root)
 
 	fbuilder := metadata.NewFileMetadataBuilder(schema, props, nil)
-	faccessor, err := fbuilder.Finish()
+	faccessor, err := fbuilder.Finish(metadata.KeyValueMetadata{})
 	require.NoError(t, err)
 	assert.EqualValues(t, parquet.V1_0, faccessor.Version())
 }
@@ -266,7 +266,7 @@ func TestKeyValueMetadata(t *testing.T) {
 	kvmeta.Append("test_key", "test_value")
 
 	fbuilder := metadata.NewFileMetadataBuilder(schema, props, kvmeta)
-	faccessor, err := fbuilder.Finish()
+	faccessor, err := fbuilder.Finish(kvmeta)
 	require.NoError(t, err)
 
 	assert.True(t, faccessor.KeyValueMetadata().Equals(kvmeta))