Repository: parquet-go/parquet-go
Pull request #203 (2024-12-24 21:02:29)

reader: Add ability to read back Key Value metadata
This PR adds a GetKeyValueMetadata method to NewGenericReader and NewReader, to be able to read back the key value metadata set whilst writing the file.

saswatamcode: Sorry for the delay on this @achille-roussel! I think exposing File on reader makes sense.
I'll try to update the PR next week!

Patch:
diff --git a/reader.go b/reader.go
index 4d89166..f28cadb 100644
--- a/reader.go
+++ b/reader.go
@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"io"
 	"reflect"
+
+	"github.com/parquet-go/parquet-go/format"
 )
 
 // GenericReader is similar to a Reader but uses a type parameter to define the
@@ -51,6 +53,7 @@ func NewGenericReader[T any](input io.ReaderAt, options ...ReaderOption) *Generi
 	r := &GenericReader[T]{
 		base: Reader{
 			file: reader{
+				file:     f,
 				schema:   c.Schema,
 				rowGroup: rowGroup,
 			},
@@ -134,6 +137,11 @@ func (r *GenericReader[T]) Close() error {
 	return r.base.Close()
 }
 
+// File returns a FileView of the underlying parquet file.
+func (r *GenericReader[T]) File() FileView {
+	return r.base.File()
+}
+
 // readRows reads the next rows from the reader into the given rows slice up to len(rows).
 //
 // The returned values are safe to reuse across readRows calls and do not share
@@ -274,6 +282,7 @@ func NewReader(input io.ReaderAt, options ...ReaderOption) *Reader {
 
 	r := &Reader{
 		file: reader{
+			file:     f,
 			schema:   f.schema,
 			rowGroup: fileRowGroupOf(f),
 		},
@@ -477,6 +486,7 @@ func (r *Reader) Close() error {
 // read rows into Go values, potentially doing partial reads on a subset of the
 // columns due to using a converted row group view.
 type reader struct {
+	file     *File
 	schema   *Schema
 	rowGroup RowGroup
 	rows     Rows
@@ -558,3 +568,90 @@ var (
 	_ RowReader = (*reader)(nil)
 	_ RowSeeker = (*reader)(nil)
 )
+
+type readerFileView struct {
+	reader *reader
+	schema *Schema
+}
+
+// File returns a FileView of the parquet file being read.
+// Only available if Reader was created with a File.
+func (r *Reader) File() FileView {
+	if r.file.schema == nil || r.file.file == nil {
+		return nil
+	}
+	return &readerFileView{
+		&r.file,
+		r.file.schema,
+	}
+}
+
+func (r *readerFileView) Metadata() *format.FileMetaData {
+	if r.reader.file != nil {
+		return r.reader.file.Metadata()
+	}
+	return nil
+}
+
+func (r *readerFileView) Schema() *Schema {
+	return r.schema
+}
+
+func (r *readerFileView) NumRows() int64 {
+	return r.reader.rowGroup.NumRows()
+}
+
+func (r *readerFileView) Lookup(key string) (string, bool) {
+	if meta := r.Metadata(); meta != nil {
+		return lookupKeyValueMetadata(meta.KeyValueMetadata, key)
+	}
+	return "", false
+}
+
+func (r *readerFileView) Size() int64 {
+	if r.reader.file != nil {
+		return r.reader.file.Size()
+	}
+	return 0
+}
+
+func (r *readerFileView) ColumnIndexes() []format.ColumnIndex {
+	if r.reader.file != nil {
+		return r.reader.file.ColumnIndexes()
+	}
+	return nil
+}
+
+func (r *readerFileView) OffsetIndexes() []format.OffsetIndex {
+	if r.reader.file != nil {
+		return r.reader.file.OffsetIndexes()
+	}
+	return nil
+}
+
+func (r *readerFileView) Root() *Column {
+	if meta := r.Metadata(); meta != nil {
+		root, _ := openColumns(nil, meta, r.ColumnIndexes(), r.OffsetIndexes())
+		return root
+	}
+	return nil
+}
+
+func (r *readerFileView) RowGroups() []RowGroup {
+	if meta := r.Metadata(); meta != nil {
+		root := r.Root()
+		columns := make([]*Column, 0, numLeafColumnsOf(root))
+		root.forEachLeaf(func(c *Column) { columns = append(columns, c) })
+
+		fileRowGroups := make([]fileRowGroup, len(meta.RowGroups))
+		for i := range fileRowGroups {
+			fileRowGroups[i].init(nil, r.schema, columns, &meta.RowGroups[i])
+		}
+		rowGroups := make([]RowGroup, len(fileRowGroups))
+		for i := range fileRowGroups {
+			rowGroups[i] = &fileRowGroups[i]
+		}
+		return rowGroups
+	}
+	return nil
+}


Test patch:
diff --git a/reader_test.go b/reader_test.go
index d528087..8c196b6 100644
--- a/reader_test.go
+++ b/reader_test.go
@@ -786,3 +786,80 @@ func TestSeekToRowDictReadMultiplePages(t *testing.T) {
 		t.Fatalf("read != write")
 	}
 }
+
+func TestLookup(t *testing.T) {
+	type rowType struct {
+		Value int32
+	}
+
+	// Create test data with metadata
+	buf := new(bytes.Buffer)
+	schema := parquet.SchemaOf(new(rowType))
+	w := parquet.NewWriter(buf, schema)
+	w.SetKeyValueMetadata("key1", "value1")
+	w.SetKeyValueMetadata("key2", "value2")
+	w.SetKeyValueMetadata("empty", "")
+
+	// Write a sample row
+	if err := w.Write(rowType{Value: 42}); err != nil {
+		t.Fatal(err)
+	}
+	if err := w.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	tests := []struct {
+		name     string
+		key      string
+		expected string
+	}{
+		{
+			name:     "existing key 1",
+			key:      "key1",
+			expected: "value1",
+		},
+		{
+			name:     "existing key 2",
+			key:      "key2",
+			expected: "value2",
+		},
+		{
+			name:     "empty value",
+			key:      "empty",
+			expected: "",
+		},
+		{
+			name:     "non-existent key",
+			key:      "missing",
+			expected: "",
+		},
+	}
+
+	t.Run("Reader", func(t *testing.T) {
+		r := parquet.NewReader(bytes.NewReader(buf.Bytes()))
+		defer r.Close()
+
+		for _, tt := range tests {
+			t.Run(tt.name, func(t *testing.T) {
+				got, _ := r.File().Lookup(tt.key)
+				if got != tt.expected {
+					t.Errorf("Lookup(%q) = %q, want %q", tt.key, got, tt.expected)
+				}
+			})
+		}
+	})
+
+	t.Run("GenericReader", func(t *testing.T) {
+		r := parquet.NewGenericReader[rowType](bytes.NewReader(buf.Bytes()))
+		defer r.Close()
+
+		for _, tt := range tests {
+			t.Run(tt.name, func(t *testing.T) {
+				got, _ := r.File().Lookup(tt.key)
+				if got != tt.expected {
+					t.Errorf("Lookup(%q) = %q, want %q", tt.key, got, tt.expected)
+				}
+			})
+		}
+	})
+}