Skip to content

Commit a043018

Browse files
Licht-Twesm
Licht-T
authored and
wesm
committed
ARROW-1678: [Python] Implement numpy.float16 SerDe
This is the patch for [ARROW-1678](https://issues.apache.org/jira/projects/ARROW/issues/ARROW-1678). Author: Licht-T <licht-t@outlook.jp> Closes apache#1205 from Licht-T/feature-halffloat-python-serde and squashes the following commits: e397cc9 [Licht-T] Add NPY_HALF case into SerializeArray method 22c37d7 [Licht-T] Use npy_half as type of the numpy float16 method return value 319abb7 [Licht-T] Add RETURN_IF_PYERROR after PyArrayScalar_New eec4e0a [Licht-T] Change headers include order 5a985b8 [Licht-T] Fix lint issue f19e543 [Licht-T] TST: Add tests for numpy.float16 SerDe 66d5f15 [Licht-T] ENH: Set VISIT_INLINE for HalfFloat 834357e [Licht-T] ENH: Add numpy.float16 deserializer 02dc77b [Licht-T] ENH: Add numpy.float16 serializer 9d19e08 [Licht-T] ENH: Add HalfFloat type case for numpy dda98b9 [Licht-T] ENH: Add HalfFloat type traits 0be305a [Licht-T] ENH: Add the IPC internal metadata for HalfFloat 7348cc6 [Licht-T] ENH: Add HalfFloat BUILDER_CASE
1 parent 8eb2b0e commit a043018

9 files changed

Lines changed: 63 additions & 3 deletions

File tree

‎cpp/src/arrow/builder.cc‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1438,6 +1438,7 @@ Status MakeBuilder(MemoryPool* pool, const std::shared_ptr<DataType>& type,
14381438
BUILDER_CASE(TIME64, Time64Builder);
14391439
BUILDER_CASE(TIMESTAMP, TimestampBuilder);
14401440
BUILDER_CASE(BOOL, BooleanBuilder);
1441+
BUILDER_CASE(HALF_FLOAT, HalfFloatBuilder);
14411442
BUILDER_CASE(FLOAT, FloatBuilder);
14421443
BUILDER_CASE(DOUBLE, DoubleBuilder);
14431444
BUILDER_CASE(STRING, StringBuilder);

‎cpp/src/arrow/ipc/metadata-internal.cc‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -381,6 +381,10 @@ static Status TypeToFlatbuffer(FBB& fbb, const DataType& type,
381381
INT_TO_FB_CASE(64, false);
382382
case Type::INT64:
383383
INT_TO_FB_CASE(64, true);
384+
case Type::HALF_FLOAT:
385+
*out_type = flatbuf::Type_FloatingPoint;
386+
*offset = FloatToFlatbuffer(fbb, flatbuf::Precision_HALF);
387+
break;
384388
case Type::FLOAT:
385389
*out_type = flatbuf::Type_FloatingPoint;
386390
*offset = FloatToFlatbuffer(fbb, flatbuf::Precision_SINGLE);

‎cpp/src/arrow/python/arrow_to_python.cc‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,15 @@
1717

1818
#include "arrow/python/arrow_to_python.h"
1919

20+
#include "arrow/python/numpy_interop.h"
21+
2022
#include <cstdint>
2123
#include <memory>
2224
#include <vector>
2325

26+
#include <numpy/arrayobject.h>
27+
#include <numpy/arrayscalars.h>
28+
2429
#include "arrow/array.h"
2530
#include "arrow/io/interfaces.h"
2631
#include "arrow/ipc/reader.h"
@@ -121,6 +126,14 @@ Status GetValue(PyObject* context, const Array& arr, int64_t index, int32_t type
121126
*result = PyUnicode_FromStringAndSize(reinterpret_cast<const char*>(str), nchars);
122127
return CheckPyError();
123128
}
129+
case Type::HALF_FLOAT: {
130+
*result = PyArrayScalar_New(Half);
131+
RETURN_IF_PYERROR();
132+
133+
npy_half halffloat = static_cast<const HalfFloatArray&>(arr).Value(index);
134+
PyArrayScalar_ASSIGN(*result, Half, halffloat);
135+
return Status::OK();
136+
}
124137
case Type::FLOAT:
125138
*result = PyFloat_FromDouble(static_cast<const FloatArray&>(arr).Value(index));
126139
return Status::OK();

‎cpp/src/arrow/python/numpy_convert.cc‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,7 @@ Status NumPyDtypeToArrow(PyObject* dtype, std::shared_ptr<DataType>* out) {
165165
#if (NPY_UINT64 != NPY_ULONGLONG)
166166
TO_ARROW_CASE(ULONGLONG);
167167
#endif
168+
TO_ARROW_TYPE_CASE(FLOAT16, float16);
168169
TO_ARROW_TYPE_CASE(FLOAT32, float32);
169170
TO_ARROW_TYPE_CASE(FLOAT64, float64);
170171
TO_ARROW_TYPE_CASE(STRING, binary);

‎cpp/src/arrow/python/numpy_to_arrow.cc‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,6 +284,8 @@ class NumPyConverter {
284284
return VisitNative<T>();
285285
}
286286

287+
Status Visit(const HalfFloatType& type) { return VisitNative<UInt16Type>(); }
288+
287289
Status Visit(const Date32Type& type) { return VisitNative<Date32Type>(); }
288290
Status Visit(const Date64Type& type) { return VisitNative<Int64Type>(); }
289291
Status Visit(const TimestampType& type) { return VisitNative<TimestampType>(); }
@@ -1183,6 +1185,7 @@ Status NumPyConverter::ConvertLists(const std::shared_ptr<DataType>& type,
11831185
LIST_CASE(UINT64, NPY_UINT64, UInt64Type)
11841186
LIST_CASE(INT64, NPY_INT64, Int64Type)
11851187
LIST_CASE(TIMESTAMP, NPY_DATETIME, TimestampType)
1188+
LIST_CASE(HALF_FLOAT, NPY_FLOAT16, HalfFloatType)
11861189
LIST_CASE(FLOAT, NPY_FLOAT, FloatType)
11871190
LIST_CASE(DOUBLE, NPY_DOUBLE, DoubleType)
11881191
LIST_CASE(STRING, NPY_OBJECT, StringType)

‎cpp/src/arrow/python/python_to_arrow.cc‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ class SequenceBuilder {
5858
ints_(::arrow::int64(), pool),
5959
bytes_(::arrow::binary(), pool),
6060
strings_(pool),
61+
half_floats_(::arrow::float16(), pool),
6162
floats_(::arrow::float32(), pool),
6263
doubles_(::arrow::float64(), pool),
6364
date64s_(::arrow::date64(), pool),
@@ -117,6 +118,11 @@ class SequenceBuilder {
117118
return strings_.Append(data, length);
118119
}
119120

121+
/// Appending a half_float to the sequence
122+
Status AppendHalfFloat(const npy_half data) {
123+
return AppendPrimitive(data, &half_float_tag_, &half_floats_);
124+
}
125+
120126
/// Appending a float to the sequence
121127
Status AppendFloat(const float data) {
122128
return AppendPrimitive(data, &float_tag_, &floats_);
@@ -222,6 +228,7 @@ class SequenceBuilder {
222228
RETURN_NOT_OK(AddElement(int_tag_, &ints_));
223229
RETURN_NOT_OK(AddElement(string_tag_, &strings_));
224230
RETURN_NOT_OK(AddElement(bytes_tag_, &bytes_));
231+
RETURN_NOT_OK(AddElement(half_float_tag_, &half_floats_));
225232
RETURN_NOT_OK(AddElement(float_tag_, &floats_));
226233
RETURN_NOT_OK(AddElement(double_tag_, &doubles_));
227234
RETURN_NOT_OK(AddElement(date64_tag_, &date64s_));
@@ -250,6 +257,7 @@ class SequenceBuilder {
250257
Int64Builder ints_;
251258
BinaryBuilder bytes_;
252259
StringBuilder strings_;
260+
HalfFloatBuilder half_floats_;
253261
FloatBuilder floats_;
254262
DoubleBuilder doubles_;
255263
Date64Builder date64s_;
@@ -274,6 +282,7 @@ class SequenceBuilder {
274282
int8_t int_tag_ = -1;
275283
int8_t string_tag_ = -1;
276284
int8_t bytes_tag_ = -1;
285+
int8_t half_float_tag_ = -1;
277286
int8_t float_tag_ = -1;
278287
int8_t double_tag_ = -1;
279288
int8_t date64_tag_ = -1;
@@ -394,6 +403,8 @@ Status SerializeSequences(PyObject* context, std::vector<PyObject*> sequences,
394403
Status AppendScalar(PyObject* obj, SequenceBuilder* builder) {
395404
if (PyArray_IsScalar(obj, Bool)) {
396405
return builder->AppendBool(reinterpret_cast<PyBoolScalarObject*>(obj)->obval != 0);
406+
} else if (PyArray_IsScalar(obj, Half)) {
407+
return builder->AppendHalfFloat(reinterpret_cast<PyHalfScalarObject*>(obj)->obval);
397408
} else if (PyArray_IsScalar(obj, Float)) {
398409
return builder->AppendFloat(reinterpret_cast<PyFloatScalarObject*>(obj)->obval);
399410
} else if (PyArray_IsScalar(obj, Double)) {
@@ -437,6 +448,9 @@ Status Append(PyObject* context, PyObject* elem, SequenceBuilder* builder,
437448
// The bool case must precede the int case (PyInt_Check passes for bools)
438449
if (PyBool_Check(elem)) {
439450
RETURN_NOT_OK(builder->AppendBool(elem == Py_True));
451+
} else if (PyArray_DescrFromScalar(elem)->type_num == NPY_HALF) {
452+
npy_half halffloat = reinterpret_cast<PyHalfScalarObject *>(elem)->obval;
453+
RETURN_NOT_OK(builder->AppendHalfFloat(halffloat));
440454
} else if (PyFloat_Check(elem)) {
441455
RETURN_NOT_OK(builder->AppendDouble(PyFloat_AS_DOUBLE(elem)));
442456
} else if (PyLong_Check(elem)) {
@@ -523,6 +537,7 @@ Status SerializeArray(PyObject* context, PyArrayObject* array, SequenceBuilder*
523537
case NPY_INT32:
524538
case NPY_UINT64:
525539
case NPY_INT64:
540+
case NPY_HALF:
526541
case NPY_FLOAT:
527542
case NPY_DOUBLE: {
528543
RETURN_NOT_OK(builder->AppendTensor(static_cast<int32_t>(tensors_out->size())));

‎cpp/src/arrow/python/type_traits.h‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424

2525
#include "arrow/python/numpy_interop.h"
2626

27+
#include <numpy/halffloat.h>
28+
2729
#include "arrow/builder.h"
2830
#include "arrow/type.h"
2931
#include "arrow/util/logging.h"
@@ -71,6 +73,17 @@ NPY_INT_DECL(LONGLONG, Int64, int64_t);
7173
NPY_INT_DECL(ULONGLONG, UInt64, uint64_t);
7274
#endif
7375

76+
template <>
77+
struct npy_traits<NPY_FLOAT16> {
78+
typedef npy_half value_type;
79+
using TypeClass = HalfFloatType;
80+
using BuilderClass = HalfFloatBuilder;
81+
82+
static constexpr bool supports_nulls = true;
83+
84+
static inline bool isnull(npy_half v) { return v == NPY_HALF_NAN; }
85+
};
86+
7487
template <>
7588
struct npy_traits<NPY_FLOAT32> {
7689
typedef float value_type;
@@ -143,6 +156,14 @@ INT_DECL(UINT16);
143156
INT_DECL(UINT32);
144157
INT_DECL(UINT64);
145158

159+
template <>
160+
struct arrow_traits<Type::HALF_FLOAT> {
161+
static constexpr int npy_type = NPY_FLOAT16;
162+
static constexpr bool supports_nulls = true;
163+
static constexpr uint16_t na_value = NPY_HALF_NAN;
164+
typedef typename npy_traits<NPY_FLOAT16>::value_type T;
165+
};
166+
146167
template <>
147168
struct arrow_traits<Type::FLOAT> {
148169
static constexpr int npy_type = NPY_FLOAT32;

‎cpp/src/arrow/visitor_inline.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ inline Status VisitTypeInline(const DataType& type, VISITOR* visitor) {
4444
TYPE_VISIT_INLINE(UInt32Type);
4545
TYPE_VISIT_INLINE(Int64Type);
4646
TYPE_VISIT_INLINE(UInt64Type);
47+
TYPE_VISIT_INLINE(HalfFloatType);
4748
TYPE_VISIT_INLINE(FloatType);
4849
TYPE_VISIT_INLINE(DoubleType);
4950
TYPE_VISIT_INLINE(StringType);
@@ -85,6 +86,7 @@ inline Status VisitArrayInline(const Array& array, VISITOR* visitor) {
8586
ARRAY_VISIT_INLINE(UInt32Type);
8687
ARRAY_VISIT_INLINE(Int64Type);
8788
ARRAY_VISIT_INLINE(UInt64Type);
89+
ARRAY_VISIT_INLINE(HalfFloatType);
8890
ARRAY_VISIT_INLINE(FloatType);
8991
ARRAY_VISIT_INLINE(DoubleType);
9092
ARRAY_VISIT_INLINE(StringType);

‎python/pyarrow/tests/test_serialization.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ def assert_equal(obj1, obj2):
9999
{True: "hello", False: "world"}, {"hello": "world", 1: 42, 2.5: 45},
100100
{"hello": set([2, 3]), "world": set([42.0]), "this": None},
101101
np.int8(3), np.int32(4), np.int64(5),
102-
np.uint8(3), np.uint32(4), np.uint64(5), np.float32(1.9),
102+
np.uint8(3), np.uint32(4), np.uint64(5), np.float16(1.9), np.float32(1.9),
103103
np.float64(1.9), np.zeros([100, 100]),
104104
np.random.normal(size=[100, 100]), np.array(["hi", 3]),
105105
np.array(["hi", 3], dtype=object),
@@ -257,8 +257,8 @@ def test_default_dict_serialization(large_memory_map):
257257

258258
def test_numpy_serialization(large_memory_map):
259259
with pa.memory_map(large_memory_map, mode="r+") as mmap:
260-
for t in ["int8", "uint8", "int16", "uint16",
261-
"int32", "uint32", "float32", "float64"]:
260+
for t in ["int8", "uint8", "int16", "uint16", "int32", "uint32",
261+
"float16", "float32", "float64"]:
262262
obj = np.random.randint(0, 10, size=(100, 100)).astype(t)
263263
serialization_roundtrip(obj, mmap)
264264

0 commit comments

Comments
 (0)