This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new 13b6d07f feat(arrow/compute): add IEEE 754 total-order float sorting
option (#1430)
13b6d07f is described below
commit 13b6d07f83095ad02dcd8747a0cc724f5168a22f
Author: Matt Topol <[email protected]>
AuthorDate: Wed Oct 28 20:06:20 2026 +0000
feat(arrow/compute): add IEEE 754 total-order float sorting option (#1430)
### Rationale for this change
The sort kernels order floats the way Arrow C++ does: numerically with
`-0 == +0`, and every NaN (whatever its sign) grouped next to the nulls
according to `NullPlacement`, regardless of `Order`. Some consumers need
a deterministic total order over floats that distinguishes `-0`/`+0` and
signed NaNs and treats NaN as an ordinary value, matching Java's
`Double.compare` / IEEE 754 `totalOrder`.
### What changes are included in this PR?
- New `SortKey.FloatTotalOrder bool` field:
- `false` (default): existing behavior, unchanged.
- `true`: IEEE 754 totalOrder for float16/float32/float64:
`-NaN < -Inf < negative values < -0 < +0 < positive values < +Inf <
+NaN`.
NaNs sort as values (not null-like), and `Descending` reverses the whole
order including NaNs.
- Nulls are unaffected: they still follow `NullPlacement` in both modes.
- Non-float columns ignore the field.
- In total-order mode the float comparators compare order-preserving
integer keys derived from the raw bits, and report no null-like values
so NaNs are not partitioned out.
- The single-chunk, no-null radix fast path (previously integers only)
now also handles total-order floats by radix sorting the materialized
integer keys.
### Are these changes tested?
Yes:
- `TestSortIndicesFloatTotalOrder`: float64/32/16 × ascending/descending
× nulls first/last, over the single-array, chunked (per-chunk sort +
merge) and multi-key record batch paths, with -NaN, NaN, ±Inf,
interleaved ±0 and a null.
- `TestSortIndicesFloatDefaultOrderingUnchanged`: pins the default
behavior.
- `TestSortIndicesSingleChunkRadixFloatTotalOrder`: radix fast path
(sliced chunk, 4096 rows) for all three float widths, checked against an
independent reference comparator.
### Are there any user-facing changes?
Yes, a new opt-in `SortKey.FloatTotalOrder` field. The default is
unchanged.
---
arrow/compute/internal/kernels/vector_sort.go | 33 +++-
.../internal/kernels/vector_sort_bench_test.go | 71 ++++++++
.../internal/kernels/vector_sort_physical.go | 63 +++++++
.../compute/internal/kernels/vector_sort_radix.go | 39 ++++
.../internal/kernels/vector_sort_radix_test.go | 199 +++++++++++++++++++++
.../internal/kernels/vector_sort_support.go | 19 ++
arrow/compute/vector_sort.go | 3 +-
arrow/compute/vector_sort_test.go | 167 +++++++++++++++++
8 files changed, 588 insertions(+), 6 deletions(-)
diff --git a/arrow/compute/internal/kernels/vector_sort.go
b/arrow/compute/internal/kernels/vector_sort.go
index 025c3dd3..7f09beb0 100644
--- a/arrow/compute/internal/kernels/vector_sort.go
+++ b/arrow/compute/internal/kernels/vector_sort.go
@@ -53,11 +53,24 @@ func (SortOptions) TypeName() string { return "SortOptions"
}
type SortState = SortOptions
-// SortKey defines a column to sort by with its ordering and null placement
options.
+// SortKey defines a column to sort by with its ordering, null placement and
float ordering options.
type SortKey struct {
ColumnIndex int
Order SortOrder
NullPlacement NullPlacement
+ // FloatTotalOrder orders float16/float32/float64 columns by the IEEE
754 totalOrder
+ // predicate: -NaN < -Inf < negative values < -0 < +0 < positive values
< +Inf < +NaN.
+ // NaNs are ordered as values (not null-like) and Descending reverses
the whole order,
+ // NaNs included; nulls still follow NullPlacement. When false (the
default), floats
+ // compare numerically with -0 == +0 and every NaN, regardless of sign,
is grouped next
+ // to the nulls according to NullPlacement, independent of Order (Arrow
C++ behavior).
+ // Non-float columns ignore it.
+ //
+ // Total order is usually faster than the default ordering, since each
value maps to an
+ // integer key and no NaN checks are needed. When it is the only sort
key and the column
+ // is a single chunk of at least 2048 rows without nulls, the sort is a
radix sort over
+ // those keys, which allocates a temporary key buffer the size of the
column's values.
+ FloatTotalOrder bool
}
// Chunk-aware sort_indices: logical row IDs 0..n-1, no chunk concatenation.
Structure follows
@@ -161,8 +174,8 @@ func newFixedSizeBinaryComparator(chunks []arrow.Array,
numRows int, vn bool, ro
}
// createChunkedComparator builds a column comparator for these chunks (one
Arrow type for all chunks).
-// rowMap may be shared by aligned multi-key columns.
-func createChunkedComparator(chunks []arrow.Array, numRows int, rowMap
logicalRowMap) (columnComparator, error) {
+// rowMap may be shared by aligned multi-key columns. floatTotalOrder enables
or disables IEEE 754 totalOrder for floats.
+func createChunkedComparator(chunks []arrow.Array, numRows int, rowMap
logicalRowMap, floatTotalOrder bool) (columnComparator, error) {
if len(chunks) == 0 {
return nil, fmt.Errorf("%w: cannot create comparator for empty
chunk list", arrow.ErrInvalid)
}
@@ -202,10 +215,19 @@ func createChunkedComparator(chunks []arrow.Array,
numRows int, rowMap logicalRo
case arrow.UINT64:
return newPhysicalSortUint64Column(chunks, numRows,
validityNulls, rowMap), nil
case arrow.FLOAT16:
+ if floatTotalOrder {
+ return
&physicalSortFloat16TotalOrderColumn{physicalSortFloatTotalOrderBase{newPhysicalColumnBase(chunks,
numRows, validityNulls, rowMap)}}, nil
+ }
return newPhysicalSortFloat16Column(chunks, numRows,
validityNulls, rowMap), nil
case arrow.FLOAT32:
+ if floatTotalOrder {
+ return
&physicalSortFloat32TotalOrderColumn{physicalSortFloatTotalOrderBase{newPhysicalColumnBase(chunks,
numRows, validityNulls, rowMap)}}, nil
+ }
return newPhysicalSortFloat32Column(chunks, numRows,
validityNulls, rowMap), nil
case arrow.FLOAT64:
+ if floatTotalOrder {
+ return
&physicalSortFloat64TotalOrderColumn{physicalSortFloatTotalOrderBase{newPhysicalColumnBase(chunks,
numRows, validityNulls, rowMap)}}, nil
+ }
return newPhysicalSortFloat64Column(chunks, numRows,
validityNulls, rowMap), nil
case arrow.DECIMAL32:
return newPhysicalSortDecimal32Column(chunks, numRows,
validityNulls, rowMap), nil
@@ -420,7 +442,7 @@ func SortIndices(ctx *exec.KernelCtx, columns
[]*arrow.Chunked, keys []SortKey)
comparators := make([]columnComparator, len(columns))
for i, col := range columns {
- comp, err := createChunkedComparator(col.Chunks(), nRows,
sharedRowMap)
+ comp, err := createChunkedComparator(col.Chunks(), nRows,
sharedRowMap, keys[i].FloatTotalOrder)
if err != nil {
return nil, err
}
@@ -456,7 +478,8 @@ func SortIndices(ctx *exec.KernelCtx, columns
[]*arrow.Chunked, keys []SortKey)
k0 := keys[0]
c0 := comparators[0]
if nRows >= radixSortMinRows &&
!c0.columnHasValidityNulls() &&
- sortIndicesSingleIntRadix(ctx, chunks[0],
indices, k0.Order) {
+ (sortIndicesSingleIntRadix(ctx, chunks[0],
indices, k0.Order) ||
+ (k0.FloatTotalOrder &&
sortIndicesSingleFloatTotalOrderRadix(ctx, chunks[0], indices, k0.Order))) {
out.Buffers[1].WrapBuffer(buf)
return out, nil
}
diff --git a/arrow/compute/internal/kernels/vector_sort_bench_test.go
b/arrow/compute/internal/kernels/vector_sort_bench_test.go
index 31f735c7..621e8741 100644
--- a/arrow/compute/internal/kernels/vector_sort_bench_test.go
+++ b/arrow/compute/internal/kernels/vector_sort_bench_test.go
@@ -21,6 +21,7 @@ package kernels
import (
"context"
"fmt"
+ "math"
"testing"
"github.com/apache/arrow-go/v18/arrow"
@@ -215,3 +216,73 @@ func BenchmarkSortIndices_Int64_Distributions(b
*testing.B) {
}
}
}
+
+// makeChunkedFloat64Split returns n float64 rows in numChunks contiguous
arrays. Every
+// nullEvery-th row is null and every nanEvery-th row is NaN (0 disables
either).
+func makeChunkedFloat64Split(tb testing.TB, mem memory.Allocator, n,
numChunks, nullEvery, nanEvery int) *arrow.Chunked {
+ tb.Helper()
+ chunks := make([]arrow.Array, 0, numChunks)
+ base, rem, global := n/numChunks, n%numChunks, 0
+ for c := range numChunks {
+ sz := base
+ if c < rem {
+ sz++
+ }
+ bld := array.NewFloat64Builder(mem)
+ for i := range sz {
+ g := global + i
+ switch {
+ case nullEvery > 0 && g%nullEvery == 0:
+ bld.AppendNull()
+ case nanEvery > 0 && g%nanEvery == 0:
+ bld.Append(math.NaN())
+ default:
+ x := uint64(g)
+
bld.Append(float64(int64((x*6364136223846793005)^(x>>12))) / 1e6)
+ }
+ }
+ chunks = append(chunks, bld.NewArray())
+ bld.Release()
+ global += sz
+ }
+ ch := arrow.NewChunked(arrow.PrimitiveTypes.Float64, chunks)
+ for _, arr := range chunks {
+ arr.Release()
+ }
+ tb.Cleanup(func() { ch.Release() })
+ return ch
+}
+
+// BenchmarkSortIndices_Float64 compares the default float ordering with
+// SortKey.FloatTotalOrder. Without nulls and with a single chunk, total order
+// takes the radix path; nulls or multiple chunks use the comparator sort.
+func BenchmarkSortIndices_Float64(b *testing.B) {
+ const rows = 65536
+ for _, tc := range []struct {
+ name string
+ chunks, nulls, nans int
+ }{
+ {"chunks=1", 1, 0, 0},
+ {"chunks=1/nans", 1, 0, 97},
+ {"chunks=16", 16, 0, 0},
+ {"chunks=16/nulls/nans", 16, 31, 97},
+ } {
+ for _, totalOrder := range []bool{false, true} {
+ b.Run(fmt.Sprintf("%s/totalOrder=%t", tc.name,
totalOrder), func(b *testing.B) {
+ ctx, mem := newBenchKernelCtx(b)
+ col := makeChunkedFloat64Split(b, mem, rows,
tc.chunks, tc.nulls, tc.nans)
+ keys := []SortKey{{ColumnIndex: 0, Order:
Ascending, NullPlacement: NullsAtEnd, FloatTotalOrder: totalOrder}}
+ columns := []*arrow.Chunked{col}
+
+ b.ReportAllocs()
+ for b.Loop() {
+ res, err := SortIndices(ctx, columns,
keys)
+ if err != nil {
+ b.Fatal(err)
+ }
+ res.Release()
+ }
+ })
+ }
+ }
+}
diff --git a/arrow/compute/internal/kernels/vector_sort_physical.go
b/arrow/compute/internal/kernels/vector_sort_physical.go
index 60d52b87..642253be 100644
--- a/arrow/compute/internal/kernels/vector_sort_physical.go
+++ b/arrow/compute/internal/kernels/vector_sort_physical.go
@@ -543,6 +543,69 @@ func (c *physicalSortFloat64Column)
columnHasValidityNulls() bool {
return c.base.columnHasValidityNulls()
}
+// Total-order float comparators (SortKey.FloatTotalOrder) compare values by
their IEEE 754
+// totalOrder key. They are separate types so the default float comparators
above carry no
+// per-comparison check for the option. NaNs are ordinary values, so these
columns report no
+// null-like values; validity nulls still follow NullPlacement.
+type physicalSortFloatTotalOrderBase struct{ base physicalColumnBase }
+
+func (c *physicalSortFloatTotalOrderBase) isNullAt(row uint64) bool {
+ return c.base.isNullAtGlobal(row)
+}
+func (c *physicalSortFloatTotalOrderBase) hasNullLikeValues() bool { return
false }
+func (c *physicalSortFloatTotalOrderBase) isNullLikeAt(uint64) bool { return
false }
+func (c *physicalSortFloatTotalOrderBase) columnHasValidityNulls() bool {
+ return c.base.columnHasValidityNulls()
+}
+
+type physicalSortFloat16TotalOrderColumn struct {
+ physicalSortFloatTotalOrderBase
+}
+
+func (c *physicalSortFloat16TotalOrderColumn) compareRowsForKey(i, j uint64,
key SortKey) int {
+ ai, aj, li, lj := c.base.pair(i, j)
+ a := ai.(*array.Float16)
+ b := aj.(*array.Float16)
+ if c.base.validityNulls {
+ if v, stop := compareKeyedNulls(a.IsNull(li), b.IsNull(lj),
key); stop {
+ return v
+ }
+ }
+ return compareOrdered(key.Order, totalOrderKey16(a.Value(li).Uint16()),
totalOrderKey16(b.Value(lj).Uint16()))
+}
+
+type physicalSortFloat32TotalOrderColumn struct {
+ physicalSortFloatTotalOrderBase
+}
+
+func (c *physicalSortFloat32TotalOrderColumn) compareRowsForKey(i, j uint64,
key SortKey) int {
+ ai, aj, li, lj := c.base.pair(i, j)
+ a := ai.(*array.Float32)
+ b := aj.(*array.Float32)
+ if c.base.validityNulls {
+ if v, stop := compareKeyedNulls(a.IsNull(li), b.IsNull(lj),
key); stop {
+ return v
+ }
+ }
+ return compareOrdered(key.Order,
totalOrderKey32(math.Float32bits(a.Value(li))),
totalOrderKey32(math.Float32bits(b.Value(lj))))
+}
+
+type physicalSortFloat64TotalOrderColumn struct {
+ physicalSortFloatTotalOrderBase
+}
+
+func (c *physicalSortFloat64TotalOrderColumn) compareRowsForKey(i, j uint64,
key SortKey) int {
+ ai, aj, li, lj := c.base.pair(i, j)
+ a := ai.(*array.Float64)
+ b := aj.(*array.Float64)
+ if c.base.validityNulls {
+ if v, stop := compareKeyedNulls(a.IsNull(li), b.IsNull(lj),
key); stop {
+ return v
+ }
+ }
+ return compareOrdered(key.Order,
totalOrderKey64(math.Float64bits(a.Value(li))),
totalOrderKey64(math.Float64bits(b.Value(lj))))
+}
+
type physicalSortDecimal32Column struct{ base physicalColumnBase }
func newPhysicalSortDecimal32Column(chunks []arrow.Array, numRows int, vn
bool, rowMap logicalRowMap) *physicalSortDecimal32Column {
diff --git a/arrow/compute/internal/kernels/vector_sort_radix.go
b/arrow/compute/internal/kernels/vector_sort_radix.go
index 579f2512..7e313865 100644
--- a/arrow/compute/internal/kernels/vector_sort_radix.go
+++ b/arrow/compute/internal/kernels/vector_sort_radix.go
@@ -19,9 +19,13 @@
package kernels
import (
+ "math"
+ "unsafe"
+
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/array"
"github.com/apache/arrow-go/v18/arrow/compute/exec"
+ "github.com/apache/arrow-go/v18/arrow/float16"
)
const radixSortMinRows = 2048
@@ -69,6 +73,41 @@ func sortIndicesSingleIntRadix(ctx *exec.KernelCtx, arr
arrow.Array, indices []u
return true
}
+// sortIndicesSingleFloatTotalOrderRadix is the SortKey.FloatTotalOrder
counterpart of
+// sortIndicesSingleIntRadix for a single chunk of non-null floats. The IEEE
754
+// totalOrder maps each float to a signed integer key, so the floats are radix
+// sorted through a materialized key buffer.
+func sortIndicesSingleFloatTotalOrderRadix(ctx *exec.KernelCtx, arr
arrow.Array, indices []uint64, order SortOrder) bool {
+ switch values := arr.(type) {
+ case *array.Float16:
+ radixSortFloatTotalOrder(ctx, indices, values.Values(), 2,
1<<15, order,
+ func(v float16.Num) int16 { return
totalOrderKey16(v.Uint16()) })
+ case *array.Float32:
+ radixSortFloatTotalOrder(ctx, indices, values.Float32Values(),
4, 1<<31, order,
+ func(v float32) int32 { return
totalOrderKey32(math.Float32bits(v)) })
+ case *array.Float64:
+ radixSortFloatTotalOrder(ctx, indices, values.Float64Values(),
8, 1<<63, order,
+ func(v float64) int64 { return
totalOrderKey64(math.Float64bits(v)) })
+ default:
+ return false
+ }
+ return true
+}
+
+func radixSortFloatTotalOrder[F any, K int16 | int32 | int64](ctx
*exec.KernelCtx, indices []uint64, values []F, numBytes int, signBit uint64,
order SortOrder, toKey func(F) K) {
+ if len(values) < 2 {
+ return
+ }
+ var zero K
+ keyBuf := ctx.Allocate(len(values) * int(unsafe.Sizeof(zero)))
+ defer keyBuf.Release()
+ keys := arrow.GetData[K](keyBuf.Buf())[:len(values)]
+ for i, v := range values {
+ keys[i] = toKey(v)
+ }
+ radixSortIntegerIndices(ctx, indices, keys, numBytes, signBit, order)
+}
+
// radixSortIntegerIndices performs a stable LSD radix sort over the logical
row
// indices. Flipping the sign bit orders signed integers correctly; inverting
the
// key reverses the sort order without reversing equal-value rows.
diff --git a/arrow/compute/internal/kernels/vector_sort_radix_test.go
b/arrow/compute/internal/kernels/vector_sort_radix_test.go
index c0388e67..47d2d3ce 100644
--- a/arrow/compute/internal/kernels/vector_sort_radix_test.go
+++ b/arrow/compute/internal/kernels/vector_sort_radix_test.go
@@ -22,12 +22,14 @@ import (
"cmp"
"context"
"fmt"
+ "math"
"slices"
"testing"
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/array"
"github.com/apache/arrow-go/v18/arrow/compute/exec"
+ "github.com/apache/arrow-go/v18/arrow/float16"
"github.com/apache/arrow-go/v18/arrow/memory"
)
@@ -147,3 +149,200 @@ func TestSortIndicesSingleChunkRadixSliced(t *testing.T) {
})
}
}
+
+// specTotalOrder is an independent reference for IEEE 754 totalOrder on
non-payload NaNs:
+// -NaN < -Inf < negatives < -0 < +0 < positives < +Inf < +NaN.
+func specTotalOrder(a, b float64) int {
+ rank := func(v float64) int {
+ switch {
+ case math.IsNaN(v) && math.Signbit(v):
+ return 0
+ case math.IsNaN(v):
+ return 2
+ }
+ return 1
+ }
+ if c := cmp.Compare(rank(a), rank(b)); c != 0 || rank(a) != 1 {
+ return c
+ }
+ if c := cmp.Compare(a, b); c != 0 {
+ return c
+ }
+ // Equal numerically: only -0 vs +0 can still differ.
+ return cmp.Compare(boolInt(!math.Signbit(a)), boolInt(!math.Signbit(b)))
+}
+
+func boolInt(b bool) int {
+ if b {
+ return 1
+ }
+ return 0
+}
+
+// TestSortIndicesSingleChunkRadixFloatTotalOrder covers the radix fast path
(single sliced
+// chunk, no nulls, >= radixSortMinRows) for SortKey.FloatTotalOrder against
the reference order.
+func TestSortIndicesSingleChunkRadixFloatTotalOrder(t *testing.T) {
+ const start, n = 37, 4096
+ special := []float64{math.NaN(),
math.Float64frombits(0xFFF8000000000000), math.Inf(1), math.Inf(-1),
+ 0, math.Copysign(0, -1), math.SmallestNonzeroFloat64,
-math.SmallestNonzeroFloat64, math.MaxFloat64, -math.MaxFloat64}
+ values := make([]float64, start+n+19)
+ for i := range values {
+ x := uint64(i)*6364136223846793005 + 1442695040888963407
+ if x%5 == 0 {
+ values[i] = special[int(x>>33)%len(special)]
+ } else {
+ values[i] = float64(int64(x>>40)%2001-1000) / 8
+ }
+ }
+
+ for _, dt := range []arrow.DataType{arrow.PrimitiveTypes.Float64,
arrow.PrimitiveTypes.Float32, arrow.FixedWidthTypes.Float16} {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ ctx := &exec.KernelCtx{Ctx:
exec.WithAllocator(context.Background(), mem)}
+ bldr := array.NewBuilder(mem, dt)
+ // Values as stored in the column type (narrowing keeps NaN
signs, may round magnitudes).
+ stored := make([]float64, len(values))
+ for i, v := range values {
+ switch b := bldr.(type) {
+ case *array.Float64Builder:
+ b.Append(v)
+ stored[i] = v
+ case *array.Float32Builder:
+ b.Append(float32(v))
+ stored[i] = float64(float32(v))
+ case *array.Float16Builder:
+ h := float16.New(float32(v))
+ b.Append(h)
+ stored[i] = float64(h.Float32())
+ }
+ }
+ full := bldr.NewArray()
+ bldr.Release()
+ sliced := array.NewSlice(full, start, start+n)
+ col := arrow.NewChunked(dt, []arrow.Array{sliced})
+
+ for _, order := range []SortOrder{Ascending, Descending} {
+ t.Run(fmt.Sprintf("%s/%v", dt, order), func(t
*testing.T) {
+ result, err := SortIndices(ctx,
[]*arrow.Chunked{col}, []SortKey{{Order: order, FloatTotalOrder: true}})
+ if err != nil {
+ t.Fatal(err)
+ }
+ defer result.Release()
+ got :=
arrow.GetData[uint64](result.Buffers[1].Buf)[:n]
+ want := make([]uint64, n)
+ for i := range want {
+ want[i] = uint64(i)
+ }
+ slices.SortStableFunc(want, func(a, b uint64)
int {
+ c :=
specTotalOrder(stored[start+int(a)], stored[start+int(b)])
+ if order == Descending {
+ return -c
+ }
+ return c
+ })
+ if !slices.Equal(got, want) {
+ t.Fatalf("radix total-order sort
differs from reference stable sort")
+ }
+ })
+ }
+ col.Release()
+ sliced.Release()
+ full.Release()
+ mem.AssertSize(t, 0)
+ }
+}
+
+// totalOrderBitPatterns lists raw bit patterns in ascending IEEE 754
totalOrder for a float
+// format with the given exponent and mantissa widths. It covers NaN payloads
and
+// signaling vs quiet NaNs: for positive NaNs, signaling < quiet and larger
payloads sort
+// later; negative NaNs mirror that.
+func totalOrderBitPatterns(expBits, mantBits uint) []uint64 {
+ sign := uint64(1) << (expBits + mantBits)
+ exp := ((uint64(1) << expBits) - 1) << mantBits
+ quiet := uint64(1) << (mantBits - 1)
+ one := (((uint64(1) << (expBits - 1)) - 1) << mantBits) // 1.0
+ pos := []uint64{
+ 0, // +0
+ 1, // smallest subnormal
+ one, // 1.0
+ exp - 1, // largest finite
+ exp, // +Inf
+ exp | 1, // sNaN, payload 1
+ exp | 2, // sNaN, payload 2
+ exp | quiet, // canonical qNaN
+ exp | quiet | 1, // qNaN, payload 1
+ exp | (quiet<<1 - 1), // qNaN, all mantissa bits set
+ }
+ out := make([]uint64, 0, 2*len(pos))
+ for i := len(pos) - 1; i >= 0; i-- {
+ out = append(out, pos[i]|sign) // negatives, most negative first
+ }
+ return append(out, pos...)
+}
+
+// TestSortIndicesFloatTotalOrderNaNPayloads checks NaN payload and
signaling/quiet ordering
+// with raw bit patterns, on both the comparator path (below radixSortMinRows)
and the
+// radix fast path.
+func TestSortIndicesFloatTotalOrderNaNPayloads(t *testing.T) {
+ formats := []struct {
+ dt arrow.DataType
+ expBits, mantBits uint
+ }{
+ {arrow.PrimitiveTypes.Float64, 11, 52},
+ {arrow.PrimitiveTypes.Float32, 8, 23},
+ {arrow.FixedWidthTypes.Float16, 5, 10},
+ }
+ for _, f := range formats {
+ sorted := totalOrderBitPatterns(f.expBits, f.mantBits)
+ for _, reps := range []int{1, radixSortMinRows/len(sorted) + 1}
{
+ for _, order := range []SortOrder{Ascending,
Descending} {
+ t.Run(fmt.Sprintf("%s/reps=%d/%v", f.dt, reps,
order), func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+ ctx := &exec.KernelCtx{Ctx:
exec.WithAllocator(context.Background(), mem)}
+
+ // Deterministically shuffled input,
repeated reps times.
+ bits := make([]uint64, 0,
len(sorted)*reps)
+ for r := range reps {
+ for i := range sorted {
+ bits = append(bits,
sorted[(i*7+r)%len(sorted)])
+ }
+ }
+ bldr := array.NewBuilder(mem, f.dt)
+ for _, b := range bits {
+ switch bb := bldr.(type) {
+ case *array.Float64Builder:
+
bb.Append(math.Float64frombits(b))
+ case *array.Float32Builder:
+
bb.Append(math.Float32frombits(uint32(b)))
+ case *array.Float16Builder:
+
bb.Append(float16.FromBits(uint16(b)))
+ }
+ }
+ arr := bldr.NewArray()
+ bldr.Release()
+ defer arr.Release()
+ col := arrow.NewChunked(f.dt,
[]arrow.Array{arr})
+ defer col.Release()
+
+ result, err := SortIndices(ctx,
[]*arrow.Chunked{col}, []SortKey{{Order: order, FloatTotalOrder: true}})
+ if err != nil {
+ t.Fatal(err)
+ }
+ defer result.Release()
+ got :=
arrow.GetData[uint64](result.Buffers[1].Buf)[:len(bits)]
+
+ rank := make(map[uint64]int,
len(sorted))
+ for i, b := range sorted {
+ rank[b] = i
+ }
+ for k := 1; k < len(got); k++ {
+ prev, cur :=
rank[bits[got[k-1]]], rank[bits[got[k]]]
+ if (order == Ascending && prev
> cur) || (order == Descending && prev < cur) {
+ t.Fatalf("position %d:
bits %#x sorted before %#x", k, bits[got[k-1]], bits[got[k]])
+ }
+ }
+ })
+ }
+ }
+ }
+}
diff --git a/arrow/compute/internal/kernels/vector_sort_support.go
b/arrow/compute/internal/kernels/vector_sort_support.go
index b762ea9a..c08471a6 100644
--- a/arrow/compute/internal/kernels/vector_sort_support.go
+++ b/arrow/compute/internal/kernels/vector_sort_support.go
@@ -136,6 +136,25 @@ func compareFloatNaNs(order SortOrder, viNaN, vjNaN bool)
(int, bool) {
return 0, false
}
+// The totalOrderKey helpers map the raw bits of an IEEE 754 float to a signed
integer whose
+// natural order is the IEEE 754 totalOrder of the floats. Non-negative floats
(sign bit clear)
+// already order correctly as signed integers; for negative floats the
magnitude bits are
+// inverted so larger magnitudes sort first. Arithmetic shift of the sign
yields the mask.
+func totalOrderKey16(bits uint16) int16 {
+ s := int16(bits)
+ return s ^ int16(uint16(s>>15)>>1)
+}
+
+func totalOrderKey32(bits uint32) int32 {
+ s := int32(bits)
+ return s ^ int32(uint32(s>>31)>>1)
+}
+
+func totalOrderKey64(bits uint64) int64 {
+ s := int64(bits)
+ return s ^ int64(uint64(s>>63)>>1)
+}
+
func sortFloat64[T float16.Num | float32 | float64](v T) float64 {
switch x := any(v).(type) {
case float16.Num:
diff --git a/arrow/compute/vector_sort.go b/arrow/compute/vector_sort.go
index bad14e4f..cc29c1ca 100644
--- a/arrow/compute/vector_sort.go
+++ b/arrow/compute/vector_sort.go
@@ -88,7 +88,8 @@ const (
SortNullsAtStart = kernels.NullsAtStart
)
-// SortKey defines a column to sort by with its ordering and null placement
options.
+// SortKey defines a column to sort by with its ordering, null placement and
float ordering options.
+// See kernels.SortKey.FloatTotalOrder for IEEE 754 totalOrder sorting of
float columns.
type SortKey = kernels.SortKey
// SortOptions defines the desired sort order for the input.
diff --git a/arrow/compute/vector_sort_test.go
b/arrow/compute/vector_sort_test.go
index a4922baa..5714cd03 100644
--- a/arrow/compute/vector_sort_test.go
+++ b/arrow/compute/vector_sort_test.go
@@ -20,7 +20,10 @@ package compute_test
import (
"context"
+ "fmt"
"math"
+ "slices"
+ "strconv"
"strings"
"testing"
@@ -1882,3 +1885,167 @@ func TestSortIndicesUUIDLexicographic(t *testing.T) {
{Order: kernels.Descending, NullPlacement:
kernels.NullsAtStart},
}, []uint64{1, 0, 3, 2})
}
+
+// floatTotalOrderLabels is the ascending IEEE 754 totalOrder of
floatTotalOrderInput's non-null
+// values; the input deliberately interleaves +0 and -0 so ties would keep
input order.
+var floatTotalOrderLabels = []string{"-NaN", "-Inf", "-2", "-0", "-0", "0",
"0", "1", "+Inf", "NaN"}
+
+func floatTotalOrderInput() []float64 {
+ negNaN := math.Float64frombits(0xFFF8000000000000)
+ negZero := math.Copysign(0, -1)
+ return []float64{1, math.NaN(), 0, math.Inf(-1), -2, negNaN, negZero,
math.Inf(1), 0, negZero}
+}
+
+func floatLabel(v float64) string {
+ switch {
+ case math.IsNaN(v) && math.Signbit(v):
+ return "-NaN"
+ case math.IsNaN(v):
+ return "NaN"
+ case math.IsInf(v, 1):
+ return "+Inf"
+ case math.IsInf(v, -1):
+ return "-Inf"
+ case v == 0 && math.Signbit(v):
+ return "-0"
+ }
+ return strconv.FormatFloat(v, 'g', -1, 64)
+}
+
+func expectedTotalOrderLabels(key compute.SortKey) []string {
+ vals := slices.Clone(floatTotalOrderLabels)
+ if key.Order == compute.SortOrderDescending {
+ slices.Reverse(vals)
+ }
+ if key.NullPlacement == compute.SortNullsAtStart {
+ return append([]string{"null"}, vals...)
+ }
+ return append(vals, "null")
+}
+
+// buildFloatTotalOrderArray builds the input (plus one null in the middle) as
the given float type.
+func buildFloatTotalOrderArray(t *testing.T, mem memory.Allocator, dt
arrow.DataType, vals []float64) arrow.Array {
+ t.Helper()
+ bldr := array.NewBuilder(mem, dt)
+ defer bldr.Release()
+ for i, v := range vals {
+ if i == len(vals)/2 {
+ bldr.AppendNull()
+ }
+ switch b := bldr.(type) {
+ case *array.Float64Builder:
+ b.Append(v)
+ case *array.Float32Builder:
+ b.Append(float32(v))
+ case *array.Float16Builder:
+ // float16.New goes through float32, which keeps the
NaN sign bit.
+ b.Append(float16.New(float32(v)))
+ }
+ }
+ return bldr.NewArray()
+}
+
+func labelsAt(arr arrow.Array, indices []uint64) []string {
+ out := make([]string, len(indices))
+ for k, i := range indices {
+ if arr.IsNull(int(i)) {
+ out[k] = "null"
+ continue
+ }
+ switch a := arr.(type) {
+ case *array.Float64:
+ out[k] = floatLabel(a.Value(int(i)))
+ case *array.Float32:
+ out[k] = floatLabel(float64(a.Value(int(i))))
+ case *array.Float16:
+ out[k] = floatLabel(float64(a.Value(int(i)).Float32()))
+ }
+ }
+ return out
+}
+
+func floatTotalOrderKeys() []compute.SortKey {
+ var keys []compute.SortKey
+ for _, o := range []kernels.SortOrder{compute.SortOrderAscending,
compute.SortOrderDescending} {
+ for _, np := range
[]kernels.NullPlacement{compute.SortNullsAtEnd, compute.SortNullsAtStart} {
+ keys = append(keys, compute.SortKey{Order: o,
NullPlacement: np, FloatTotalOrder: true})
+ }
+ }
+ return keys
+}
+
+func TestSortIndicesFloatTotalOrder(t *testing.T) {
+ mem := memory.NewGoAllocator()
+ ctx := context.Background()
+ vals := floatTotalOrderInput()
+
+ for _, dt := range []arrow.DataType{arrow.PrimitiveTypes.Float64,
arrow.PrimitiveTypes.Float32, arrow.FixedWidthTypes.Float16} {
+ for _, key := range floatTotalOrderKeys() {
+ want := expectedTotalOrderLabels(key)
+ name := fmt.Sprintf("%s/order=%d/nulls=%d", dt,
key.Order, key.NullPlacement)
+
+ t.Run(name+"/array", func(t *testing.T) {
+ arr := buildFloatTotalOrderArray(t, mem, dt,
vals)
+ defer arr.Release()
+ idx, err := compute.SortIndicesArray(ctx, arr,
key)
+ require.NoError(t, err)
+ defer idx.Release()
+ assert.Equal(t, want, labelsAt(arr,
idx.(*array.Uint64).Uint64Values()))
+ })
+
+ t.Run(name+"/chunked", func(t *testing.T) {
+ // Two chunks exercise the per-chunk sort
followed by the merge.
+ a := buildFloatTotalOrderArray(t, mem, dt,
vals[:5])
+ defer a.Release()
+ b := buildFloatTotalOrderArray(t, mem, dt,
vals[5:])
+ defer b.Release()
+ chunked := arrow.NewChunked(dt,
[]arrow.Array{a, b})
+ defer chunked.Release()
+ flat, err :=
array.Concatenate(chunked.Chunks(), mem)
+ require.NoError(t, err)
+ defer flat.Release()
+ idx, err := compute.SortIndicesChunked(ctx,
chunked, key)
+ require.NoError(t, err)
+ defer idx.Release()
+ got := labelsAt(flat,
idx.(*array.Uint64).Uint64Values())
+ // One null per chunk.
+ if key.NullPlacement ==
compute.SortNullsAtStart {
+ assert.Equal(t,
append([]string{"null"}, want...), got)
+ } else {
+ assert.Equal(t,
append(slices.Clone(want), "null"), got)
+ }
+ })
+
+ t.Run(name+"/record_batch", func(t *testing.T) {
+ // Leading float key with a constant secondary
key exercises the multi-key radix path.
+ arr := buildFloatTotalOrderArray(t, mem, dt,
vals)
+ defer arr.Release()
+ ib := array.NewInt32Builder(mem)
+ defer ib.Release()
+ ib.AppendValues(make([]int32, arr.Len()), nil)
+ ints := ib.NewArray()
+ defer ints.Release()
+ schema := arrow.NewSchema([]arrow.Field{{Name:
"f", Type: dt, Nullable: true}, {Name: "i", Type: arrow.PrimitiveTypes.Int32}},
nil)
+ rec := array.NewRecordBatch(schema,
[]arrow.Array{arr, ints}, int64(arr.Len()))
+ defer rec.Release()
+ k1 := compute.SortKey{ColumnIndex: 1}
+ idx, err := compute.SortIndicesRecordBatch(ctx,
rec, []compute.SortKey{key, k1})
+ require.NoError(t, err)
+ defer idx.Release()
+ assert.Equal(t, want, labelsAt(arr,
idx.(*array.Uint64).Uint64Values()))
+ })
+ }
+ }
+}
+
+func TestSortIndicesFloatDefaultOrderingUnchanged(t *testing.T) {
+ // The default ordering keeps NaNs null-like, ignores their sign, and
treats -0 == +0.
+ mem := memory.NewGoAllocator()
+ arr := buildFloatTotalOrderArray(t, mem, arrow.PrimitiveTypes.Float64,
floatTotalOrderInput())
+ defer arr.Release()
+ idx, err := compute.SortIndicesArray(context.Background(), arr,
compute.SortKey{Order: compute.SortOrderDescending, NullPlacement:
compute.SortNullsAtStart})
+ require.NoError(t, err)
+ defer idx.Release()
+ assert.Equal(t, []string{"null", "NaN", "-NaN", "+Inf", "1", "0", "-0",
"0", "-0", "-2", "-Inf"},
+ labelsAt(arr, idx.(*array.Uint64).Uint64Values()))
+}
