diff --git a/cpp/src/arrow/compute/kernels/vector_hash.cc b/cpp/src/arrow/compute/kernels/vector_hash.cc index 90ec9e365c35..f42972e18af0 100644 --- a/cpp/src/arrow/compute/kernels/vector_hash.cc +++ b/cpp/src/arrow/compute/kernels/vector_hash.cc @@ -550,6 +550,7 @@ KernelInit GetHashInit(Type::type type_id) { return HashInit>; case Type::INT16: case Type::UINT16: + case Type::HALF_FLOAT: return HashInit>; case Type::INT32: case Type::UINT32: @@ -700,6 +701,13 @@ void AddHashKernels(VectorFunction* func, VectorKernel base, OutputType out_ty) DCHECK_OK(func->AddKernel(base)); } + // float16() is not part of PrimitiveTypes() (FloatingPointTypes() only covers + // float32 and float64; see GH-43017), so it must be registered explicitly. Like + // float32 and float64, it is hashed by its raw bit pattern (via UInt16Type). + base.init = GetHashInit(Type::HALF_FLOAT); + base.signature = KernelSignature::Make({float16()}, out_ty); + DCHECK_OK(func->AddKernel(base)); + // Parametric types that we want matching to be dependent only on type id auto parametric_types = {Type::TIME32, Type::TIME64, Type::TIMESTAMP, Type::DURATION, Type::FIXED_SIZE_BINARY}; diff --git a/cpp/src/arrow/compute/kernels/vector_hash_test.cc b/cpp/src/arrow/compute/kernels/vector_hash_test.cc index b0fa296e0074..5b066ab782f4 100644 --- a/cpp/src/arrow/compute/kernels/vector_hash_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_hash_test.cc @@ -403,6 +403,44 @@ TEST_F(TestHashKernel, DictEncodeBoolean) { ArrayFromJSON(boolean(), "[true]"), ArrayFromJSON(int32(), "[0, null, 0]")); } +// float16 is not part of PrimitiveTypes(), so it is not covered by +// TestHashKernelPrimitive above and gets its own coverage here. Like float32 and +// float64, it is hashed by its raw bit pattern. +TEST_F(TestHashKernel, UniqueHalfFloat) { + CheckUnique(ArrayFromJSON(float16(), "[1.5, 2.5, null, 1.5, 3.5, null]"), + ArrayFromJSON(float16(), "[1.5, 2.5, null, 3.5]")); + + // No nulls + CheckUnique(ArrayFromJSON(float16(), "[1.5, 2.5, 1.5, 3.5]"), + ArrayFromJSON(float16(), "[1.5, 2.5, 3.5]")); + + // Sliced + CheckUnique(ArrayFromJSON(float16(), "[1.5, 2.5, null, 3.5, 2.5, 1.5]")->Slice(1, 4), + ArrayFromJSON(float16(), "[2.5, null, 3.5]")); +} + +TEST_F(TestHashKernel, ValueCountsHalfFloat) { + CheckValueCounts(ArrayFromJSON(float16(), "[1.5, 2.5, null, 1.5, 3.5, null]"), + ArrayFromJSON(float16(), "[1.5, 2.5, null, 3.5]"), + ArrayFromJSON(int64(), "[2, 1, 2, 1]")); +} + +TEST_F(TestHashKernel, DictEncodeHalfFloat) { + CheckDictEncode(ArrayFromJSON(float16(), "[1.5, 2.5, null, 1.5, 3.5, null]"), + ArrayFromJSON(float16(), "[1.5, 2.5, 3.5]"), + ArrayFromJSON(int32(), "[0, 1, null, 0, 2, null]")); + + // No nulls + CheckDictEncode(ArrayFromJSON(float16(), "[1.5, 2.5, 1.5, 3.5]"), + ArrayFromJSON(float16(), "[1.5, 2.5, 3.5]"), + ArrayFromJSON(int32(), "[0, 1, 0, 2]")); + + // Sliced + CheckDictEncode( + ArrayFromJSON(float16(), "[1.5, 2.5, null, 3.5, 2.5, 1.5]")->Slice(1, 4), + ArrayFromJSON(float16(), "[2.5, 3.5]"), ArrayFromJSON(int32(), "[0, null, 1, 0]")); +} + template class TestHashKernelBinaryTypes : public TestHashKernel { protected: