From 2a8e834a54d609234a28fcc587abdcaa901ea981 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Mon, 30 Jun 2025 15:15:43 -0400 Subject: [PATCH 1/4] feat(arrow/compute): support some float16 casts --- arrow/array/float16_builder.go | 4 +- arrow/compute/cast_test.go | 8 +- .../compute/internal/kernels/cast_numeric.go | 46 ++++++++- arrow/compute/internal/kernels/helpers.go | 6 +- .../compute/internal/kernels/numeric_cast.go | 61 +++++++++--- arrow/float16/float16.go | 41 ++++---- arrow/float16/float16_test.go | 98 +++++++++---------- arrow/internal/testing/tools/data_types.go | 2 +- arrow/type_traits.go | 34 +++---- parquet/schema/reflection.go | 2 +- 10 files changed, 194 insertions(+), 108 deletions(-) diff --git a/arrow/array/float16_builder.go b/arrow/array/float16_builder.go index 60cc91b75..9a524113d 100644 --- a/arrow/array/float16_builder.go +++ b/arrow/array/float16_builder.go @@ -87,7 +87,7 @@ func (b *Float16Builder) AppendNulls(n int) { func (b *Float16Builder) AppendEmptyValue() { b.Reserve(1) - b.UnsafeAppend(float16.Num{}) + b.UnsafeAppend(float16.Num(0)) } func (b *Float16Builder) AppendEmptyValues(n int) { @@ -230,7 +230,7 @@ func (b *Float16Builder) UnmarshalOne(dec *json.Decoder) error { default: return &json.UnmarshalTypeError{ Value: fmt.Sprint(t), - Type: reflect.TypeOf(float16.Num{}), + Type: reflect.TypeOf(float16.Num(0)), Offset: dec.InputOffset(), } } diff --git a/arrow/compute/cast_test.go b/arrow/compute/cast_test.go index 4e5f0a54d..370ced16d 100644 --- a/arrow/compute/cast_test.go +++ b/arrow/compute/cast_test.go @@ -574,7 +574,7 @@ func (c *CastSuite) TestToIntDowncastUnsafe() { } func (c *CastSuite) TestFloatingToInt() { - for _, from := range []arrow.DataType{arrow.PrimitiveTypes.Float32, arrow.PrimitiveTypes.Float64} { + for _, from := range []arrow.DataType{arrow.PrimitiveTypes.Float32, arrow.PrimitiveTypes.Float64, arrow.FixedWidthTypes.Float16} { for _, to := range []arrow.DataType{arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int64} { // float to int no truncation c.checkCast(from, to, `[1.0, null, 0.0, -1.0, 5.0]`, `[1, null, 0, -1, 5]`) @@ -590,6 +590,12 @@ func (c *CastSuite) TestFloatingToInt() { } } +func (c *CastSuite) TestFloat16ToFloating() { + for _, to := range []arrow.DataType{arrow.PrimitiveTypes.Float32, arrow.PrimitiveTypes.Float64} { + c.checkCast(arrow.FixedWidthTypes.Float16, to, `[1.5, null, 0.0, -1.5, 5.5]`, `[1.5, null, 0.0, -1.5, 5.5]`) + } +} + func (c *CastSuite) TestIntToFloating() { for _, from := range []arrow.DataType{arrow.PrimitiveTypes.Uint32, arrow.PrimitiveTypes.Int32} { two24 := `[16777216, 16777217]` diff --git a/arrow/compute/internal/kernels/cast_numeric.go b/arrow/compute/internal/kernels/cast_numeric.go index a1772599f..6969d82cb 100644 --- a/arrow/compute/internal/kernels/cast_numeric.go +++ b/arrow/compute/internal/kernels/cast_numeric.go @@ -22,6 +22,7 @@ import ( "unsafe" "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/float16" ) var castNumericUnsafe func(itype, otype arrow.Type, in, out []byte, len int) = castNumericGo @@ -32,7 +33,19 @@ func DoStaticCast[InT, OutT numeric](in []InT, out []OutT) { } } -func reinterpret[T numeric](b []byte, len int) (res []T) { +func DoFloat16Cast[InT numeric](in []InT, out []float16.Num) { + for i, v := range in { + out[i] = float16.New(float32(v)) + } +} + +func DoFloat16CastToNumber[OutT numeric](in []float16.Num, out []OutT) { + for i, v := range in { + out[i] = OutT(v.Float32()) + } +} + +func reinterpret[T numeric | float16.Num](b []byte, len int) (res []T) { return unsafe.Slice((*T)(unsafe.Pointer(&b[0])), len) } @@ -54,6 +67,8 @@ func castNumberToNumberUnsafeImpl[T numeric](outT arrow.Type, in []T, out []byte DoStaticCast(in, reinterpret[int64](out, len(in))) case arrow.UINT64: DoStaticCast(in, reinterpret[uint64](out, len(in))) + case arrow.FLOAT16: + DoFloat16Cast(in, reinterpret[float16.Num](out, len(in))) case arrow.FLOAT32: DoStaticCast(in, reinterpret[float32](out, len(in))) case arrow.FLOAT64: @@ -61,6 +76,33 @@ func castNumberToNumberUnsafeImpl[T numeric](outT arrow.Type, in []T, out []byte } } +func castFloat16ToNumberUnsafeImpl(outT arrow.Type, in []float16.Num, out []byte) { + switch outT { + case arrow.INT8: + DoFloat16CastToNumber(in, reinterpret[int8](out, len(in))) + case arrow.UINT8: + DoFloat16CastToNumber(in, reinterpret[uint8](out, len(in))) + case arrow.INT16: + DoFloat16CastToNumber(in, reinterpret[int16](out, len(in))) + case arrow.UINT16: + DoFloat16CastToNumber(in, reinterpret[uint16](out, len(in))) + case arrow.INT32: + DoFloat16CastToNumber(in, reinterpret[int32](out, len(in))) + case arrow.UINT32: + DoFloat16CastToNumber(in, reinterpret[uint32](out, len(in))) + case arrow.INT64: + DoFloat16CastToNumber(in, reinterpret[int64](out, len(in))) + case arrow.UINT64: + DoFloat16CastToNumber(in, reinterpret[uint64](out, len(in))) + case arrow.FLOAT16: + copy(reinterpret[float16.Num](out, len(in)), in) + case arrow.FLOAT32: + DoFloat16CastToNumber(in, reinterpret[float32](out, len(in))) + case arrow.FLOAT64: + DoFloat16CastToNumber(in, reinterpret[float64](out, len(in))) + } +} + func castNumericGo(itype, otype arrow.Type, in, out []byte, len int) { switch itype { case arrow.INT8: @@ -79,6 +121,8 @@ func castNumericGo(itype, otype arrow.Type, in, out []byte, len int) { castNumberToNumberUnsafeImpl(otype, reinterpret[int64](in, len), out) case arrow.UINT64: castNumberToNumberUnsafeImpl(otype, reinterpret[uint64](in, len), out) + case arrow.FLOAT16: + castFloat16ToNumberUnsafeImpl(otype, reinterpret[float16.Num](in, len), out) case arrow.FLOAT32: castNumberToNumberUnsafeImpl(otype, reinterpret[float32](in, len), out) case arrow.FLOAT64: diff --git a/arrow/compute/internal/kernels/helpers.go b/arrow/compute/internal/kernels/helpers.go index 4a9ead12d..ef5f0bb48 100644 --- a/arrow/compute/internal/kernels/helpers.go +++ b/arrow/compute/internal/kernels/helpers.go @@ -695,7 +695,11 @@ func castNumberToNumberUnsafe(in, out *exec.ArraySpan) { inputOffset := in.Type.(arrow.FixedWidthDataType).Bytes() * int(in.Offset) outputOffset := out.Type.(arrow.FixedWidthDataType).Bytes() * int(out.Offset) - castNumericUnsafe(in.Type.ID(), out.Type.ID(), in.Buffers[1].Buf[inputOffset:], out.Buffers[1].Buf[outputOffset:], int(in.Len)) + if in.Type.ID() == arrow.FLOAT16 || out.Type.ID() == arrow.FLOAT16 { + castNumericGo(in.Type.ID(), out.Type.ID(), in.Buffers[1].Buf[inputOffset:], out.Buffers[1].Buf[outputOffset:], int(in.Len)) + } else { + castNumericUnsafe(in.Type.ID(), out.Type.ID(), in.Buffers[1].Buf[inputOffset:], out.Buffers[1].Buf[outputOffset:], int(in.Len)) + } } func MaxDecimalDigitsForInt(id arrow.Type) (int32, error) { diff --git a/arrow/compute/internal/kernels/numeric_cast.go b/arrow/compute/internal/kernels/numeric_cast.go index 1e76709e1..ea68864ab 100644 --- a/arrow/compute/internal/kernels/numeric_cast.go +++ b/arrow/compute/internal/kernels/numeric_cast.go @@ -28,6 +28,7 @@ import ( "github.com/apache/arrow-go/v18/arrow/compute/exec" "github.com/apache/arrow-go/v18/arrow/decimal128" "github.com/apache/arrow-go/v18/arrow/decimal256" + "github.com/apache/arrow-go/v18/arrow/float16" "github.com/apache/arrow-go/v18/arrow/internal/debug" "github.com/apache/arrow-go/v18/internal/bitutils" "golang.org/x/exp/constraints" @@ -506,7 +507,7 @@ func CastFloat64ToDecimal(ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.E return executor(ctx, batch, out) } -func CastDecimalToFloating[OutT constraints.Float](ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.ExecResult) error { +func CastDecimalToFloating[OutT constraints.Float | float16.Num](ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.ExecResult) error { var ( executor exec.ArrayKernelExec ) @@ -514,20 +515,34 @@ func CastDecimalToFloating[OutT constraints.Float](ctx *exec.KernelCtx, batch *e switch dt := batch.Values[0].Array.Type.(type) { case *arrow.Decimal128Type: scale := dt.Scale - executor = ScalarUnaryNotNull(func(_ *exec.KernelCtx, v decimal128.Num, err *error) OutT { - return OutT(v.ToFloat64(scale)) - }) + switch out.Type.ID() { + case arrow.FLOAT16: + executor = ScalarUnaryNotNull(func(_ *exec.KernelCtx, v decimal128.Num, err *error) OutT { + return OutT(float16.New(v.ToFloat32(scale))) + }) + default: + executor = ScalarUnaryNotNull(func(_ *exec.KernelCtx, v decimal128.Num, err *error) OutT { + return OutT(v.ToFloat64(scale)) + }) + } case *arrow.Decimal256Type: scale := dt.Scale - executor = ScalarUnaryNotNull(func(_ *exec.KernelCtx, v decimal256.Num, err *error) OutT { - return OutT(v.ToFloat64(scale)) - }) + switch out.Type.ID() { + case arrow.FLOAT16: + executor = ScalarUnaryNotNull(func(_ *exec.KernelCtx, v decimal128.Num, err *error) OutT { + return OutT(float16.New(v.ToFloat32(scale))) + }) + default: + executor = ScalarUnaryNotNull(func(_ *exec.KernelCtx, v decimal256.Num, err *error) OutT { + return OutT(v.ToFloat64(scale)) + }) + } } return executor(ctx, batch, out) } -func boolToNum[T numeric](_ *exec.KernelCtx, in []byte, out []T) error { +func boolToNum[T numeric | float16.Num](_ *exec.KernelCtx, in []byte, out []T) error { var ( zero T one = T(1) @@ -543,13 +558,25 @@ func boolToNum[T numeric](_ *exec.KernelCtx, in []byte, out []T) error { return nil } -func checkFloatTrunc[InT constraints.Float, OutT arrow.IntType | arrow.UintType](in, out *exec.ArraySpan) error { - wasTrunc := func(out OutT, in InT) bool { +func wasTrunc[InT constraints.Float | float16.Num, OutT arrow.IntType | arrow.UintType](out OutT, in InT) bool { + switch v := any(in).(type) { + case float16.Num: + return float16.New(float32(out)) != v + default: return InT(out) != in } - wasTruncMaybeNull := func(out OutT, in InT, isValid bool) bool { +} + +func wasTruncMaybeNull[InT constraints.Float | float16.Num, OutT arrow.IntType | arrow.UintType](out OutT, in InT, isValid bool) bool { + switch v := any(in).(type) { + case float16.Num: + return isValid && (float16.New(float32(out)) != v) + default: return isValid && (InT(out) != in) } +} + +func checkFloatTrunc[InT constraints.Float | float16.Num, OutT arrow.IntType | arrow.UintType](in, out *exec.ArraySpan) error { getError := func(val InT) error { return fmt.Errorf("%w: float value %f was truncated converting to %s", arrow.ErrInvalid, val, out.Type) @@ -598,7 +625,7 @@ func checkFloatTrunc[InT constraints.Float, OutT arrow.IntType | arrow.UintType] return nil } -func checkFloatToIntTruncImpl[T constraints.Float](in, out *exec.ArraySpan) error { +func checkFloatToIntTruncImpl[T constraints.Float | float16.Num](in, out *exec.ArraySpan) error { switch out.Type.ID() { case arrow.INT8: return checkFloatTrunc[T, int8](in, out) @@ -623,6 +650,8 @@ func checkFloatToIntTruncImpl[T constraints.Float](in, out *exec.ArraySpan) erro func checkFloatToIntTrunc(in, out *exec.ArraySpan) error { switch in.Type.ID() { + case arrow.FLOAT16: + return checkFloatToIntTruncImpl[float16.Num](in, out) case arrow.FLOAT32: return checkFloatToIntTruncImpl[float32](in, out) case arrow.FLOAT64: @@ -729,7 +758,7 @@ func getParseStringExec[OffsetT int32 | int64](out arrow.Type) exec.ArrayKernelE panic("invalid type for getParseStringExec") } -func addCommonNumberCasts[T numeric](outTy arrow.DataType, kernels []exec.ScalarKernel) []exec.ScalarKernel { +func addCommonNumberCasts[T numeric | float16.Num](outTy arrow.DataType, kernels []exec.ScalarKernel) []exec.ScalarKernel { kernels = append(kernels, GetCommonCastKernels(outTy.ID(), exec.NewOutputType(outTy))...) kernels = append(kernels, exec.NewScalarKernel( @@ -759,7 +788,7 @@ func GetCastToInteger[T arrow.IntType | arrow.UintType](outType arrow.DataType) CastIntToInt, nil)) } - for _, inTy := range floatingTypes { + for _, inTy := range append(floatingTypes, arrow.FixedWidthTypes.Float16) { kernels = append(kernels, exec.NewScalarKernel( []exec.InputType{exec.NewExactInput(inTy)}, output, CastFloatingToInteger, nil)) @@ -775,7 +804,7 @@ func GetCastToInteger[T arrow.IntType | arrow.UintType](outType arrow.DataType) return kernels } -func GetCastToFloating[T constraints.Float](outType arrow.DataType) []exec.ScalarKernel { +func GetCastToFloating[T constraints.Float | float16.Num](outType arrow.DataType) []exec.ScalarKernel { kernels := make([]exec.ScalarKernel, 0) output := exec.NewOutputType(outType) @@ -785,7 +814,7 @@ func GetCastToFloating[T constraints.Float](outType arrow.DataType) []exec.Scala CastIntegerToFloating, nil)) } - for _, inTy := range floatingTypes { + for _, inTy := range append(floatingTypes, arrow.FixedWidthTypes.Float16) { kernels = append(kernels, exec.NewScalarKernel( []exec.InputType{exec.NewExactInput(inTy)}, output, CastFloatingToFloating, nil)) diff --git a/arrow/float16/float16.go b/arrow/float16/float16.go index 0aa4df8a0..e0ac69334 100644 --- a/arrow/float16/float16.go +++ b/arrow/float16/float16.go @@ -18,6 +18,7 @@ package float16 import ( "encoding/binary" + "fmt" "math" "strconv" ) @@ -26,12 +27,10 @@ import ( // stored on 16 bits. // // See https://en.wikipedia.org/wiki/Half-precision_floating-point_format for more informations. -type Num struct { - bits uint16 -} +type Num uint16 var ( - MaxNum = Num{bits: 0b0111101111111111} + MaxNum = Num(0b0111101111111111) MinNum = MaxNum.Negate() ) @@ -55,14 +54,18 @@ func New(f float32) Num { res = 0 fc = 0 } - return Num{bits: (sn << 15) | uint16(res<<10) | fc} + return Num((sn << 15) | uint16(res<<10) | fc) +} + +func (f Num) Format(s fmt.State, verb rune) { + fmt.Fprintf(s, fmt.FormatString(s, verb), f.Float32()) } func (f Num) Float32() float32 { - sn := uint32((f.bits >> 15) & 0x1) - exp := (f.bits >> 10) & 0x1f + sn := uint32((f >> 15) & 0x1) + exp := (f >> 10) & 0x1f res := uint32(exp) + 127 - 15 - fc := uint32(f.bits & 0x3ff) + fc := uint32(f & 0x3ff) switch exp { case 0: res = 0 @@ -73,7 +76,7 @@ func (f Num) Float32() float32 { } func (n Num) Negate() Num { - return Num{bits: n.bits ^ 0x8000} + return Num(n ^ 0x8000) } func (n Num) Add(rhs Num) Num { @@ -171,29 +174,29 @@ func (n Num) Sign() int { return 1 } -func (n Num) Signbit() bool { return (n.bits & 0x8000) != 0 } +func (n Num) Signbit() bool { return (n & 0x8000) != 0 } -func (n Num) IsNaN() bool { return (n.bits & 0x7fff) > 0x7c00 } +func (n Num) IsNaN() bool { return (n & 0x7fff) > 0x7c00 } -func (n Num) IsInf() bool { return (n.bits & 0x7c00) == 0x7c00 } +func (n Num) IsInf() bool { return (n & 0x7c00) == 0x7c00 } -func (n Num) IsZero() bool { return (n.bits & 0x7fff) == 0 } +func (n Num) IsZero() bool { return (n & 0x7fff) == 0 } -func (f Num) Uint16() uint16 { return f.bits } +func (f Num) Uint16() uint16 { return uint16(f) } func (f Num) String() string { return strconv.FormatFloat(float64(f.Float32()), 'g', -1, 32) } -func Inf() Num { return Num{bits: 0x7c00} } +func Inf() Num { return Num(0x7c00) } -func NaN() Num { return Num{bits: 0x7fff} } +func NaN() Num { return Num(0x7fff) } -func FromBits(src uint16) Num { return Num{bits: src} } +func FromBits(src uint16) Num { return Num(src) } func FromLEBytes(src []byte) Num { - return Num{bits: binary.LittleEndian.Uint16(src)} + return Num(binary.LittleEndian.Uint16(src)) } func (f Num) PutLEBytes(dst []byte) { - binary.LittleEndian.PutUint16(dst, f.bits) + binary.LittleEndian.PutUint16(dst, uint16(f)) } func (f Num) ToLEBytes() []byte { diff --git a/arrow/float16/float16_test.go b/arrow/float16/float16_test.go index cfde440c5..271b70708 100644 --- a/arrow/float16/float16_test.go +++ b/arrow/float16/float16_test.go @@ -25,20 +25,20 @@ import ( func TestFloat16(t *testing.T) { cases := map[Num]float32{ - {bits: 0x3c00}: 1, - {bits: 0x4000}: 2, - {bits: 0xc000}: -2, - {bits: 0x0000}: 0, - {bits: 0x5b8f}: 241.875, - {bits: 0xdb8f}: -241.875, - {bits: 0x48c8}: 9.5625, - {bits: 0xc8c8}: -9.5625, + 0x3c00: 1, + 0x4000: 2, + 0xc000: -2, + 0x0000: 0, + 0x5b8f: 241.875, + 0xdb8f: -241.875, + 0x48c8: 9.5625, + 0xc8c8: -9.5625, } for k, v := range cases { f := k.Float32() assert.Equal(t, v, f, "float32 values should be the same") i := New(v) - assert.Equal(t, k.bits, i.bits, "float16 values should be the same") + assert.Equal(t, k, i, "float16 values should be the same") assert.Equal(t, k.Uint16(), i.Uint16(), "float16 values should be the same") assert.Equal(t, k.String(), fmt.Sprintf("%v", v), "string representation differ") } @@ -50,9 +50,9 @@ func TestAdd(t *testing.T) { rhs Num want Num }{ - {Num{bits: 0x0000}, Num{bits: 0x0000}, Num{bits: 0x0000}}, // 0 + 0 = 0 - {Num{bits: 0x3c00}, Num{bits: 0x4000}, Num{bits: 0x4200}}, // 1 + 2 = 3 - {Num{bits: 0x4248}, Num{bits: 0x3245}, Num{bits: 0x42AC}}, // 3.141 + 0.196 = 3.336 + {Num(0x0000), Num(0x0000), Num(0x0000)}, // 0 + 0 = 0 + {Num(0x3c00), Num(0x4000), Num(0x4200)}, // 1 + 2 = 3 + {Num(0x4248), Num(0x3245), Num(0x42AC)}, // 3.141 + 0.196 = 3.336 } { t.Run("add", func(t *testing.T) { n := tc.n.Add(tc.rhs) @@ -69,9 +69,9 @@ func TestSub(t *testing.T) { rhs Num want Num }{ - {Num{bits: 0x0000}, Num{bits: 0x0000}, Num{bits: 0x0000}}, // 0 - 0 = 0 - {Num{bits: 0x3c00}, Num{bits: 0x4000}, Num{bits: 0xBC00}}, // 1 - 2 = -1 - {Num{bits: 0x4248}, Num{bits: 0x3245}, Num{bits: 0x41E3}}, // 3.141 - 0.196 = 2.944 + {Num(0x0000), Num(0x0000), Num(0x0000)}, // 0 - 0 = 0 + {Num(0x3c00), Num(0x4000), Num(0xBC00)}, // 1 - 2 = -1 + {Num(0x4248), Num(0x3245), Num(0x41E3)}, // 3.141 - 0.196 = 2.944 } { t.Run("sub", func(t *testing.T) { n := tc.n.Sub(tc.rhs) @@ -88,9 +88,9 @@ func TestMul(t *testing.T) { rhs Num want Num }{ - {Num{bits: 0x0000}, Num{bits: 0x0000}, Num{bits: 0x0000}}, // 0 * 0 = 0 - {Num{bits: 0x3c00}, Num{bits: 0x4000}, Num{bits: 0x4000}}, // 1 * 2 = 2 - {Num{bits: 0x4248}, Num{bits: 0x3245}, Num{bits: 0x38EC}}, // 3.141 * 0.196 = 0.6153 + {Num(0x0000), Num(0x0000), Num(0x0000)}, // 0 * 0 = 0 + {Num(0x3c00), Num(0x4000), Num(0x4000)}, // 1 * 2 = 2 + {Num(0x4248), Num(0x3245), Num(0x38EC)}, // 3.141 * 0.196 = 0.6153 } { t.Run("mul", func(t *testing.T) { n := tc.n.Mul(tc.rhs) @@ -107,9 +107,9 @@ func TestDiv(t *testing.T) { rhs Num want Num }{ - {Num{bits: 0x0000}, Num{bits: 0x3c00}, Num{bits: 0x0000}}, // 0 / 1 = 0 - {Num{bits: 0x3c00}, Num{bits: 0x4000}, Num{bits: 0x3800}}, // 1 / 2 = 0.5 - {Num{bits: 0x4248}, Num{bits: 0x3245}, Num{bits: 0x4C01}}, // 3.141 * 0.196 = 16.02 + {Num(0x0000), Num(0x3c00), Num(0x0000)}, // 0 / 1 = 0 + {Num(0x3c00), Num(0x4000), Num(0x3800)}, // 1 / 2 = 0.5 + {Num(0x4248), Num(0x3245), Num(0x4C01)}, // 3.141 * 0.196 = 16.02 } { t.Run("div", func(t *testing.T) { n := tc.n.Div(tc.rhs) @@ -126,9 +126,9 @@ func TestGreater(t *testing.T) { rhs Num want bool }{ - {Num{bits: 0x3c00}, Num{bits: 0x4000}, false}, // 1 > 2 = false - {Num{bits: 0x4900}, Num{bits: 0x4900}, false}, // 10 == 10 = false - {Num{bits: 0x4248}, Num{bits: 0x3245}, true}, // 3.141 > 0.196 = true + {Num(0x3c00), Num(0x4000), false}, // 1 > 2 = false + {Num(0x4900), Num(0x4900), false}, // 10 == 10 = false + {Num(0x4248), Num(0x3245), true}, // 3.141 > 0.196 = true } { t.Run("greater", func(t *testing.T) { n := tc.n.Greater(tc.rhs) @@ -145,9 +145,9 @@ func TestLess(t *testing.T) { rhs Num want bool }{ - {Num{bits: 0x3c00}, Num{bits: 0x4000}, true}, // 1 < 2 = true - {Num{bits: 0x4900}, Num{bits: 0x4900}, false}, // 10 == 10 = false - {Num{bits: 0x4248}, Num{bits: 0x3245}, false}, // 3.141 < 0.196 = false + {Num(0x3c00), Num(0x4000), true}, // 1 < 2 = true + {Num(0x4900), Num(0x4900), false}, // 10 == 10 = false + {Num(0x4248), Num(0x3245), false}, // 3.141 < 0.196 = false } { t.Run("less", func(t *testing.T) { n := tc.n.Less(tc.rhs) @@ -164,9 +164,9 @@ func TestCmp(t *testing.T) { rhs Num want int }{ - {Num{bits: 0x3c00}, Num{bits: 0x4000}, -1}, // cmp(1, 2) = -1 - {Num{bits: 0x4900}, Num{bits: 0x4900}, 0}, // cmp(10, 10) = 0 - {Num{bits: 0x4248}, Num{bits: 0x3245}, 1}, // cmp(3.141, 0.196) = 1 + {Num(0x3c00), Num(0x4000), -1}, // cmp(1, 2) = -1 + {Num(0x4900), Num(0x4900), 0}, // cmp(10, 10) = 0 + {Num(0x4248), Num(0x3245), 1}, // cmp(3.141, 0.196) = 1 } { t.Run("cmp", func(t *testing.T) { n := tc.n.Cmp(tc.rhs) @@ -183,8 +183,8 @@ func TestMax(t *testing.T) { rhs []Num want Num }{ - {Num{bits: 0x3c00}, []Num{{bits: 0x4000}, {bits: 0x4580}, {bits: 0x3C00}, {bits: 0x4247}}, Num{bits: 0x4580}}, // max(2, 5.5, 1, 3.14) = 5.5 - {Num{bits: 0x4248}, []Num{{bits: 0xC000}, {bits: 0xC580}, {bits: 0x3C00}, {bits: 0x4247}}, Num{bits: 0x4248}}, // max(-2, -5.5, 1, 3.14) = 3.14 + {Num(0x3c00), []Num{Num(0x4000), Num(0x4580), Num(0x3C00), Num(0x4247)}, Num(0x4580)}, // max(2, 5.5, 1, 3.14) = 5.5 + {Num(0x4248), []Num{Num(0xC000), Num(0xC580), Num(0x3C00), Num(0x4247)}, Num(0x4248)}, // max(-2, -5.5, 1, 3.14) = 3.14 } { t.Run("max", func(t *testing.T) { n := Max(tc.n, tc.rhs...) @@ -201,8 +201,8 @@ func TestMin(t *testing.T) { rhs []Num want Num }{ - {Num{bits: 0x3c00}, []Num{{bits: 0x4000}, {bits: 0x4580}, {bits: 0x3C00}, {bits: 0x4247}}, Num{bits: 0x3C00}}, // min(2, 5.5, 1, 3.14) = 1 - {Num{bits: 0x4248}, []Num{{bits: 0x4000}, {bits: 0xC580}, {bits: 0xBC00}, {bits: 0x4247}}, Num{bits: 0xC580}}, // min(2, -5.5, -1, 3.14) = -5.5 + {Num(0x3c00), []Num{Num(0x4000), Num(0x4580), Num(0x3C00), Num(0x4247)}, Num(0x3C00)}, // min(2, 5.5, 1, 3.14) = 1 + {Num(0x4248), []Num{Num(0x4000), Num(0xC580), Num(0xBC00), Num(0x4247)}, Num(0xC580)}, // min(2, -5.5, -1, 3.14) = -5.5 } { t.Run("min", func(t *testing.T) { n := Min(tc.n, tc.rhs...) @@ -218,9 +218,9 @@ func TestAbs(t *testing.T) { n Num want Num }{ - {Num{bits: 0x4580}, Num{bits: 0x4580}}, // 5.5 - {Num{bits: 0x0000}, Num{bits: 0x0000}}, // 0 - {Num{bits: 0xC580}, Num{bits: 0x4580}}, // -5.5 + {Num(0x4580), Num(0x4580)}, // 5.5 + {Num(0x0000), Num(0x0000)}, // 0 + {Num(0xC580), Num(0x4580)}, // -5.5 } { t.Run("abs", func(t *testing.T) { n := tc.n.Abs() @@ -236,10 +236,10 @@ func TestSign(t *testing.T) { n Num want int }{ - {Num{bits: 0x4580}, 1}, // 5.5 - {Num{bits: 0x0000}, 0}, // 0 - {Num{bits: 0x8000}, 0}, // -0 - {Num{bits: 0xC580}, -1}, // -5.5 + {Num(0x4580), 1}, // 5.5 + {Num(0x0000), 0}, // 0 + {Num(0x8000), 0}, // -0 + {Num(0xC580), -1}, // -5.5 } { t.Run("sign", func(t *testing.T) { n := tc.n.Sign() @@ -255,10 +255,10 @@ func TestSignbit(t *testing.T) { n Num want bool }{ - {Num{bits: 0x4580}, false}, // 5.5 - {Num{bits: 0x0000}, false}, // 0 - {Num{bits: 0x8000}, true}, // -0 - {Num{bits: 0xC580}, true}, // -5.5 + {Num(0x4580), false}, // 5.5 + {Num(0x0000), false}, // 0 + {Num(0x8000), true}, // -0 + {Num(0xC580), true}, // -5.5 } { t.Run("signbit", func(t *testing.T) { n := tc.n.Signbit() @@ -278,10 +278,10 @@ func TestIsNaN(t *testing.T) { {NaN().Negate(), true}, {Inf(), false}, {Inf().Negate(), false}, - {Num{bits: 0x7c01}, true}, // nan - {Num{bits: 0xfc01}, true}, // -nan - {Num{bits: 0x7e00}, true}, // nan - {Num{bits: 0xfe00}, true}, // -nan + {Num(0x7c01), true}, // nan + {Num(0xfc01), true}, // -nan + {Num(0x7e00), true}, // nan + {Num(0xfe00), true}, // -nan } { t.Run("isnan", func(t *testing.T) { n := tc.n.IsNaN() diff --git a/arrow/internal/testing/tools/data_types.go b/arrow/internal/testing/tools/data_types.go index 5ffd44b48..f494c1ffd 100644 --- a/arrow/internal/testing/tools/data_types.go +++ b/arrow/internal/testing/tools/data_types.go @@ -42,7 +42,7 @@ var typMap = map[reflect.Type]arrow.DataType{ reflect.TypeOf(arrow.Date32(0)): arrow.FixedWidthTypes.Date32, reflect.TypeOf(arrow.Date64(0)): arrow.FixedWidthTypes.Date64, reflect.TypeOf(true): arrow.FixedWidthTypes.Boolean, - reflect.TypeOf(float16.Num{}): arrow.FixedWidthTypes.Float16, + reflect.TypeOf(float16.Num(0)): arrow.FixedWidthTypes.Float16, reflect.TypeOf([]byte{}): arrow.BinaryTypes.Binary, } diff --git a/arrow/type_traits.go b/arrow/type_traits.go index 7185ef25a..1eb30eb2e 100644 --- a/arrow/type_traits.go +++ b/arrow/type_traits.go @@ -124,23 +124,23 @@ func GetData[T FixedWidthType | ViewHeader](in []byte) []T { } var typMap = map[reflect.Type]DataType{ - reflect.TypeOf(false): FixedWidthTypes.Boolean, - reflect.TypeOf(int8(0)): PrimitiveTypes.Int8, - reflect.TypeOf(int16(0)): PrimitiveTypes.Int16, - reflect.TypeOf(int32(0)): PrimitiveTypes.Int32, - reflect.TypeOf(int64(0)): PrimitiveTypes.Int64, - reflect.TypeOf(uint8(0)): PrimitiveTypes.Uint8, - reflect.TypeOf(uint16(0)): PrimitiveTypes.Uint16, - reflect.TypeOf(uint32(0)): PrimitiveTypes.Uint32, - reflect.TypeOf(uint64(0)): PrimitiveTypes.Uint64, - reflect.TypeOf(float32(0)): PrimitiveTypes.Float32, - reflect.TypeOf(float64(0)): PrimitiveTypes.Float64, - reflect.TypeOf(string("")): BinaryTypes.String, - reflect.TypeOf(Date32(0)): FixedWidthTypes.Date32, - reflect.TypeOf(Date64(0)): FixedWidthTypes.Date64, - reflect.TypeOf(true): FixedWidthTypes.Boolean, - reflect.TypeOf(float16.Num{}): FixedWidthTypes.Float16, - reflect.TypeOf([]byte{}): BinaryTypes.Binary, + reflect.TypeOf(false): FixedWidthTypes.Boolean, + reflect.TypeOf(int8(0)): PrimitiveTypes.Int8, + reflect.TypeOf(int16(0)): PrimitiveTypes.Int16, + reflect.TypeOf(int32(0)): PrimitiveTypes.Int32, + reflect.TypeOf(int64(0)): PrimitiveTypes.Int64, + reflect.TypeOf(uint8(0)): PrimitiveTypes.Uint8, + reflect.TypeOf(uint16(0)): PrimitiveTypes.Uint16, + reflect.TypeOf(uint32(0)): PrimitiveTypes.Uint32, + reflect.TypeOf(uint64(0)): PrimitiveTypes.Uint64, + reflect.TypeOf(float32(0)): PrimitiveTypes.Float32, + reflect.TypeOf(float64(0)): PrimitiveTypes.Float64, + reflect.TypeOf(string("")): BinaryTypes.String, + reflect.TypeOf(Date32(0)): FixedWidthTypes.Date32, + reflect.TypeOf(Date64(0)): FixedWidthTypes.Date64, + reflect.TypeOf(true): FixedWidthTypes.Boolean, + reflect.TypeOf(float16.Num(0)): FixedWidthTypes.Float16, + reflect.TypeOf([]byte{}): BinaryTypes.Binary, } // GetDataType returns the appropriate DataType for the given type T diff --git a/parquet/schema/reflection.go b/parquet/schema/reflection.go index 6082c2988..e7c5252e4 100644 --- a/parquet/schema/reflection.go +++ b/parquet/schema/reflection.go @@ -376,7 +376,7 @@ func typeToNode(name string, typ reflect.Type, repType parquet.Repetition, info } return Must(MapOf(name, key, value, repType, fieldID)) case reflect.Struct: - if typ == reflect.TypeOf(float16.Num{}) { + if typ == reflect.TypeOf(float16.Num(0)) { return MustPrimitive(NewPrimitiveNodeLogical(name, repType, Float16LogicalType{}, parquet.Types.FixedLenByteArray, 2, fieldID)) } // structs are Group nodes From 02b67d7f3695edc6d39bcb5270adaea11df0b694 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Mon, 30 Jun 2025 16:15:59 -0400 Subject: [PATCH 2/4] fix schema logical types --- parquet/schema/reflection.go | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/parquet/schema/reflection.go b/parquet/schema/reflection.go index e7c5252e4..9c6738658 100644 --- a/parquet/schema/reflection.go +++ b/parquet/schema/reflection.go @@ -376,9 +376,6 @@ func typeToNode(name string, typ reflect.Type, repType parquet.Repetition, info } return Must(MapOf(name, key, value, repType, fieldID)) case reflect.Struct: - if typ == reflect.TypeOf(float16.Num(0)) { - return MustPrimitive(NewPrimitiveNodeLogical(name, repType, Float16LogicalType{}, parquet.Types.FixedLenByteArray, 2, fieldID)) - } // structs are Group nodes fields := make(FieldList, 0) for i := 0; i < typ.NumField(); i++ { @@ -502,6 +499,10 @@ func typeToNode(name string, typ reflect.Type, repType parquet.Repetition, info return MustPrimitive(NewPrimitiveNodeLogical(name, repType, NewIntLogicalType(bitwidth, true), ptyp, 0, fieldID)) case reflect.Uint, reflect.Uint32, reflect.Uint8, reflect.Uint16, reflect.Uint64: + if typ == reflect.TypeOf(float16.Num(0)) { + return MustPrimitive(NewPrimitiveNodeLogical(name, repType, Float16LogicalType{}, parquet.Types.FixedLenByteArray, 2, fieldID)) + } + // handle unsigned integer types and default to the corresponding logical type for it. ptyp := parquet.Types.Int32 if typ.Bits() == 64 { From b2798b37b9411c286a97cff00f8470c945369a97 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Thu, 3 Jul 2025 12:41:39 -0400 Subject: [PATCH 3/4] avoid breaking change --- arrow/array/float16_builder.go | 4 +- .../compute/internal/kernels/numeric_cast.go | 136 +++++++++++++----- arrow/float16/float16.go | 36 ++--- arrow/float16/float16_test.go | 96 ++++++------- arrow/internal/testing/tools/data_types.go | 2 +- arrow/type_traits.go | 34 ++--- parquet/schema/reflection.go | 7 +- 7 files changed, 194 insertions(+), 121 deletions(-) diff --git a/arrow/array/float16_builder.go b/arrow/array/float16_builder.go index 9a524113d..60cc91b75 100644 --- a/arrow/array/float16_builder.go +++ b/arrow/array/float16_builder.go @@ -87,7 +87,7 @@ func (b *Float16Builder) AppendNulls(n int) { func (b *Float16Builder) AppendEmptyValue() { b.Reserve(1) - b.UnsafeAppend(float16.Num(0)) + b.UnsafeAppend(float16.Num{}) } func (b *Float16Builder) AppendEmptyValues(n int) { @@ -230,7 +230,7 @@ func (b *Float16Builder) UnmarshalOne(dec *json.Decoder) error { default: return &json.UnmarshalTypeError{ Value: fmt.Sprint(t), - Type: reflect.TypeOf(float16.Num(0)), + Type: reflect.TypeOf(float16.Num{}), Offset: dec.InputOffset(), } } diff --git a/arrow/compute/internal/kernels/numeric_cast.go b/arrow/compute/internal/kernels/numeric_cast.go index ea68864ab..7681b02e2 100644 --- a/arrow/compute/internal/kernels/numeric_cast.go +++ b/arrow/compute/internal/kernels/numeric_cast.go @@ -507,7 +507,7 @@ func CastFloat64ToDecimal(ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.E return executor(ctx, batch, out) } -func CastDecimalToFloating[OutT constraints.Float | float16.Num](ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.ExecResult) error { +func CastDecimalToFloat16(ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.ExecResult) error { var ( executor exec.ArrayKernelExec ) @@ -515,34 +515,41 @@ func CastDecimalToFloating[OutT constraints.Float | float16.Num](ctx *exec.Kerne switch dt := batch.Values[0].Array.Type.(type) { case *arrow.Decimal128Type: scale := dt.Scale - switch out.Type.ID() { - case arrow.FLOAT16: - executor = ScalarUnaryNotNull(func(_ *exec.KernelCtx, v decimal128.Num, err *error) OutT { - return OutT(float16.New(v.ToFloat32(scale))) - }) - default: - executor = ScalarUnaryNotNull(func(_ *exec.KernelCtx, v decimal128.Num, err *error) OutT { - return OutT(v.ToFloat64(scale)) - }) - } + executor = ScalarUnaryNotNull(func(_ *exec.KernelCtx, v decimal128.Num, err *error) float16.Num { + return float16.New(v.ToFloat32(scale)) + }) case *arrow.Decimal256Type: scale := dt.Scale - switch out.Type.ID() { - case arrow.FLOAT16: - executor = ScalarUnaryNotNull(func(_ *exec.KernelCtx, v decimal128.Num, err *error) OutT { - return OutT(float16.New(v.ToFloat32(scale))) - }) - default: - executor = ScalarUnaryNotNull(func(_ *exec.KernelCtx, v decimal256.Num, err *error) OutT { - return OutT(v.ToFloat64(scale)) - }) - } + executor = ScalarUnaryNotNull(func(_ *exec.KernelCtx, v decimal256.Num, err *error) float16.Num { + return float16.New(v.ToFloat32(scale)) + }) } return executor(ctx, batch, out) } -func boolToNum[T numeric | float16.Num](_ *exec.KernelCtx, in []byte, out []T) error { +func CastDecimalToFloating[OutT constraints.Float](ctx *exec.KernelCtx, batch *exec.ExecSpan, out *exec.ExecResult) error { + var ( + executor exec.ArrayKernelExec + ) + + switch dt := batch.Values[0].Array.Type.(type) { + case *arrow.Decimal128Type: + scale := dt.Scale + executor = ScalarUnaryNotNull(func(_ *exec.KernelCtx, v decimal128.Num, err *error) OutT { + return OutT(v.ToFloat64(scale)) + }) + case *arrow.Decimal256Type: + scale := dt.Scale + executor = ScalarUnaryNotNull(func(_ *exec.KernelCtx, v decimal256.Num, err *error) OutT { + return OutT(v.ToFloat64(scale)) + }) + } + + return executor(ctx, batch, out) +} + +func boolToNum[T numeric](_ *exec.KernelCtx, in []byte, out []T) error { var ( zero T one = T(1) @@ -558,12 +565,32 @@ func boolToNum[T numeric | float16.Num](_ *exec.KernelCtx, in []byte, out []T) e return nil } +func boolToFloat16(_ *exec.KernelCtx, in []byte, out []float16.Num) error { + var ( + zero float16.Num + one = float16.New(1) + ) + + for i := range out { + if bitutil.BitIsSet(in, i) { + out[i] = one + } else { + out[i] = zero + } + } + return nil +} + func wasTrunc[InT constraints.Float | float16.Num, OutT arrow.IntType | arrow.UintType](out OutT, in InT) bool { switch v := any(in).(type) { case float16.Num: return float16.New(float32(out)) != v + case float32: + return float32(out) != v + case float64: + return float64(out) != v default: - return InT(out) != in + return false } } @@ -571,8 +598,12 @@ func wasTruncMaybeNull[InT constraints.Float | float16.Num, OutT arrow.IntType | switch v := any(in).(type) { case float16.Num: return isValid && (float16.New(float32(out)) != v) + case float32: + return isValid && (float32(out) != v) + case float64: + return isValid && (float64(out) != v) default: - return isValid && (InT(out) != in) + return false } } @@ -758,7 +789,27 @@ func getParseStringExec[OffsetT int32 | int64](out arrow.Type) exec.ArrayKernelE panic("invalid type for getParseStringExec") } -func addCommonNumberCasts[T numeric | float16.Num](outTy arrow.DataType, kernels []exec.ScalarKernel) []exec.ScalarKernel { +func addFloat16Casts(outTy arrow.DataType, kernels []exec.ScalarKernel) []exec.ScalarKernel { + kernels = append(kernels, GetCommonCastKernels(outTy.ID(), exec.NewOutputType(outTy))...) + + kernels = append(kernels, exec.NewScalarKernel( + []exec.InputType{exec.NewExactInput(arrow.FixedWidthTypes.Boolean)}, + exec.NewOutputType(outTy), ScalarUnaryBoolArg(boolToFloat16), nil)) + + for _, inTy := range []arrow.DataType{arrow.BinaryTypes.Binary, arrow.BinaryTypes.String} { + kernels = append(kernels, exec.NewScalarKernel( + []exec.InputType{exec.NewExactInput(inTy)}, exec.NewOutputType(outTy), + getParseStringExec[int32](outTy.ID()), nil)) + } + for _, inTy := range []arrow.DataType{arrow.BinaryTypes.LargeBinary, arrow.BinaryTypes.LargeString} { + kernels = append(kernels, exec.NewScalarKernel( + []exec.InputType{exec.NewExactInput(inTy)}, exec.NewOutputType(outTy), + getParseStringExec[int64](outTy.ID()), nil)) + } + return kernels +} + +func addCommonNumberCasts[T numeric](outTy arrow.DataType, kernels []exec.ScalarKernel) []exec.ScalarKernel { kernels = append(kernels, GetCommonCastKernels(outTy.ID(), exec.NewOutputType(outTy))...) kernels = append(kernels, exec.NewScalarKernel( @@ -820,13 +871,34 @@ func GetCastToFloating[T constraints.Float | float16.Num](outType arrow.DataType CastFloatingToFloating, nil)) } - kernels = addCommonNumberCasts[T](outType, kernels) - kernels = append(kernels, exec.NewScalarKernel( - []exec.InputType{exec.NewIDInput(arrow.DECIMAL128)}, output, - CastDecimalToFloating[T], nil)) - kernels = append(kernels, exec.NewScalarKernel( - []exec.InputType{exec.NewIDInput(arrow.DECIMAL256)}, output, - CastDecimalToFloating[T], nil)) + var z T + switch any(z).(type) { + case float16.Num: + kernels = addFloat16Casts(outType, kernels) + kernels = append(kernels, exec.NewScalarKernel( + []exec.InputType{exec.NewIDInput(arrow.DECIMAL128)}, output, + CastDecimalToFloat16, nil)) + kernels = append(kernels, exec.NewScalarKernel( + []exec.InputType{exec.NewIDInput(arrow.DECIMAL256)}, output, + CastDecimalToFloat16, nil)) + case float32: + kernels = addCommonNumberCasts[float32](outType, kernels) + kernels = append(kernels, exec.NewScalarKernel( + []exec.InputType{exec.NewIDInput(arrow.DECIMAL128)}, output, + CastDecimalToFloating[float32], nil)) + kernels = append(kernels, exec.NewScalarKernel( + []exec.InputType{exec.NewIDInput(arrow.DECIMAL256)}, output, + CastDecimalToFloating[float32], nil)) + case float64: + kernels = addCommonNumberCasts[float64](outType, kernels) + kernels = append(kernels, exec.NewScalarKernel( + []exec.InputType{exec.NewIDInput(arrow.DECIMAL128)}, output, + CastDecimalToFloating[float64], nil)) + kernels = append(kernels, exec.NewScalarKernel( + []exec.InputType{exec.NewIDInput(arrow.DECIMAL256)}, output, + CastDecimalToFloating[float64], nil)) + } + return kernels } diff --git a/arrow/float16/float16.go b/arrow/float16/float16.go index e0ac69334..f6c276b6c 100644 --- a/arrow/float16/float16.go +++ b/arrow/float16/float16.go @@ -27,10 +27,12 @@ import ( // stored on 16 bits. // // See https://en.wikipedia.org/wiki/Half-precision_floating-point_format for more informations. -type Num uint16 +type Num struct { + bits uint16 +} var ( - MaxNum = Num(0b0111101111111111) + MaxNum = Num{bits: 0b0111101111111111} MinNum = MaxNum.Negate() ) @@ -54,7 +56,7 @@ func New(f float32) Num { res = 0 fc = 0 } - return Num((sn << 15) | uint16(res<<10) | fc) + return Num{bits: (sn << 15) | uint16(res<<10) | fc} } func (f Num) Format(s fmt.State, verb rune) { @@ -62,10 +64,10 @@ func (f Num) Format(s fmt.State, verb rune) { } func (f Num) Float32() float32 { - sn := uint32((f >> 15) & 0x1) - exp := (f >> 10) & 0x1f + sn := uint32((f.bits >> 15) & 0x1) + exp := (f.bits >> 10) & 0x1f res := uint32(exp) + 127 - 15 - fc := uint32(f & 0x3ff) + fc := uint32(f.bits & 0x3ff) switch exp { case 0: res = 0 @@ -76,7 +78,7 @@ func (f Num) Float32() float32 { } func (n Num) Negate() Num { - return Num(n ^ 0x8000) + return Num{bits: n.bits ^ 0x8000} } func (n Num) Add(rhs Num) Num { @@ -174,29 +176,29 @@ func (n Num) Sign() int { return 1 } -func (n Num) Signbit() bool { return (n & 0x8000) != 0 } +func (n Num) Signbit() bool { return (n.bits & 0x8000) != 0 } -func (n Num) IsNaN() bool { return (n & 0x7fff) > 0x7c00 } +func (n Num) IsNaN() bool { return (n.bits & 0x7fff) > 0x7c00 } -func (n Num) IsInf() bool { return (n & 0x7c00) == 0x7c00 } +func (n Num) IsInf() bool { return (n.bits & 0x7c00) == 0x7c00 } -func (n Num) IsZero() bool { return (n & 0x7fff) == 0 } +func (n Num) IsZero() bool { return (n.bits & 0x7fff) == 0 } -func (f Num) Uint16() uint16 { return uint16(f) } +func (f Num) Uint16() uint16 { return uint16(f.bits) } func (f Num) String() string { return strconv.FormatFloat(float64(f.Float32()), 'g', -1, 32) } -func Inf() Num { return Num(0x7c00) } +func Inf() Num { return Num{bits: 0x7c00} } -func NaN() Num { return Num(0x7fff) } +func NaN() Num { return Num{bits: 0x7fff} } -func FromBits(src uint16) Num { return Num(src) } +func FromBits(src uint16) Num { return Num{bits: src} } func FromLEBytes(src []byte) Num { - return Num(binary.LittleEndian.Uint16(src)) + return Num{bits: binary.LittleEndian.Uint16(src)} } func (f Num) PutLEBytes(dst []byte) { - binary.LittleEndian.PutUint16(dst, uint16(f)) + binary.LittleEndian.PutUint16(dst, f.bits) } func (f Num) ToLEBytes() []byte { diff --git a/arrow/float16/float16_test.go b/arrow/float16/float16_test.go index 271b70708..7f2ef5cca 100644 --- a/arrow/float16/float16_test.go +++ b/arrow/float16/float16_test.go @@ -25,14 +25,14 @@ import ( func TestFloat16(t *testing.T) { cases := map[Num]float32{ - 0x3c00: 1, - 0x4000: 2, - 0xc000: -2, - 0x0000: 0, - 0x5b8f: 241.875, - 0xdb8f: -241.875, - 0x48c8: 9.5625, - 0xc8c8: -9.5625, + {bits: 0x3c00}: 1, + {bits: 0x4000}: 2, + {bits: 0xc000}: -2, + {bits: 0x0000}: 0, + {bits: 0x5b8f}: 241.875, + {bits: 0xdb8f}: -241.875, + {bits: 0x48c8}: 9.5625, + {bits: 0xc8c8}: -9.5625, } for k, v := range cases { f := k.Float32() @@ -50,9 +50,9 @@ func TestAdd(t *testing.T) { rhs Num want Num }{ - {Num(0x0000), Num(0x0000), Num(0x0000)}, // 0 + 0 = 0 - {Num(0x3c00), Num(0x4000), Num(0x4200)}, // 1 + 2 = 3 - {Num(0x4248), Num(0x3245), Num(0x42AC)}, // 3.141 + 0.196 = 3.336 + {Num{bits: 0x0000}, Num{bits: 0x0000}, Num{bits: 0x0000}}, // 0 + 0 = 0 + {Num{bits: 0x3c00}, Num{bits: 0x4000}, Num{bits: 0x4200}}, // 1 + 2 = 3 + {Num{bits: 0x4248}, Num{bits: 0x3245}, Num{bits: 0x42AC}}, // 3.141 + 0.196 = 3.336 } { t.Run("add", func(t *testing.T) { n := tc.n.Add(tc.rhs) @@ -69,9 +69,9 @@ func TestSub(t *testing.T) { rhs Num want Num }{ - {Num(0x0000), Num(0x0000), Num(0x0000)}, // 0 - 0 = 0 - {Num(0x3c00), Num(0x4000), Num(0xBC00)}, // 1 - 2 = -1 - {Num(0x4248), Num(0x3245), Num(0x41E3)}, // 3.141 - 0.196 = 2.944 + {Num{bits: 0x0000}, Num{bits: 0x0000}, Num{bits: 0x0000}}, // 0 - 0 = 0 + {Num{bits: 0x3c00}, Num{bits: 0x4000}, Num{bits: 0xBC00}}, // 1 - 2 = -1 + {Num{bits: 0x4248}, Num{bits: 0x3245}, Num{bits: 0x41E3}}, // 3.141 - 0.196 = 2.944 } { t.Run("sub", func(t *testing.T) { n := tc.n.Sub(tc.rhs) @@ -88,9 +88,9 @@ func TestMul(t *testing.T) { rhs Num want Num }{ - {Num(0x0000), Num(0x0000), Num(0x0000)}, // 0 * 0 = 0 - {Num(0x3c00), Num(0x4000), Num(0x4000)}, // 1 * 2 = 2 - {Num(0x4248), Num(0x3245), Num(0x38EC)}, // 3.141 * 0.196 = 0.6153 + {Num{bits: 0x0000}, Num{bits: 0x0000}, Num{bits: 0x0000}}, // 0 * 0 = 0 + {Num{bits: 0x3c00}, Num{bits: 0x4000}, Num{bits: 0x4000}}, // 1 * 2 = 2 + {Num{bits: 0x4248}, Num{bits: 0x3245}, Num{bits: 0x38EC}}, // 3.141 * 0.196 = 0.6153 } { t.Run("mul", func(t *testing.T) { n := tc.n.Mul(tc.rhs) @@ -107,9 +107,9 @@ func TestDiv(t *testing.T) { rhs Num want Num }{ - {Num(0x0000), Num(0x3c00), Num(0x0000)}, // 0 / 1 = 0 - {Num(0x3c00), Num(0x4000), Num(0x3800)}, // 1 / 2 = 0.5 - {Num(0x4248), Num(0x3245), Num(0x4C01)}, // 3.141 * 0.196 = 16.02 + {Num{bits: 0x0000}, Num{bits: 0x3c00}, Num{bits: 0x0000}}, // 0 / 1 = 0 + {Num{bits: 0x3c00}, Num{bits: 0x4000}, Num{bits: 0x3800}}, // 1 / 2 = 0.5 + {Num{bits: 0x4248}, Num{bits: 0x3245}, Num{bits: 0x4C01}}, // 3.141 * 0.196 = 16.02 } { t.Run("div", func(t *testing.T) { n := tc.n.Div(tc.rhs) @@ -126,9 +126,9 @@ func TestGreater(t *testing.T) { rhs Num want bool }{ - {Num(0x3c00), Num(0x4000), false}, // 1 > 2 = false - {Num(0x4900), Num(0x4900), false}, // 10 == 10 = false - {Num(0x4248), Num(0x3245), true}, // 3.141 > 0.196 = true + {Num{bits: 0x3c00}, Num{bits: 0x4000}, false}, // 1 > 2 = false + {Num{bits: 0x4900}, Num{bits: 0x4900}, false}, // 10 == 10 = false + {Num{bits: 0x4248}, Num{bits: 0x3245}, true}, // 3.141 > 0.196 = true } { t.Run("greater", func(t *testing.T) { n := tc.n.Greater(tc.rhs) @@ -145,9 +145,9 @@ func TestLess(t *testing.T) { rhs Num want bool }{ - {Num(0x3c00), Num(0x4000), true}, // 1 < 2 = true - {Num(0x4900), Num(0x4900), false}, // 10 == 10 = false - {Num(0x4248), Num(0x3245), false}, // 3.141 < 0.196 = false + {Num{bits: 0x3c00}, Num{bits: 0x4000}, true}, // 1 < 2 = true + {Num{bits: 0x4900}, Num{bits: 0x4900}, false}, // 10 == 10 = false + {Num{bits: 0x4248}, Num{bits: 0x3245}, false}, // 3.141 < 0.196 = false } { t.Run("less", func(t *testing.T) { n := tc.n.Less(tc.rhs) @@ -164,9 +164,9 @@ func TestCmp(t *testing.T) { rhs Num want int }{ - {Num(0x3c00), Num(0x4000), -1}, // cmp(1, 2) = -1 - {Num(0x4900), Num(0x4900), 0}, // cmp(10, 10) = 0 - {Num(0x4248), Num(0x3245), 1}, // cmp(3.141, 0.196) = 1 + {Num{bits: 0x3c00}, Num{bits: 0x4000}, -1}, // cmp(1, 2) = -1 + {Num{bits: 0x4900}, Num{bits: 0x4900}, 0}, // cmp(10, 10) = 0 + {Num{bits: 0x4248}, Num{bits: 0x3245}, 1}, // cmp(3.141, 0.196) = 1 } { t.Run("cmp", func(t *testing.T) { n := tc.n.Cmp(tc.rhs) @@ -183,8 +183,8 @@ func TestMax(t *testing.T) { rhs []Num want Num }{ - {Num(0x3c00), []Num{Num(0x4000), Num(0x4580), Num(0x3C00), Num(0x4247)}, Num(0x4580)}, // max(2, 5.5, 1, 3.14) = 5.5 - {Num(0x4248), []Num{Num(0xC000), Num(0xC580), Num(0x3C00), Num(0x4247)}, Num(0x4248)}, // max(-2, -5.5, 1, 3.14) = 3.14 + {Num{bits: 0x3c00}, []Num{Num{bits: 0x4000}, Num{bits: 0x4580}, Num{bits: 0x3C00}, Num{bits: 0x4247}}, Num{bits: 0x4580}}, // max(2, 5.5, 1, 3.14) = 5.5 + {Num{bits: 0x4248}, []Num{Num{bits: 0xC000}, Num{bits: 0xC580}, Num{bits: 0x3C00}, Num{bits: 0x4247}}, Num{bits: 0x4248}}, // max(-2, -5.5, 1, 3.14) = 3.14 } { t.Run("max", func(t *testing.T) { n := Max(tc.n, tc.rhs...) @@ -201,8 +201,8 @@ func TestMin(t *testing.T) { rhs []Num want Num }{ - {Num(0x3c00), []Num{Num(0x4000), Num(0x4580), Num(0x3C00), Num(0x4247)}, Num(0x3C00)}, // min(2, 5.5, 1, 3.14) = 1 - {Num(0x4248), []Num{Num(0x4000), Num(0xC580), Num(0xBC00), Num(0x4247)}, Num(0xC580)}, // min(2, -5.5, -1, 3.14) = -5.5 + {Num{bits: 0x3c00}, []Num{Num{bits: 0x4000}, Num{bits: 0x4580}, Num{bits: 0x3C00}, Num{bits: 0x4247}}, Num{bits: 0x3C00}}, // min(2, 5.5, 1, 3.14) = 1 + {Num{bits: 0x4248}, []Num{Num{bits: 0x4000}, Num{bits: 0xC580}, Num{bits: 0xBC00}, Num{bits: 0x4247}}, Num{bits: 0xC580}}, // min(2, -5.5, -1, 3.14) = -5.5 } { t.Run("min", func(t *testing.T) { n := Min(tc.n, tc.rhs...) @@ -218,9 +218,9 @@ func TestAbs(t *testing.T) { n Num want Num }{ - {Num(0x4580), Num(0x4580)}, // 5.5 - {Num(0x0000), Num(0x0000)}, // 0 - {Num(0xC580), Num(0x4580)}, // -5.5 + {Num{bits: 0x4580}, Num{bits: 0x4580}}, // 5.5 + {Num{bits: 0x0000}, Num{bits: 0x0000}}, // 0 + {Num{bits: 0xC580}, Num{bits: 0x4580}}, // -5.5 } { t.Run("abs", func(t *testing.T) { n := tc.n.Abs() @@ -236,10 +236,10 @@ func TestSign(t *testing.T) { n Num want int }{ - {Num(0x4580), 1}, // 5.5 - {Num(0x0000), 0}, // 0 - {Num(0x8000), 0}, // -0 - {Num(0xC580), -1}, // -5.5 + {Num{bits: 0x4580}, 1}, // 5.5 + {Num{bits: 0x0000}, 0}, // 0 + {Num{bits: 0x8000}, 0}, // -0 + {Num{bits: 0xC580}, -1}, // -5.5 } { t.Run("sign", func(t *testing.T) { n := tc.n.Sign() @@ -255,10 +255,10 @@ func TestSignbit(t *testing.T) { n Num want bool }{ - {Num(0x4580), false}, // 5.5 - {Num(0x0000), false}, // 0 - {Num(0x8000), true}, // -0 - {Num(0xC580), true}, // -5.5 + {Num{bits: 0x4580}, false}, // 5.5 + {Num{bits: 0x0000}, false}, // 0 + {Num{bits: 0x8000}, true}, // -0 + {Num{bits: 0xC580}, true}, // -5.5 } { t.Run("signbit", func(t *testing.T) { n := tc.n.Signbit() @@ -278,10 +278,10 @@ func TestIsNaN(t *testing.T) { {NaN().Negate(), true}, {Inf(), false}, {Inf().Negate(), false}, - {Num(0x7c01), true}, // nan - {Num(0xfc01), true}, // -nan - {Num(0x7e00), true}, // nan - {Num(0xfe00), true}, // -nan + {Num{bits: 0x7c01}, true}, // nan + {Num{bits: 0xfc01}, true}, // -nan + {Num{bits: 0x7e00}, true}, // nan + {Num{bits: 0xfe00}, true}, // -nan } { t.Run("isnan", func(t *testing.T) { n := tc.n.IsNaN() diff --git a/arrow/internal/testing/tools/data_types.go b/arrow/internal/testing/tools/data_types.go index f494c1ffd..5ffd44b48 100644 --- a/arrow/internal/testing/tools/data_types.go +++ b/arrow/internal/testing/tools/data_types.go @@ -42,7 +42,7 @@ var typMap = map[reflect.Type]arrow.DataType{ reflect.TypeOf(arrow.Date32(0)): arrow.FixedWidthTypes.Date32, reflect.TypeOf(arrow.Date64(0)): arrow.FixedWidthTypes.Date64, reflect.TypeOf(true): arrow.FixedWidthTypes.Boolean, - reflect.TypeOf(float16.Num(0)): arrow.FixedWidthTypes.Float16, + reflect.TypeOf(float16.Num{}): arrow.FixedWidthTypes.Float16, reflect.TypeOf([]byte{}): arrow.BinaryTypes.Binary, } diff --git a/arrow/type_traits.go b/arrow/type_traits.go index 1eb30eb2e..7185ef25a 100644 --- a/arrow/type_traits.go +++ b/arrow/type_traits.go @@ -124,23 +124,23 @@ func GetData[T FixedWidthType | ViewHeader](in []byte) []T { } var typMap = map[reflect.Type]DataType{ - reflect.TypeOf(false): FixedWidthTypes.Boolean, - reflect.TypeOf(int8(0)): PrimitiveTypes.Int8, - reflect.TypeOf(int16(0)): PrimitiveTypes.Int16, - reflect.TypeOf(int32(0)): PrimitiveTypes.Int32, - reflect.TypeOf(int64(0)): PrimitiveTypes.Int64, - reflect.TypeOf(uint8(0)): PrimitiveTypes.Uint8, - reflect.TypeOf(uint16(0)): PrimitiveTypes.Uint16, - reflect.TypeOf(uint32(0)): PrimitiveTypes.Uint32, - reflect.TypeOf(uint64(0)): PrimitiveTypes.Uint64, - reflect.TypeOf(float32(0)): PrimitiveTypes.Float32, - reflect.TypeOf(float64(0)): PrimitiveTypes.Float64, - reflect.TypeOf(string("")): BinaryTypes.String, - reflect.TypeOf(Date32(0)): FixedWidthTypes.Date32, - reflect.TypeOf(Date64(0)): FixedWidthTypes.Date64, - reflect.TypeOf(true): FixedWidthTypes.Boolean, - reflect.TypeOf(float16.Num(0)): FixedWidthTypes.Float16, - reflect.TypeOf([]byte{}): BinaryTypes.Binary, + reflect.TypeOf(false): FixedWidthTypes.Boolean, + reflect.TypeOf(int8(0)): PrimitiveTypes.Int8, + reflect.TypeOf(int16(0)): PrimitiveTypes.Int16, + reflect.TypeOf(int32(0)): PrimitiveTypes.Int32, + reflect.TypeOf(int64(0)): PrimitiveTypes.Int64, + reflect.TypeOf(uint8(0)): PrimitiveTypes.Uint8, + reflect.TypeOf(uint16(0)): PrimitiveTypes.Uint16, + reflect.TypeOf(uint32(0)): PrimitiveTypes.Uint32, + reflect.TypeOf(uint64(0)): PrimitiveTypes.Uint64, + reflect.TypeOf(float32(0)): PrimitiveTypes.Float32, + reflect.TypeOf(float64(0)): PrimitiveTypes.Float64, + reflect.TypeOf(string("")): BinaryTypes.String, + reflect.TypeOf(Date32(0)): FixedWidthTypes.Date32, + reflect.TypeOf(Date64(0)): FixedWidthTypes.Date64, + reflect.TypeOf(true): FixedWidthTypes.Boolean, + reflect.TypeOf(float16.Num{}): FixedWidthTypes.Float16, + reflect.TypeOf([]byte{}): BinaryTypes.Binary, } // GetDataType returns the appropriate DataType for the given type T diff --git a/parquet/schema/reflection.go b/parquet/schema/reflection.go index 9c6738658..6082c2988 100644 --- a/parquet/schema/reflection.go +++ b/parquet/schema/reflection.go @@ -376,6 +376,9 @@ func typeToNode(name string, typ reflect.Type, repType parquet.Repetition, info } return Must(MapOf(name, key, value, repType, fieldID)) case reflect.Struct: + if typ == reflect.TypeOf(float16.Num{}) { + return MustPrimitive(NewPrimitiveNodeLogical(name, repType, Float16LogicalType{}, parquet.Types.FixedLenByteArray, 2, fieldID)) + } // structs are Group nodes fields := make(FieldList, 0) for i := 0; i < typ.NumField(); i++ { @@ -499,10 +502,6 @@ func typeToNode(name string, typ reflect.Type, repType parquet.Repetition, info return MustPrimitive(NewPrimitiveNodeLogical(name, repType, NewIntLogicalType(bitwidth, true), ptyp, 0, fieldID)) case reflect.Uint, reflect.Uint32, reflect.Uint8, reflect.Uint16, reflect.Uint64: - if typ == reflect.TypeOf(float16.Num(0)) { - return MustPrimitive(NewPrimitiveNodeLogical(name, repType, Float16LogicalType{}, parquet.Types.FixedLenByteArray, 2, fieldID)) - } - // handle unsigned integer types and default to the corresponding logical type for it. ptyp := parquet.Types.Int32 if typ.Bits() == 64 { From a22cd326db820c27b23ae599be4b879a62bc1688 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Thu, 3 Jul 2025 12:49:45 -0400 Subject: [PATCH 4/4] lint --- arrow/float16/float16_test.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/arrow/float16/float16_test.go b/arrow/float16/float16_test.go index 7f2ef5cca..9857eda5c 100644 --- a/arrow/float16/float16_test.go +++ b/arrow/float16/float16_test.go @@ -183,8 +183,8 @@ func TestMax(t *testing.T) { rhs []Num want Num }{ - {Num{bits: 0x3c00}, []Num{Num{bits: 0x4000}, Num{bits: 0x4580}, Num{bits: 0x3C00}, Num{bits: 0x4247}}, Num{bits: 0x4580}}, // max(2, 5.5, 1, 3.14) = 5.5 - {Num{bits: 0x4248}, []Num{Num{bits: 0xC000}, Num{bits: 0xC580}, Num{bits: 0x3C00}, Num{bits: 0x4247}}, Num{bits: 0x4248}}, // max(-2, -5.5, 1, 3.14) = 3.14 + {Num{bits: 0x3c00}, []Num{{bits: 0x4000}, {bits: 0x4580}, {bits: 0x3C00}, {bits: 0x4247}}, Num{bits: 0x4580}}, // max(2, 5.5, 1, 3.14) = 5.5 + {Num{bits: 0x4248}, []Num{{bits: 0xC000}, {bits: 0xC580}, {bits: 0x3C00}, {bits: 0x4247}}, Num{bits: 0x4248}}, // max(-2, -5.5, 1, 3.14) = 3.14 } { t.Run("max", func(t *testing.T) { n := Max(tc.n, tc.rhs...) @@ -201,8 +201,8 @@ func TestMin(t *testing.T) { rhs []Num want Num }{ - {Num{bits: 0x3c00}, []Num{Num{bits: 0x4000}, Num{bits: 0x4580}, Num{bits: 0x3C00}, Num{bits: 0x4247}}, Num{bits: 0x3C00}}, // min(2, 5.5, 1, 3.14) = 1 - {Num{bits: 0x4248}, []Num{Num{bits: 0x4000}, Num{bits: 0xC580}, Num{bits: 0xBC00}, Num{bits: 0x4247}}, Num{bits: 0xC580}}, // min(2, -5.5, -1, 3.14) = -5.5 + {Num{bits: 0x3c00}, []Num{{bits: 0x4000}, {bits: 0x4580}, {bits: 0x3C00}, {bits: 0x4247}}, Num{bits: 0x3C00}}, // min(2, 5.5, 1, 3.14) = 1 + {Num{bits: 0x4248}, []Num{{bits: 0x4000}, {bits: 0xC580}, {bits: 0xBC00}, {bits: 0x4247}}, Num{bits: 0xC580}}, // min(2, -5.5, -1, 3.14) = -5.5 } { t.Run("min", func(t *testing.T) { n := Min(tc.n, tc.rhs...)