feat: support conversion of chunked arrays

Author: ahmed-mez

parquet/pqarrow/column_readers.go

@@ -319,7 +319,7 @@ func (sr *structReader) BuildArray(lenBound int64) (*arrow.Chunked, error) {
 			return nil, err
 		}
 
-		childArrData[i], err = chunksToSingle(field)
+		childArrData[i], err = chunksToSingle(field, sr.rctx.mem)
 		field.Release() // release field before checking
 		if err != nil {
 			return nil, err
@@ -442,7 +442,7 @@ func (lr *listReader) BuildArray(lenBound int64) (*arrow.Chunked, error) {
 		validityBuffer.Resize(int(bitutil.BytesForBits(validityIO.Read)))
 	}
 
-	item, err := chunksToSingle(arr)
+	item, err := chunksToSingle(arr, lr.rctx.mem)
 	if err != nil {
 		return nil, err
 	}
@@ -489,18 +489,25 @@ func newFixedSizeListReader(rctx *readerCtx, field *arrow.Field, info file.Level
 }
 
 // helper function to combine chunks into a single array.
-//
-// nested data conversion for chunked array outputs not yet implemented
-func chunksToSingle(chunked *arrow.Chunked) (arrow.ArrayData, error) {
+func chunksToSingle(chunked *arrow.Chunked, mem memory.Allocator) (arrow.ArrayData, error) {
 	switch len(chunked.Chunks()) {
 	case 0:
 		return array.NewData(chunked.DataType(), 0, []*memory.Buffer{nil, nil}, nil, 0, 0), nil
 	case 1:
 		data := chunked.Chunk(0).Data()
 		data.Retain() // we pass control to the caller
 		return data, nil
-	default: // if an item reader yields a chunked array, this is not yet implemented
-		return nil, arrow.ErrNotImplemented
+	default:
+		// concatenate multiple chunks into a single array
+		concatenated, err := array.Concatenate(chunked.Chunks(), mem)
+		if err != nil {
+			return nil, err
+		}
+		defer concatenated.Release()
+
+		data := concatenated.Data()
+		data.Retain()
+		return data, nil
 	}
 }
 

parquet/pqarrow/column_readers_test.go

@@ -0,0 +1,150 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pqarrow
+
+import (
+	"testing"
+
+	"github.com/apache/arrow-go/v18/arrow"
+	"github.com/apache/arrow-go/v18/arrow/array"
+	"github.com/apache/arrow-go/v18/arrow/memory"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestChunksToSingle(t *testing.T) {
+	mem := memory.NewGoAllocator()
+
+	t.Run("empty chunked array", func(t *testing.T) {
+		chunked := arrow.NewChunked(arrow.PrimitiveTypes.Int32, []arrow.Array{})
+		defer chunked.Release()
+
+		result, err := chunksToSingle(chunked, mem)
+		require.NoError(t, err)
+		defer result.Release()
+
+		assert.Equal(t, 0, result.Len())
+	})
+
+	t.Run("single chunk", func(t *testing.T) {
+		bldr := array.NewInt32Builder(mem)
+		defer bldr.Release()
+		bldr.AppendValues([]int32{1, 2, 3, 4, 5}, nil)
+		arr := bldr.NewInt32Array()
+		defer arr.Release()
+
+		chunked := arrow.NewChunked(arrow.PrimitiveTypes.Int32, []arrow.Array{arr})
+		defer chunked.Release()
+
+		result, err := chunksToSingle(chunked, mem)
+		require.NoError(t, err)
+		defer result.Release()
+
+		assert.Equal(t, 5, result.Len())
+	})
+
+	t.Run("multiple chunks", func(t *testing.T) {
+		bldr := array.NewInt32Builder(mem)
+		defer bldr.Release()
+
+		bldr.AppendValues([]int32{1, 2, 3}, nil)
+		chunk1 := bldr.NewInt32Array()
+		defer chunk1.Release()
+
+		bldr.AppendValues([]int32{4, 5, 6}, nil)
+		chunk2 := bldr.NewInt32Array()
+		defer chunk2.Release()
+
+		bldr.AppendValues([]int32{7, 8, 9, 10}, nil)
+		chunk3 := bldr.NewInt32Array()
+		defer chunk3.Release()
+
+		chunked := arrow.NewChunked(arrow.PrimitiveTypes.Int32, []arrow.Array{chunk1, chunk2, chunk3})
+		defer chunked.Release()
+
+		result, err := chunksToSingle(chunked, mem)
+		require.NoError(t, err)
+		defer result.Release()
+
+		assert.Equal(t, 10, result.Len())
+
+		// Verify concatenated values
+		resultArr := array.MakeFromData(result).(*array.Int32)
+		defer resultArr.Release()
+		for i, expected := range []int32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10} {
+			assert.Equal(t, expected, resultArr.Value(i))
+		}
+	})
+
+	t.Run("multiple chunks with nulls", func(t *testing.T) {
+		bldr := array.NewInt32Builder(mem)
+		defer bldr.Release()
+
+		bldr.AppendValues([]int32{1, 0, 3}, []bool{true, false, true})
+		chunk1 := bldr.NewInt32Array()
+		defer chunk1.Release()
+
+		bldr.AppendValues([]int32{4, 0, 6}, []bool{true, false, true})
+		chunk2 := bldr.NewInt32Array()
+		defer chunk2.Release()
+
+		chunked := arrow.NewChunked(arrow.PrimitiveTypes.Int32, []arrow.Array{chunk1, chunk2})
+		defer chunked.Release()
+
+		result, err := chunksToSingle(chunked, mem)
+		require.NoError(t, err)
+		defer result.Release()
+
+		assert.Equal(t, 6, result.Len())
+		assert.Equal(t, 2, result.NullN())
+
+		resultArr := array.MakeFromData(result).(*array.Int32)
+		defer resultArr.Release()
+		assert.False(t, resultArr.IsValid(1))
+		assert.False(t, resultArr.IsValid(4))
+		assert.Equal(t, int32(1), resultArr.Value(0))
+		assert.Equal(t, int32(3), resultArr.Value(2))
+	})
+
+	t.Run("multiple chunks string type", func(t *testing.T) {
+		bldr := array.NewStringBuilder(mem)
+		defer bldr.Release()
+
+		bldr.AppendValues([]string{"hello", "world"}, nil)
+		chunk1 := bldr.NewStringArray()
+		defer chunk1.Release()
+
+		bldr.AppendValues([]string{"arrow", "parquet"}, nil)
+		chunk2 := bldr.NewStringArray()
+		defer chunk2.Release()
+
+		chunked := arrow.NewChunked(arrow.BinaryTypes.String, []arrow.Array{chunk1, chunk2})
+		defer chunked.Release()
+
+		result, err := chunksToSingle(chunked, mem)
+		require.NoError(t, err)
+		defer result.Release()
+
+		assert.Equal(t, 4, result.Len())
+
+		resultArr := array.MakeFromData(result).(*array.String)
+		defer resultArr.Release()
+		assert.Equal(t, "hello", resultArr.Value(0))
+		assert.Equal(t, "parquet", resultArr.Value(3))
+	})
+}

parquet/pqarrow/file_reader.go

@@ -527,6 +527,7 @@ func (fr *FileReader) GetRecordReader(ctx context.Context, colIndices, rowGroups
 		parallel:     fr.Props.Parallel,
 		sc:           sc,
 		fieldReaders: readers,
+		mem:          fr.mem,
 	}
 	rr.refCount.Add(1)
 	return rr, nil
@@ -721,6 +722,7 @@ type recordReader struct {
 	fieldReaders []*ColumnReader
 	cur          arrow.RecordBatch
 	err          error
+	mem          memory.Allocator
 
 	refCount atomic.Int64
 }
@@ -789,7 +791,7 @@ func (r *recordReader) next() bool {
 			return io.EOF
 		}
 
-		arrdata, err := chunksToSingle(data)
+		arrdata, err := chunksToSingle(data, r.mem)
 		if err != nil {
 			return err
 		}