Repository: parquet-go/parquet-go (go)

Issue:
Enhancement: add Writer.WriteColumns([][]Value)
My application converts data from one format into Parquet. I don't have static Go types as the shape of input data is dynamic. So the process it to first convert the input schema to a Parquet schema, and then to convert input data -- which is not columnar -- into a Parquet data.

That translation step mains a `[][]Value` as columns (so even though structurally the same as `[]Row`, the data therein is basically transposed). As I process a record, I append values to the relevant column in this `[][]Value`.

Currently, when I need to write a batch of values to the Parquet file, I have to _transpose_ these values and produce a `[]Row`. But the very first thing that `Writer` (and `GenericWriter`) do is to transpose it back into columns. So it feels like a waste to maintain the transpose code and definitely a waste of CPU cycles (and potentially allocating slices) to transpose and then immediately un-transpose.

If the `Writer` and `GenericWriter` types exposed a `WriteColumns([][]Value)` method, then I could directly pass the data as my application has computed it, no transposing and un-transposing.

I am happy to contribute such a change if it would be accepted.

Pull request description:
Writer.ColumnWriters now returns concrete ColumnWriter type instead of ValueWriter
I know this isn't perfectly backwards-compatible since it changes the return type of a method. But it's a very-newly-introduced method and this repo is still pre-v1.0, so this seemed okay.

This basically expands the interface of what you get back from `Writer.ColumnWriters()` so that you can actually flush pages from the buffer or have the writer auto-flush based on the page size (by using `WriteRows` instead of the lower level `WriteValues`).

This adds a new test that verifies equivalence between using `Writer.WriteRows`, `ColumnWriter.WriteRows`, and `ColumnWriter.WriteValues`.

