Skip to content

Commit

Permalink
pqarrow/arrowutils: Handle empty structs correctly
Browse files Browse the repository at this point in the history
  • Loading branch information
metalmatze committed Nov 25, 2024
1 parent ab75300 commit 6e043a8
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 0 deletions.
6 changes: 6 additions & 0 deletions pqarrow/arrowutils/sort.go
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
30 changes: 30 additions & 0 deletions pqarrow/arrowutils/sort_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down

0 comments on commit 6e043a8

Please sign in to comment.