Update Python bindings, refine error types

Author: wesm

cpp/src/arrow/compute/api.h

@@ -27,3 +27,4 @@
 #include "arrow/compute/function.h"       // IWYU pragma: export
 #include "arrow/compute/kernel.h"         // IWYU pragma: export
 #include "arrow/compute/registry.h"       // IWYU pragma: export
+#include "arrow/datum.h"                  // IWYU pragma: export

cpp/src/arrow/compute/cast.cc

@@ -129,8 +129,9 @@ Result<const ScalarKernel*> CastFunction::DispatchExact(
   }
 
   if (candidate_kernels.size() == 0) {
-    return Status::Invalid("Function ", this->name(),
-                           " has no kernel matching input type ", values[0].ToString());
+    return Status::NotImplemented("Function ", this->name(),
+                                  " has no kernel matching input type ",
+                                  values[0].ToString());
   } else if (candidate_kernels.size() == 1) {
     // One match, return it
     return candidate_kernels[0];
@@ -173,7 +174,8 @@ Result<std::shared_ptr<const CastFunction>> GetCastFunction(
   internal::EnsureInitCastTable();
   auto it = internal::g_cast_table.find(static_cast<int>(to_type->id()));
   if (it == internal::g_cast_table.end()) {
-    return Status::Invalid("No cast function available to cast to ", to_type->ToString());
+    return Status::NotImplemented("No cast function available to cast to ",
+                                  to_type->ToString());
   }
   return it->second;
 }

cpp/src/arrow/compute/kernels/scalar_cast_internal.cc

@@ -63,8 +63,7 @@ void ZeroCopyCastExec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
   }
 }
 