This addresses the [latest observations/comment](https://github.com/parquet-go/parquet-go/issues/241#issuecomment-2682237551) on #241.

Closes #241.

Patch:
diff --git a/writer.go b/writer.go
index 88f7640..aca06c0 100644
--- a/writer.go
+++ b/writer.go
@@ -221,7 +221,7 @@ func (w *GenericWriter[T]) Schema() *Schema {
 	return w.base.Schema()
 }
 
-func (w *GenericWriter[T]) ColumnWriters() []ValueWriter {
+func (w *GenericWriter[T]) ColumnWriters() []*ColumnWriter {
 	return w.base.ColumnWriters()
 }
 
@@ -491,7 +491,7 @@ func (w *Writer) SetKeyValueMetadata(key, value string) {
 // ColumnWriters returns writers for each column. This allows applications to
 // write values directly to each column instead of having to first assemble
 // values into rows to use WriteRows.
-func (w *Writer) ColumnWriters() []ValueWriter { return w.writer.valueWriters }
+func (w *Writer) ColumnWriters() []*ColumnWriter { return w.writer.columns }
 
 type writerFileView struct {
 	writer *writer
@@ -572,11 +572,10 @@ type writer struct {
 	createdBy string
 	metadata  []format.KeyValue
 
-	columns      []*writerColumn
-	valueWriters []ValueWriter
-	columnChunk  []format.ColumnChunk
-	columnIndex  []format.ColumnIndex
-	offsetIndex  []format.OffsetIndex
+	columns     []*ColumnWriter
+	columnChunk []format.ColumnChunk
+	columnIndex []format.ColumnIndex
+	offsetIndex []format.OffsetIndex
 
 	columnOrders   []format.ColumnOrder
 	schemaElements []format.SchemaElement
@@ -677,7 +676,7 @@ func newWriter(output io.Writer, config *WriterConfig) *writer {
 			columnType = dictionary.Type()
 		}
 
-		c := &writerColumn{
+		c := &ColumnWriter{
 			buffers:            buffers,
 			pool:               config.ColumnPageBuffers,
 			columnPath:         leaf.path,
@@ -764,10 +763,6 @@ func newWriter(output io.Writer, config *WriterConfig) *writer {
 	for i, c := range w.columns {
 		w.columnOrders[i] = *c.columnType.ColumnOrder()
 	}
-	w.valueWriters = make([]ValueWriter, len(w.columns))
-	for i, c := range w.columns {
-		w.valueWriters[i] = c
-	}
 
 	return w
 }
@@ -1067,7 +1062,7 @@ func (w *writer) WriteRows(rows []Row) (int, error) {
 
 		for i, values := range w.values {
 			if len(values) > 0 {
-				if err := w.columns[i].writeRows(values); err != nil {
+				if _, err := w.columns[i].WriteRowValues(values); err != nil {
 					return 0, err
 				}
 			}
@@ -1123,7 +1118,7 @@ func (w *writer) writeRows(numRows int, write func(i, j int) (int, error)) (int,
 // The WriteValues method is intended to work in pair with WritePage to allow
 // programs to target writing values to specific columns of of the writer.
 func (w *writer) WriteValues(values []Value) (numValues int, err error) {
-	return w.columns[values[0].Column()].WriteValues(values)
+	return w.columns[values[0].Column()].writeValues(values)
 }
 
 // One writerBuffers is used by each writer instance, the memory buffers here
@@ -1221,7 +1216,8 @@ func (wb *writerBuffers) swapPageAndScratchBuffers() {
 	wb.page, wb.scratch = wb.scratch, wb.page[:0]
 }
 
-type writerColumn struct {
+// ColumnWriter writes values for a single column to underlying medium.
+type ColumnWriter struct {
 	pool       BufferPool
 	pageBuffer io.ReadWriteSeeker
 	numPages   int
@@ -1259,7 +1255,7 @@ type writerColumn struct {
 	offsetIndex *format.OffsetIndex
 }
 
-func (c *writerColumn) reset() {
+func (c *ColumnWriter) reset() {
 	if c.columnBuffer != nil {
 		c.columnBuffer.Reset()
 	}
@@ -1291,7 +1287,7 @@ func (c *writerColumn) reset() {
 	c.offsetIndex.PageLocations = c.offsetIndex.PageLocations[:0]
 }
 
-func (c *writerColumn) totalRowCount() int64 {
+func (c *ColumnWriter) totalRowCount() int64 {
 	n := c.numRows
 	if c.columnBuffer != nil {
 		n += int64(c.columnBuffer.Len())
@@ -1299,7 +1295,10 @@ func (c *writerColumn) totalRowCount() int64 {
 	return n
 }
 
-func (c *writerColumn) flush() (err error) {
+func (c *ColumnWriter) flush() (err error) {
+	if c.columnBuffer == nil {
+		return nil
+	}
 	if c.columnBuffer.Len() > 0 {
 		defer c.columnBuffer.Reset()
 		_, err = c.writeDataPage(c.columnBuffer.Page())
@@ -1307,7 +1306,7 @@ func (c *writerColumn) flush() (err error) {
 	return err
 }
 
-func (c *writerColumn) flushFilterPages() (err error) {
+func (c *ColumnWriter) flushFilterPages() (err error) {
 	if c.columnFilter == nil {
 		return nil
 	}
@@ -1409,7 +1408,7 @@ func (c *writerColumn) flushFilterPages() (err error) {
 	return nil
 }
 
-func (c *writerColumn) resizeBloomFilter(numValues int64) {
+func (c *ColumnWriter) resizeBloomFilter(numValues int64) {
 	filterSize := c.columnFilter.Size(numValues)
 	if cap(c.filter) < filterSize {
 		c.filter = make([]byte, filterSize)
@@ -1421,7 +1420,7 @@ func (c *writerColumn) resizeBloomFilter(numValues int64) {
 	}
 }
 
-func (c *writerColumn) newColumnBuffer() ColumnBuffer {
+func (c *ColumnWriter) newColumnBuffer() ColumnBuffer {
 	column := c.columnType.NewColumnBuffer(int(c.bufferIndex), c.columnType.EstimateNumValues(int(c.bufferSize)))
 	switch {
 	case c.maxRepetitionLevel > 0:
@@ -1432,29 +1431,41 @@ func (c *writerColumn) newColumnBuffer() ColumnBuffer {
 	return column
 }
 
-func (c *writerColumn) writeRows(rows []Value) error {
+// WriteRowValues writes entire rows to the column. On success, this returns the
+// number of rows written (not the number of values).
+//
+// Unlike ValueWriter, where arbitrary values may be written regardless of row
+// boundaries, this method requires whole rows. This is because the written
+// values may be automatically flushed to a data page, based on the writer's
+// configured page buffer size, and a single row is not permitted to span two
+// pages.
+func (c *ColumnWriter) WriteRowValues(rows []Value) (int, error) {
+	var startingRows int64
 	if c.columnBuffer == nil {
 		// Lazily create the row group column so we don't need to allocate it if
 		// rows are not written individually to the column.
 		c.columnBuffer = c.newColumnBuffer()
+	} else {
+		startingRows = int64(c.columnBuffer.Len())
 	}
 	if _, err := c.columnBuffer.WriteValues(rows); err != nil {
-		return err
+		return 0, err
 	}
+	numRows := int(int64(c.columnBuffer.Len()) - startingRows)
 	if c.columnBuffer.Size() >= int64(c.bufferSize) {
-		return c.flush()
+		return numRows, c.flush()
 	}
-	return nil
+	return numRows, nil
 }
 
-func (c *writerColumn) WriteValues(values []Value) (numValues int, err error) {
+func (c *ColumnWriter) writeValues(values []Value) (numValues int, err error) {
 	if c.columnBuffer == nil {
 		c.columnBuffer = c.newColumnBuffer()
 	}
 	return c.columnBuffer.WriteValues(values)
 }
 
-func (c *writerColumn) writeBloomFilter(w io.Writer) error {
+func (c *ColumnWriter) writeBloomFilter(w io.Writer) error {
 	e := thrift.NewEncoder(c.header.protocol.NewWriter(w))
 	h := bloomFilterHeader(c.columnFilter)
 	h.NumBytes = int32(len(c.filter))
@@ -1465,7 +1476,7 @@ func (c *writerColumn) writeBloomFilter(w io.Writer) error {
 	return err
 }
 
-func (c *writerColumn) writeDataPage(page Page) (int64, error) {
+func (c *ColumnWriter) writeDataPage(page Page) (int64, error) {
 	numValues := page.NumValues()
 	if numValues == 0 {
 		return 0, nil
@@ -1578,7 +1589,7 @@ func (c *writerColumn) writeDataPage(page Page) (int64, error) {
 	return numValues, nil
 }
 
-func (c *writerColumn) writeDictionaryPage(output io.Writer, dict Dictionary) (err error) {
+func (c *ColumnWriter) writeDictionaryPage(output io.Writer, dict Dictionary) (err error) {
 	buf := c.buffers
 	buf.reset()
 
@@ -1623,14 +1634,14 @@ func (c *writerColumn) writeDictionaryPage(output io.Writer, dict Dictionary) (e
 	return nil
 }
 
-func (w *writerColumn) writePageToFilter(page Page) (err error) {
+func (w *ColumnWriter) writePageToFilter(page Page) (err error) {
 	pageType := page.Type()
 	pageData := page.Data()
 	w.filter, err = pageType.Encode(w.filter, pageData, w.columnFilter.Encoding())
 	return err
 }
 
-func (c *writerColumn) writePageTo(size int64, writeTo func(io.Writer) (int64, error)) (err error) {
+func (c *ColumnWriter) writePageTo(size int64, writeTo func(io.Writer) (int64, error)) (err error) {
 	if c.pageBuffer == nil {
 		c.pageBuffer = c.pool.GetBuffer()
 		defer func() {
@@ -1654,7 +1665,7 @@ func (c *writerColumn) writePageTo(size int64, writeTo func(io.Writer) (int64, e
 	return nil
 }
 
-func (c *writerColumn) makePageStatistics(page Page) format.Statistics {
+func (c *ColumnWriter) makePageStatistics(page Page) format.Statistics {
 	numNulls := page.NumNulls()
 	minValue, maxValue, _ := page.Bounds()
 	minValueBytes := minValue.Bytes()
@@ -1668,7 +1679,7 @@ func (c *writerColumn) makePageStatistics(page Page) format.Statistics {
 	}
 }
 
-func (c *writerColumn) recordPageStats(headerSize int32, header *format.PageHeader, page Page) {
+func (c *ColumnWriter) recordPageStats(headerSize int32, header *format.PageHeader, page Page) {
 	uncompressedSize := headerSize + header.UncompressedPageSize
 	compressedSize := headerSize + header.CompressedPageSize
 
@@ -1815,8 +1826,6 @@ var (
 	_ RowWriter   = (*writer)(nil)
 	_ ValueWriter = (*writer)(nil)
 
-	_ ValueWriter = (*writerColumn)(nil)
-
 	_ io.ReaderFrom   = (*offsetTrackingWriter)(nil)
 	_ io.StringWriter = (*offsetTrackingWriter)(nil)
 )


Test patch:
diff --git a/writer_test.go b/writer_test.go
index 16509aa..ff2292b 100644
--- a/writer_test.go
+++ b/writer_test.go
@@ -2,6 +2,7 @@ package parquet_test
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"io"
 	"math/rand"
@@ -19,6 +20,7 @@ import (
 
 	"github.com/parquet-go/parquet-go"
 	"github.com/parquet-go/parquet-go/compress"
+	"github.com/parquet-go/parquet-go/internal/unsafecast"
 )
 
 const (
@@ -100,6 +102,156 @@ func benchmarkGenericWriter[Row generator[Row]](b *testing.B) {
 	})
 }
 
+func TestWriteRowsColumns(t *testing.T) {
+	type simpleFlat struct {
+		A, B, C int64
+		S, T, U string
+	}
+	t.Run("simple", func(t *testing.T) {
+		testWriteRowsColumns[simpleFlat](t)
+	})
+	type complexNested struct {
+		A  []int64
+		S  []string
+		M  map[string]string
+		R1 []struct {
+			B  []int64
+			T  []string
+			R2 []simpleFlat
+		}
+	}
+	t.Run("complex", func(t *testing.T) {
+		testWriteRowsColumns[complexNested](t)
+	})
+}
+
+func testWriteRowsColumns[T any](t *testing.T) {
+	var zero T
+	const numRows = 1000
+	data := rowsOf(numRows, zero)
+	rowSlice := asRowSlice[T](t, data)
+	colSlice := transpose(rowSlice)
+	testCases := []struct {
+		name  string
+		write func(*testing.T, *parquet.GenericWriter[T], [][]parquet.Value) (int, error)
+		data  [][]parquet.Value
+	}{
+		{
+			name: "WriteRows",
+			write: func(_ *testing.T, w *parquet.GenericWriter[T], vals [][]parquet.Value) (int, error) {
+				// annoying that we can't use normal type conversions between
+				// []Row and [][]Value since they are structurally identical
+				rows := unsafecast.Slice[parquet.Row](vals)
+				return w.WriteRows(rows)
+			},
+			data: unsafecast.Slice[[]parquet.Value](rowSlice),
+		},
+		{
+			name: "ColumnWriter.WriteRowValues",
+			write: func(t *testing.T, w *parquet.GenericWriter[T], vals [][]parquet.Value) (int, error) {
+				t.Helper()
+				var numRows int
+				for i, col := range vals {
+					num, err := w.ColumnWriters()[i].WriteRowValues(col)
+					if err != nil {
+						return 0, err
+					}
+					if i == 0 {
+						numRows = num
+					} else if numRows != num {
+						t.Errorf("column %d disagrees with number of rows written in column %d", i, 0)
+					}
+				}
+				return numRows, nil
+			},
+			data: colSlice,
+		},
+	}
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			var buf bytes.Buffer
+			w := parquet.NewGenericWriter[T](&buf)
+			n, err := testCase.write(t, w, testCase.data)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if n != numRows {
+				t.Errorf("wrote %d rows, but expected to write %d", n, numRows)
+			}
+			if err := w.Close(); err != nil {
+				t.Fatal(err)
+			}
+			// Now read the data back and make sure it matches the input data.
+			r := parquet.NewGenericReader[T](bytes.NewReader(buf.Bytes()))
+			roundTripped := readAllRows(t, r, numRows)
+			assertRowsEqualByRow(t, roundTripped, rowSlice)
+		})
+	}
+}
+
+func asRowSlice[T any](t testing.TB, data rows) []parquet.Row {
+	t.Helper()
+
+	typedRows := make([]T, len(data))
+	for i := range data {
+		typedRows[i] = data[i].(T)
+	}
+
+	var buf bytes.Buffer
+	w := parquet.NewGenericWriter[T](&buf)
+	n, err := w.Write(typedRows)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n != len(typedRows) {
+		t.Errorf("wrote %d rows, but expected to write %d", n, len(typedRows))
+	}
+	if err := w.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	r := parquet.NewGenericReader[T](bytes.NewReader(buf.Bytes()))
+	return readAllRows(t, r, len(typedRows))
+}
+
+func readAllRows[T any](t testing.TB, r *parquet.GenericReader[T], numRows int) []parquet.Row {
+	t.Helper()
+	if r.NumRows() != int64(numRows) {
+		t.Errorf("reader reports %d rows, but expected %d", r.NumRows(), numRows)
+	}
+	rows := make([]parquet.Row, numRows)
+	n, err := r.ReadRows(rows)
+	if n != numRows {
+		t.Errorf("wrote %d rows, but expected to write %d", n, numRows)
+	}
+	if err != nil && !errors.Is(err, io.EOF) {
+		t.Fatal(err)
+	}
+	if err := r.Close(); err != nil {
+		t.Fatal(err)
+	}
+	return rows
+}
+
+func transpose(rows []parquet.Row) [][]parquet.Value {
+	var cols [][]parquet.Value
+	for i, row := range rows {
+		if i == 0 {
+			var columnCount int
+			row.Range(func(columnIndex int, columnValues []parquet.Value) bool {
+				columnCount++
+				return true
+			})
+			cols = make([][]parquet.Value, columnCount)
+		}
+		row.Range(func(columnIndex int, columnValues []parquet.Value) bool {
+			cols[columnIndex] = append(cols[columnIndex], columnValues...)
+			return true
+		})
+	}
+	return cols
+}
+
 func TestIssueSegmentio272(t *testing.T) {
 	type T2 struct {
 		X string `parquet:",dict,optional"`