From 6e043a8abfe861dfa0630fb5bbba2d8046a5f23c Mon Sep 17 00:00:00 2001
From: Matthias Loibl <mail@matthiasloibl.com>
Date: Sat, 23 Nov 2024 19:02:43 +0100
Subject: [PATCH] pqarrow/arrowutils: Handle empty structs correctly

---
 pqarrow/arrowutils/sort.go      |  6 ++++++
 pqarrow/arrowutils/sort_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/pqarrow/arrowutils/sort.go b/pqarrow/arrowutils/sort.go
index 09aa16f85..9f7c14d6b 100644
--- a/pqarrow/arrowutils/sort.go
+++ b/pqarrow/arrowutils/sort.go
@@ -279,6 +279,12 @@ func TakeListColumn(ctx context.Context, a *array.List, idx int, arr []arrow.Arr
 func TakeStructColumn(ctx context.Context, a *array.Struct, idx int, arr []arrow.Array, indices *array.Int32) error {
 	aType := a.Data().DataType().(*arrow.StructType)
 
+	// Immediately, return this struct if it has no fields/columns
+	if a.NumField() == 0 {
+		arr[idx] = a
+		return nil
+	}
+
 	cols := make([]arrow.Array, a.NumField())
 	names := make([]string, a.NumField())
 	defer func() {
diff --git a/pqarrow/arrowutils/sort_test.go b/pqarrow/arrowutils/sort_test.go
index 2ba5e2797..1bcf3150b 100644
--- a/pqarrow/arrowutils/sort_test.go
+++ b/pqarrow/arrowutils/sort_test.go
@@ -585,6 +585,36 @@ func TestReorderRecord(t *testing.T) {
 		require.Equal(t, "[3 2 5 1 4]", readRunEndEncodedDictionary(resultStruct.Field(1).(*array.RunEndEncoded)))
 		require.Equal(t, "[3 2 5 1 4]", resultStruct.Field(2).(*array.Int64).String())
 	})
+
+	t.Run("StructEmpty", func(t *testing.T) {
+		mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+		defer mem.AssertSize(t, 0)
+
+		b := array.NewRecordBuilder(mem, arrow.NewSchema(
+			[]arrow.Field{
+				{
+					Name: "struct",
+					Type: arrow.StructOf(),
+				},
+			}, &arrow.Metadata{},
+		))
+		defer b.Release()
+		b.Field(0).AppendNulls(5)
+
+		r := b.NewRecord()
+		defer r.Release()
+
+		indices := array.NewInt32Builder(mem)
+		indices.AppendValues([]int32{2, 1, 4, 0, 3}, nil)
+		by := indices.NewInt32Array()
+		defer by.Release()
+
+		result, err := Take(compute.WithAllocator(context.Background(), mem), r, by)
+		require.Nil(t, err)
+		defer result.Release()
+		resultStruct := result.Column(0).(*array.Struct)
+		resultStruct.Len()
+	})
 }
 
 // Use all supported sort field.