Repository: segmentio/parquet-go
Pull request #131 (2022-04-09 04:28:37)

add parquet.(*Buffer).ColumnBuffer
This PR adds a `ColumnBuffer` method to the `parquet.Buffer` type. I left a description for the reason why this method is useful in its documentation.



Patch:
diff --git a/README.md b/README.md
index d5e56a0..8aac2cd 100644
--- a/README.md
+++ b/README.md
@@ -450,7 +450,7 @@ of Go values:
 func writeColumns(buffer *parquet.Buffer, columns [3][]interface{}) error {
     values := make([]parquet.Value, len(columns[0]))
     for i := range columns {
-        c := buffer.Column(i).(parquet.ColumnBuffer)
+        c := buffer.ColumnBuffer(i)
         for j, v := range columns[i] {
             values[j] = parquet.ValueOf(v)
         }
@@ -467,10 +467,12 @@ func writeColumns(buffer *parquet.Buffer, ids []int64, values []float32) error {
     if len(ids) != len(values) {
         return fmt.Errorf("number of ids and values mismatch: ids=%d values=%d", len(ids), len(values))
     }
-    if err := buffer.(parquet.Int64Writer).WriteInt64s(ids); err != nil {
+    col0 := buffer.ColumnBuffer(0)
+    col1 := buffer.ColumnBuffer(1)
+    if err := col0.(parquet.Int64Writer).WriteInt64s(ids); err != nil {
         return err
     }
-    if err := buffer.(parquet.FloatWriter).WriteFloats(values); err != nil {
+    if err := col1.(parquet.FloatWriter).WriteFloats(values); err != nil {
         return err
     }
     return nil
diff --git a/buffer.go b/buffer.go
index 4496a44..2304a90 100644
--- a/buffer.go
+++ b/buffer.go
@@ -118,6 +118,16 @@ func (buf *Buffer) NumColumns() int { return len(buf.columns) }
 // The method panics if i is negative or beyond the last column index in buf.
 func (buf *Buffer) Column(i int) ColumnChunk { return buf.columns[i] }
 
+// ColumnBuffer returns the buffer column at index i.
+//
+// This method is similar to Column, but returns a ColumnBuffer type instead of
+// a Column type (the latter being read-only); calling ColumnBuffer or Column
+// with the same index returns the same underlying object, but with different
+// types, which removes the need for making a type assertion if the program
+// needed to write directly to the column buffers. The presence of the Column
+// method is still required to satisfy the RowGroup interface.
+func (buf *Buffer) ColumnBuffer(i int) ColumnBuffer { return buf.columns[i] }
+
 // Schema returns the schema of the buffer.
 //
 // The schema is either configured by passing a Schema in the option list when


Test patch:
diff --git a/buffer_test.go b/buffer_test.go
index d467b94..13ee5ef 100644
--- a/buffer_test.go
+++ b/buffer_test.go
@@ -273,7 +273,7 @@ func testBuffer(t *testing.T, node parquet.Node, reader parquet.ValueReader, buf
 	sortFunc(typ, batch)
 	sort.Sort(buffer)
 
-	page := buffer.Column(0).(parquet.ColumnBuffer).Page()
+	page := buffer.ColumnBuffer(0).Page()
 	numValues := page.NumValues()
 	if numValues != int64(len(batch)) {
 		t.Fatalf("number of values mistmatch: want=%d got=%d", len(batch), numValues)
diff --git a/page_go18_test.go b/page_go18_test.go
index 2b8dc5c..e42382b 100644
--- a/page_go18_test.go
+++ b/page_go18_test.go
@@ -183,7 +183,7 @@ func testPage[T any](t *testing.T, schema *parquet.Schema, test pageTest[T]) {
 
 func testBufferPage[T any](t *testing.T, schema *parquet.Schema, test pageTest[T]) {
 	buffer := parquet.NewBuffer(schema)
-	column := buffer.Column(0).(parquet.ColumnBuffer)
+	column := buffer.ColumnBuffer(0)
 
 	w, err := test.write(column)
 	if err != nil {
@@ -201,7 +201,7 @@ func testBufferPage[T any](t *testing.T, schema *parquet.Schema, test pageTest[T
 
 func testFilePage[T any](t *testing.T, schema *parquet.Schema, test pageTest[T]) {
 	buffer := parquet.NewBuffer(schema)
-	column := buffer.Column(0).(parquet.ColumnBuffer)
+	column := buffer.ColumnBuffer(0)
 
 	w, err := test.write(column)
 	if err != nil {
diff --git a/page_test.go b/page_test.go
index ee256a9..01b2fe1 100644
--- a/page_test.go
+++ b/page_test.go
@@ -331,7 +331,7 @@ func testPage(t *testing.T, schema *parquet.Schema, test pageTest) {
 
 func testBufferPage(t *testing.T, schema *parquet.Schema, test pageTest) {
 	buffer := parquet.NewBuffer(schema)
-	column := buffer.Column(0).(parquet.ColumnBuffer)
+	column := buffer.ColumnBuffer(0)
 
 	w, err := test.write(column)
 	if err != nil {
@@ -349,7 +349,7 @@ func testBufferPage(t *testing.T, schema *parquet.Schema, test pageTest) {
 
 func testFilePage(t *testing.T, schema *parquet.Schema, test pageTest) {
 	buffer := parquet.NewBuffer(schema)
-	column := buffer.Column(0).(parquet.ColumnBuffer)
+	column := buffer.ColumnBuffer(0)
 
 	w, err := test.write(column)
 	if err != nil {