-void AddZeroCopyCast(InputType in_type, const std::shared_ptr<DataType>& out_type,
-                     CastFunction* func) {
+void AddZeroCopyCast(InputType in_type, OutputType out_type, CastFunction* func) {
   auto sig = KernelSignature::Make({in_type}, out_type);
   ScalarKernel kernel;
   kernel.exec = ZeroCopyCastExec;

cpp/src/arrow/compute/kernels/scalar_cast_internal.h

@@ -172,9 +172,12 @@ struct FromDictionaryCast {
     const Array& dictionary = *input.dictionary;
     const DataType& values_type = *dictionary.type();
 
-    // Check if values and output type match
-    DCHECK(values_type.Equals(*output->type))
-        << "Dictionary type: " << values_type << " target type: " << (*output->type);
+    // ARROW-7077
+    if (!values_type.Equals(*output->type)) {
+      ctx->SetStatus(Status::Invalid("Cannot unpack dictionary of type ", type.ToString(),
+                                     " to type ", output->type->ToString()));
+      return;
+    }
 
     FromDictUnpackHelper<T> unpack_helper;
     switch (type.index_type()->id()) {
@@ -237,8 +240,7 @@ void AddSimpleCast(InputType in_ty, OutputType out_ty, CastFunction* func) {
 
 void ZeroCopyCastExec(KernelContext* ctx, const ExecBatch& batch, Datum* out);
 
-void AddZeroCopyCast(InputType in_type, const std::shared_ptr<DataType>& out_type,
-                     CastFunction* func);
+void AddZeroCopyCast(InputType in_type, OutputType out_type, CastFunction* func);
 
 // OutputType::Resolver that returns a descr with the shape of the input
 // argument and the type from CastOptions
@@ -252,7 +254,8 @@ struct MaybeAddFromDictionary {
 
 template <typename T>
 struct MaybeAddFromDictionary<
-    T, enable_if_t<!is_boolean_type<T>::value && !is_nested_type<T>::value>> {
+    T, enable_if_t<!is_boolean_type<T>::value && !is_nested_type<T>::value &&
+                   !std::is_same<DictionaryType, T>::value>> {
   static void Add(const OutputType& out_ty, CastFunction* func) {
     // Dictionary unpacking not implemented for boolean or nested types
     DCHECK_OK(func->AddKernel(Type::DICTIONARY, {InputType::Array(Type::DICTIONARY)},

cpp/src/arrow/compute/kernels/scalar_cast_nested.cc

@@ -76,7 +76,16 @@ std::vector<std::shared_ptr<CastFunction>> GetNestedCasts() {
   AddCommonCasts<LargeListType>(kOutputTargetType, cast_large_list.get());
   AddListCast<LargeListType>(cast_large_list.get());
 
-  return {cast_list, cast_large_list};
+  // FSL is a bit incomplete at the moment
+  auto cast_fsl =
+      std::make_shared<CastFunction>("cast_fixed_size_list", Type::FIXED_SIZE_LIST);
+  AddCommonCasts<FixedSizeListType>(kOutputTargetType, cast_fsl.get());
+
+  // So is struct
+  auto cast_struct = std::make_shared<CastFunction>("cast_struct", Type::STRUCT);
+  AddCommonCasts<StructType>(kOutputTargetType, cast_struct.get());
+
+  return {cast_list, cast_large_list, cast_fsl, cast_struct};
 }
 
 }  // namespace internal

cpp/src/arrow/compute/kernels/scalar_cast_numeric.cc

@@ -546,6 +546,12 @@ std::vector<std::shared_ptr<CastFunction>> GetNumericCasts() {
   functions.push_back(GetCastToInteger<UInt32Type>("cast_uint32"));
   functions.push_back(GetCastToInteger<UInt64Type>("cast_uint64"));
 
+  // HalfFloat is a bit brain-damaged for now
+  auto cast_half_float =
+      std::make_shared<CastFunction>("cast_half_float", Type::HALF_FLOAT);
+  AddCommonCasts<HalfFloatType>(float16(), cast_half_float.get());
+  functions.push_back(cast_half_float);
+
   functions.push_back(GetCastToFloating<FloatType>("cast_float"));
   functions.push_back(GetCastToFloating<DoubleType>("cast_double"));
 

cpp/src/arrow/compute/kernels/scalar_cast_temporal.cc

@@ -134,7 +134,8 @@ struct CastFunctor<
     const auto& in_type = checked_cast<const I&>(*batch[0].type());
     const auto& out_type = checked_cast<const O&>(*output->type);
 
-    DCHECK_NE(in_type.unit(), out_type.unit()) << "Do not cast equal types";
+    // The units may be equal if the time zones are different. We might go to
+    // lengths to make this zero copy in the future but we leave it for now
 
     auto conversion = util::kTimestampConversionTable[static_cast<int>(in_type.unit())]
                                                      [static_cast<int>(out_type.unit())];
@@ -344,10 +345,7 @@ std::shared_ptr<CastFunction> GetDurationCast() {
   auto nanos = duration(TimeUnit::NANO);
 
   // Same integer representation
-  AddZeroCopyCast(/*in_type=*/int64(), /*out_type=*/seconds, func.get());
-  AddZeroCopyCast(int64(), millis, func.get());
-  AddZeroCopyCast(int64(), micros, func.get());
-  AddZeroCopyCast(int64(), nanos, func.get());
+  AddZeroCopyCast(/*in_type=*/int64(), kOutputTargetType, func.get());
 
   // Between durations
   AddCrossUnitCast<DurationType>(func.get());
@@ -359,12 +357,8 @@ std::shared_ptr<CastFunction> GetTime32Cast() {
   auto func = std::make_shared<CastFunction>("cast_time32", Type::TIME32);
   AddCommonCasts<Date32Type>(kOutputTargetType, func.get());
 
-  auto seconds = time32(TimeUnit::SECOND);
-  auto millis = time32(TimeUnit::MILLI);
-
   // Zero copy when the unit is the same or same integer representation
-  AddZeroCopyCast(int32(), seconds, func.get());
-  AddZeroCopyCast(int32(), millis, func.get());
+  AddZeroCopyCast(/*in_type=*/int32(), kOutputTargetType, func.get());
 
   // time64 -> time32
   AddSimpleCast<Time64Type, Time32Type>(InputType(Type::TIME64), kOutputTargetType,
@@ -380,12 +374,8 @@ std::shared_ptr<CastFunction> GetTime64Cast() {
   auto func = std::make_shared<CastFunction>("cast_time64", Type::TIME64);
   AddCommonCasts<Time64Type>(kOutputTargetType, func.get());
 
-  auto micros = time64(TimeUnit::MICRO);
-  auto nanos = time64(TimeUnit::NANO);
-
   // Zero copy when the unit is the same or same integer representation
-  AddZeroCopyCast(int64(), micros, func.get());
-  AddZeroCopyCast(int64(), nanos, func.get());
+  AddZeroCopyCast(/*in_type=*/int64(), kOutputTargetType, func.get());
 
   // time32 -> time64
   AddSimpleCast<Time32Type, Time64Type>(InputType(Type::TIME32), kOutputTargetType,
@@ -402,11 +392,7 @@ std::shared_ptr<CastFunction> GetTimestampCast() {
   AddCommonCasts<TimestampType>(kOutputTargetType, func.get());
 
   // Same integer representation
-  AddZeroCopyCast(/*in_type=*/int64(), /*out_type=*/timestamp(TimeUnit::SECOND),
-                  func.get());
-  AddZeroCopyCast(int64(), timestamp(TimeUnit::MILLI), func.get());
-  AddZeroCopyCast(int64(), timestamp(TimeUnit::MICRO), func.get());
-  AddZeroCopyCast(int64(), timestamp(TimeUnit::NANO), func.get());
+  AddZeroCopyCast(/*in_type=*/int64(), kOutputTargetType, func.get());
 
   // From date types
   // TODO: ARROW-8876, these casts are not implemented

cpp/src/arrow/compute/kernels/scalar_cast_test.cc

@@ -1216,7 +1216,7 @@ TEST_F(TestCast, UnsupportedTarget) {
   std::shared_ptr<Array> arr;
   ArrayFromVector<Int32Type, int32_t>(int32(), is_valid, v1, &arr);
 
-  ASSERT_RAISES(Invalid, Cast(*arr, list(utf8())));
+  ASSERT_RAISES(NotImplemented, Cast(*arr, list(utf8())));
 }
 
 TEST_F(TestCast, DateTimeZeroCopy) {
@@ -1322,8 +1322,8 @@ TEST_F(TestCast, ListToPrimitive) {
   auto from_int = ArrayFromJSON(list(int8()), "[[1, 2], [3, 4]]");
   auto from_binary = ArrayFromJSON(list(binary()), "[[\"1\", \"2\"], [\"3\", \"4\"]]");
 
-  ASSERT_RAISES(Invalid, Cast(*from_int, uint8()));
-  ASSERT_RAISES(Invalid, Cast(*from_binary, utf8()));
+  ASSERT_RAISES(NotImplemented, Cast(*from_int, uint8()));
+  ASSERT_RAISES(NotImplemented, Cast(*from_binary, utf8()));
 }
 
 TEST_F(TestCast, ListToList) {
@@ -1530,7 +1530,7 @@ TYPED_TEST(TestDictionaryCast, DISABLED_OutTypeError) {
   auto out_type = (plain_array->type()->id() == Type::INT8) ? binary() : int8();
   // Test an output type that's not part of TestTypes.
   out_type = list(in_type);
-  ASSERT_RAISES(Invalid, GetCastFunction(out_type));
+  ASSERT_RAISES(NotImplemented, GetCastFunction(out_type));
 }
 
 std::shared_ptr<Array> SmallintArrayFromJSON(const std::string& json_data) {

python/pyarrow/_compute.pyx

@@ -20,18 +20,18 @@
 from pyarrow.lib cimport (
     Array,
     wrap_datum,
-    _context,
     check_status,
     ChunkedArray
 )
 from pyarrow.includes.libarrow cimport CDatum, Sum
+from pyarrow.includes.common cimport *
 
 
 cdef _sum_array(array: Array):
     cdef CDatum out
 
     with nogil:
-        check_status(Sum(_context(), CDatum(array.sp_array), &out))
+        out = GetResultValue(Sum(CDatum(array.sp_array)))
 
     return wrap_datum(out)
 
@@ -40,7 +40,7 @@ cdef _sum_chunked_array(array: ChunkedArray):
     cdef CDatum out
 
     with nogil:
-        check_status(Sum(_context(), CDatum(array.sp_chunked_array), &out))
+        out = GetResultValue(Sum(CDatum(array.sp_chunked_array)))
 
     return wrap_datum(out)
 

python/pyarrow/array.pxi

@@ -376,19 +376,6 @@ cdef Py_ssize_t _normalize_index(Py_ssize_t index,
     return index
 
 
-cdef class _FunctionContext:
-    cdef:
-        unique_ptr[CFunctionContext] ctx
-
-    def __cinit__(self):
-        self.ctx.reset(new CFunctionContext(c_default_memory_pool()))
-
-cdef _FunctionContext _global_ctx = _FunctionContext()
-
-cdef CFunctionContext* _context() nogil:
-    return _global_ctx.ctx.get()
-
-
 cdef wrap_datum(const CDatum& datum):
     if datum.kind() == DatumType_ARRAY:
         return pyarrow_wrap_array(MakeArray(datum.array()))
@@ -705,9 +692,7 @@ cdef class Array(_PandasConvertible):
             shared_ptr[CArray] result
 
         with nogil:
-            check_status(Cast(_context(), self.ap[0], type.sp_type,
-                              options, &result))
-
+            result = GetResultValue(Cast(self.ap[0], type.sp_type, options))
         return pyarrow_wrap_array(result)
 
     def view(self, object target_type):
@@ -736,10 +721,8 @@ cdef class Array(_PandasConvertible):
         Sum the values in a numerical array.
         """
         cdef CDatum out
-
         with nogil:
-            check_status(Sum(_context(), CDatum(self.sp_array), &out))
-
+            out = GetResultValue(Sum(CDatum(self.sp_array)))
         return wrap_datum(out)
 
     def unique(self):
@@ -749,7 +732,7 @@ cdef class Array(_PandasConvertible):
         cdef shared_ptr[CArray] result
 
         with nogil:
-            check_status(Unique(_context(), CDatum(self.sp_array), &result))
+            result = GetResultValue(Unique(CDatum(self.sp_array)))
 
         return pyarrow_wrap_array(result)
 
@@ -760,9 +743,7 @@ cdef class Array(_PandasConvertible):
         cdef CDatum out
 
         with nogil:
-            check_status(DictionaryEncode(_context(), CDatum(self.sp_array),
-                                          &out))
-
+            out = GetResultValue(DictionaryEncode(CDatum(self.sp_array)))
         return wrap_datum(out)
 
     def value_counts(self):
@@ -776,8 +757,7 @@ cdef class Array(_PandasConvertible):
         cdef shared_ptr[CArray] result
 
         with nogil:
-            check_status(ValueCounts(_context(), CDatum(self.sp_array),
-                                     &result))
+            result = GetResultValue(ValueCounts(CDatum(self.sp_array)))
         return pyarrow_wrap_array(result)
 
     @staticmethod
@@ -1040,8 +1020,8 @@ cdef class Array(_PandasConvertible):
         c_indices = asarray(indices)
 
         with nogil:
-            check_status(Take(_context(), CDatum(self.sp_array),
-                              CDatum(c_indices.sp_array), options, &out))
+            out = GetResultValue(Take(CDatum(self.sp_array),
+                                      CDatum(c_indices.sp_array), options))
 
         return wrap_datum(out)
 
@@ -1091,8 +1071,8 @@ cdef class Array(_PandasConvertible):
         options = _convert_filter_option(null_selection_behavior)
 
         with nogil:
-            check_status(FilterKernel(_context(), CDatum(self.sp_array),
-                                      CDatum(mask.sp_array), options, &out))
+            out = GetResultValue(FilterKernel(CDatum(self.sp_array),
+                                              CDatum(mask.sp_array), options))
 
         return wrap_datum(out)