zeroshade commented on code in PR #1337:
URL: https://github.com/apache/arrow-go/pull/1337#discussion_r4125742171
##########
arrow/compute/internal/kernels/vector_sort.go:
##########
@@ -55,7 +55,20 @@ type SortState = SortOptions
// SortKey defines a column to sort by with its ordering and null placement
options.
type SortKey struct {
- ColumnIndex int
+ // ColumnIndex is the top-level column to sort by. Ignored when
ColumnPath
+ // is non-empty; ColumnPath[0] is used as the top-level index instead.
+ ColumnIndex int
+
+ // ColumnPath, if non-empty, addresses a field nested inside struct
+ // columns: ColumnPath[0] selects the top-level column (as ColumnIndex
+ // would), and each subsequent entry selects a child field of the
+ // preceding struct column. A struct that is null at a given row makes
+ // every descendant null at that row for sorting purposes, regardless of
+ // the descendant's own physical validity bitmap. Only struct nesting is
+ // supported; a path element pointing into a list, map, or union column
+ // is an error.
+ ColumnPath []int
Review Comment:
Adding a slice field makes `SortKey` non-comparable, so existing `key ==
other` / map-key usage stops compiling. `SortKey` has been public since
v18.6.0, so this is a breaking change in a minor release. It needs a deliberate
call. Arrow C++ `SortKey` takes a `FieldRef` target, which may be the more
consistent shape.
##########
arrow/compute/vector_sort.go:
##########
@@ -106,6 +109,97 @@ func DefaultSortKey() SortKey {
}
}
+// maskedStructField returns the specified field index with a masked validity
bitmap.
+//
+// - A null struct row will return null regardless of any child values.
+// - The returned array is a new array that callers must [Release].
+// - This mirrors the unexported
array.(*Struct).newStructFieldWithParentValidityMask.
+func maskedStructField(st *array.Struct, fieldIndex int) arrow.Array {
+ field := st.Field(fieldIndex)
+
+ // short circuit if there are no nulls in the array
+ if st.NullN() == 0 {
+ field.Retain()
+ return field
+ }
+
+ nullBitmapBytes := field.NullBitmapBytes()
+ fieldOffset := field.Data().Offset()
+
+ var maskedNullBitmapBytes []byte
+ if len(nullBitmapBytes) == 0 {
+ fieldEnd := int64(fieldOffset + field.Len())
+ maskedNullBitmapBytes = make([]byte,
int(bitutil.BytesForBits(fieldEnd)))
+ bitutil.SetBitsTo(maskedNullBitmapBytes, 0, fieldEnd, true)
+ } else {
+ maskedNullBitmapBytes = make([]byte, len(nullBitmapBytes))
+ copy(maskedNullBitmapBytes, nullBitmapBytes)
+ }
+ for i := 0; i < field.Len(); i++ {
+ if st.IsNull(i) {
+ bitutil.ClearBit(maskedNullBitmapBytes, fieldOffset+i)
+ }
+ }
+
+ // Normalize to a zero offset so maskedNullBitmapBytes (indexed from
+ // fieldOffset above) lines up with the sliced buffers.
+ sliced := array.NewSliceData(field.Data(), 0, int64(field.Len()))
+ defer sliced.Release()
+
+ origBufs := sliced.Buffers()
+ bufs := make([]*memory.Buffer, len(origBufs))
+ copy(bufs, origBufs)
+ bufs[0] = memory.NewBufferBytes(maskedNullBitmapBytes)
+
+ data := array.NewData(sliced.DataType(), sliced.Len(), bufs,
sliced.Children(), array.UnknownNullCount, 0)
Review Comment:
`NewSliceData(d, 0, n)` keeps `d.Offset()` and shares the buffers; it does
not reset the offset to zero. Passing `0` here makes the new array read values
from index 0 and misaligns the bitmap, which was built at `fieldOffset+i`.
Repro: ids `[100,101,1,2,3,4]`, ranks `[-5,-4,30,null,10,20]`, then
`batch.NewSlice(2,6)` and sort by `{0,1,0}`. The result is `[1 2 3 4]`; the
expected order is `[3 4 1 2]`. Using `sliced.Offset()` fixes it. You could also
skip the `NewSliceData` round-trip and build from `field.Data()` directly.
Please add a sliced batch/chunk test.
##########
arrow/compute/vector_sort.go:
##########
@@ -106,6 +109,97 @@ func DefaultSortKey() SortKey {
}
}
+// maskedStructField returns the specified field index with a masked validity
bitmap.
+//
+// - A null struct row will return null regardless of any child values.
+// - The returned array is a new array that callers must [Release].
+// - This mirrors the unexported
array.(*Struct).newStructFieldWithParentValidityMask.
+func maskedStructField(st *array.Struct, fieldIndex int) arrow.Array {
+ field := st.Field(fieldIndex)
+
+ // short circuit if there are no nulls in the array
+ if st.NullN() == 0 {
+ field.Retain()
+ return field
+ }
+
+ nullBitmapBytes := field.NullBitmapBytes()
+ fieldOffset := field.Data().Offset()
+
+ var maskedNullBitmapBytes []byte
+ if len(nullBitmapBytes) == 0 {
+ fieldEnd := int64(fieldOffset + field.Len())
+ maskedNullBitmapBytes = make([]byte,
int(bitutil.BytesForBits(fieldEnd)))
+ bitutil.SetBitsTo(maskedNullBitmapBytes, 0, fieldEnd, true)
+ } else {
+ maskedNullBitmapBytes = make([]byte, len(nullBitmapBytes))
+ copy(maskedNullBitmapBytes, nullBitmapBytes)
+ }
+ for i := 0; i < field.Len(); i++ {
+ if st.IsNull(i) {
+ bitutil.ClearBit(maskedNullBitmapBytes, fieldOffset+i)
+ }
+ }
+
+ // Normalize to a zero offset so maskedNullBitmapBytes (indexed from
+ // fieldOffset above) lines up with the sliced buffers.
+ sliced := array.NewSliceData(field.Data(), 0, int64(field.Len()))
+ defer sliced.Release()
+
+ origBufs := sliced.Buffers()
+ bufs := make([]*memory.Buffer, len(origBufs))
+ copy(bufs, origBufs)
+ bufs[0] = memory.NewBufferBytes(maskedNullBitmapBytes)
+
+ data := array.NewData(sliced.DataType(), sliced.Len(), bufs,
sliced.Children(), array.UnknownNullCount, 0)
+ defer data.Release()
+
+ return array.MakeFromData(data)
+}
+
+// resolveSortColumnPath walks path into col, descending through struct
+// fields. path[0] has already been consumed to select col itself (the
+// top-level column); remaining elements each select a child field of the
+// preceding struct. Every array returned for a non-empty path (including
+// intermediate ones released internally) is a new reference obtained via
+// maskedStructField, so a null ancestor struct correctly forces
+// its descendants null for sorting purposes.
+func resolveSortColumnPath(col arrow.Array, path []int) (arrow.Array, error) {
+ if len(path) == 0 {
+ col.Retain()
+ return col, nil
+ }
+
+ st, ok := col.(*array.Struct)
+ if !ok {
+ return nil, fmt.Errorf("%w: sort key column path element
requires a struct column, got %s",
+ arrow.ErrInvalid, col.DataType())
+ }
+ idx := path[0]
+ if idx < 0 || idx >= st.NumField() {
+ return nil, fmt.Errorf("%w: sort key struct field index %d out
of range", arrow.ErrIndex, idx)
Review Comment:
nit: an out-of-range top-level index returns `ErrInvalid`, but an
out-of-range nested index returns `ErrIndex`. Pick one.
##########
arrow/compute/vector_sort_test.go:
##########
@@ -697,6 +697,96 @@ func TestSortRecordBatch(t *testing.T) {
require.Error(t, err)
require.ErrorIs(t, err, arrow.ErrInvalid)
})
+
+ t.Run("NestedStructColumnPath", func(t *testing.T) {
+ nestedType := arrow.StructOf(
+ arrow.Field{Name: "id", Type:
arrow.PrimitiveTypes.Int32},
+ arrow.Field{Name: "info", Type: arrow.StructOf(
+ arrow.Field{Name: "rank", Type:
arrow.PrimitiveTypes.Int32, Nullable: true},
+ ), Nullable: true},
+ )
+ nestedSchema := arrow.NewSchema([]arrow.Field{{Name: "s", Type:
nestedType}}, nil)
+
+ bldr := array.NewStructBuilder(mem, nestedType)
+ defer bldr.Release()
+ idBldr := bldr.FieldBuilder(0).(*array.Int32Builder)
+ infoBldr := bldr.FieldBuilder(1).(*array.StructBuilder)
+ rankBldr := infoBldr.FieldBuilder(0).(*array.Int32Builder)
+
+ // Row 0: id=1, info={rank:30}
+ bldr.Append(true)
+ idBldr.Append(1)
+ infoBldr.Append(true)
+ rankBldr.Append(30)
+
+ // Row 1: id=2, info=null. StructBuilder.Append(false)
auto-appends
+ // null to info's children (rank), so rankBldr is not touched
here;
+ // rank's underlying storage slot still ends up
non-null-looking at
+ // the physical level, which is exactly what parent-validity
masking
+ // must override.
+ bldr.Append(true)
+ idBldr.Append(2)
+ infoBldr.Append(false)
Review Comment:
`StructBuilder.Append(false)` calls `AppendNull()` on every child
(`array/struct.go:395-398`), so `rank` is already physically null here and the
comment is inaccurate. As a result, this test and the multi-chunk table test
pass even when `maskedStructField` just returns `st.Field(i)` unmasked (I
checked).
Please build a case where the parent is null but the child slot is valid,
e.g. with `array.NewStructArrayWithNulls` or `NewStructData` and a hand-built
bitmap.
##########
arrow/compute/vector_sort.go:
##########
@@ -172,12 +276,37 @@ func sortIndicesImpl(ctx context.Context, opts
FunctionOptions, input Datum) (Da
sortColumns = make([]*arrow.Chunked, len(inputSortKeys))
needsRelease = make([]bool, len(inputSortKeys))
for i, key := range inputSortKeys {
- if key.ColumnIndex < 0 || int64(key.ColumnIndex) >=
tbl.NumCols() {
- return nil, fmt.Errorf("%w: sort key %d has
invalid column index %d", arrow.ErrInvalid, i, key.ColumnIndex)
+ topIdx := key.ColumnIndex
+ if len(key.ColumnPath) > 0 {
+ topIdx = key.ColumnPath[0]
+ }
+ if topIdx < 0 || int64(topIdx) >= tbl.NumCols() {
+ return nil, fmt.Errorf("%w: sort key %d has
invalid column index %d", arrow.ErrInvalid, i, topIdx)
+ }
+ chunked := tbl.Column(topIdx).Data()
+ if len(key.ColumnPath) <= 1 {
+ // Table columns are already Chunked; borrow
from the table (do not Release).
+ sortColumns[i] = chunked
+ needsRelease[i] = false
+ continue
}
- // Table columns are already Chunked; borrow from the
table (do not Release).
- sortColumns[i] = tbl.Column(key.ColumnIndex).Data()
- needsRelease[i] = false
+
+ resolvedChunks := make([]arrow.Array,
len(chunked.Chunks()))
+ var leafType arrow.DataType
+ for c, chunk := range chunked.Chunks() {
+ resolved, err := resolveSortColumnPath(chunk,
key.ColumnPath[1:])
+ if err != nil {
+ return nil, fmt.Errorf("sort key %d:
%w", i, err)
+ }
+ defer resolved.Release()
+ resolvedChunks[c] = resolved
+ leafType = resolved.DataType()
+ }
+ if leafType == nil {
+ leafType = chunked.DataType()
Review Comment:
With zero chunks, the path is never validated and `leafType` falls back to
the struct type. Validating the path once against the column type
(`compute.FieldPath.GetFieldFromType`) before the per-chunk loop would give the
correct leaf type and catch bad paths even on empty tables.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
