ARROW-15545: [Python][C++] Support casting to extension type

cpp/src/arrow/CMakeLists.txt

@@ -425,6 +425,7 @@ if(ARROW_COMPUTE)
        compute/kernels/scalar_boolean.cc
        compute/kernels/scalar_cast_boolean.cc
        compute/kernels/scalar_cast_dictionary.cc
+       compute/kernels/scalar_cast_extension.cc
        compute/kernels/scalar_cast_internal.cc
        compute/kernels/scalar_cast_nested.cc
        compute/kernels/scalar_cast_numeric.cc

cpp/src/arrow/compute/cast.cc

@@ -62,6 +62,7 @@ void InitCastTable() {
   AddCastFunctions(GetNumericCasts());
   AddCastFunctions(GetTemporalCasts());
   AddCastFunctions(GetDictionaryCasts());
+  AddCastFunctions(GetExtensionCasts());
 }
 
 void EnsureInitCastTable() { std::call_once(cast_table_initialized, InitCastTable); }

cpp/src/arrow/compute/cast_internal.h

@@ -63,6 +63,7 @@ std::vector<std::shared_ptr<CastFunction>> GetTemporalCasts();
 std::vector<std::shared_ptr<CastFunction>> GetBinaryLikeCasts();
 std::vector<std::shared_ptr<CastFunction>> GetNestedCasts();
 std::vector<std::shared_ptr<CastFunction>> GetDictionaryCasts();
+std::vector<std::shared_ptr<CastFunction>> GetExtensionCasts();
 
 ARROW_EXPORT
 Result<std::shared_ptr<CastFunction>> GetCastFunction(const DataType& to_type);

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

@@ -0,0 +1,74 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Implementation of casting to extension types
+#include "arrow/compute/kernels/common.h"
+#include "arrow/compute/kernels/scalar_cast_internal.h"
+#include "arrow/scalar.h"
+
+namespace arrow {
+namespace compute {
+namespace internal {
+
+namespace {
+Status CastToExtension(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
+  const CastOptions& options = checked_cast<const CastState*>(ctx->state())->options;
+  const auto& ext_ty = static_cast<const ExtensionType&>(*options.to_type.type);
+  auto out_ty = ext_ty.storage_type();
+
+  DCHECK(batch[0].is_array());
+  std::shared_ptr<Array> array = batch[0].array.ToArray();
+
+  // Try to prevent user errors by preventing casting between extensions w/
+  // different storage types. Provide a tip on how to accomplish same outcome.
+  std::shared_ptr<Array> result;
+  if (array->type()->id() == Type::EXTENSION) {
+    if (!array->type()->Equals(out_ty)) {
+      return Status::TypeError("Casting from '" + array->type()->ToString() +
+                               "' to different extension type '" + ext_ty.ToString() +
+                               "' not permitted. One can first cast to the storage "
+                               "type, then to the extension type.");
+    }
+    result = array;
+  } else {
+    ARROW_ASSIGN_OR_RAISE(result, Cast(*array, out_ty, options, ctx->exec_context()));
+  }
+
+  ExtensionArray extension(options.to_type.GetSharedPtr(), result);
+  out->value = std::move(extension.data());
+  return Status::OK();
+}
+
+std::shared_ptr<CastFunction> GetCastToExtension(std::string name) {
+  auto func = std::make_shared<CastFunction>(std::move(name), Type::EXTENSION);
+  for (Type::type in_ty : AllTypeIds()) {
+    DCHECK_OK(
+        func->AddKernel(in_ty, {InputType(in_ty)}, kOutputTargetType, CastToExtension));
+  }
+  return func;
+}
+
+};  // namespace
+
+std::vector<std::shared_ptr<CastFunction>> GetExtensionCasts() {
+  auto func = GetCastToExtension("cast_extension");
+  return {func};
+}
+
+}  // namespace internal
+}  // namespace compute
+}  // namespace arrow

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

@@ -225,7 +225,8 @@ TEST(Cast, CanCast) {
   ExpectCanCast(smallint(), {int16()});  // cast storage
   ExpectCanCast(smallint(),
                 kNumericTypes);  // any cast which is valid for storage is supported
-  ExpectCannotCast(null(), {smallint()});  // FIXME missing common cast from null
+  ExpectCanCast(null(), {smallint()});
+  ExpectCanCast(tinyint(), {smallint()});  // cast between compatible storage types
 
   ExpectCanCast(date32(), {utf8(), large_utf8()});
   ExpectCanCast(date64(), {utf8(), large_utf8()});
@@ -2728,6 +2729,13 @@ std::shared_ptr<Array> SmallintArrayFromJSON(const std::string& json_data) {
   return MakeArray(ext_data);
 }
 
+std::shared_ptr<Array> TinyintArrayFromJSON(const std::string& json_data) {
+  auto arr = ArrayFromJSON(int8(), json_data);
+  auto ext_data = arr->data()->Copy();
+  ext_data->type = tinyint();
+  return MakeArray(ext_data);
+}
+
 TEST(Cast, ExtensionTypeToIntDowncast) {
   auto smallint = std::make_shared<SmallintType>();
   ExtensionTypeGuard smallint_guard(smallint);
@@ -2765,6 +2773,68 @@ TEST(Cast, ExtensionTypeToIntDowncast) {
   }
 }
 
+TEST(Cast, PrimitiveToExtension) {
+  {
+    auto primitive_array = ArrayFromJSON(uint8(), "[0, 1, 3]");
+    auto extension_array = SmallintArrayFromJSON("[0, 1, 3]");
+    CastOptions options;
+    options.to_type = smallint();
+    CheckCast(primitive_array, extension_array, options);
+  }
+  {
+    CastOptions options;
+    options.to_type = smallint();
+    CheckCastFails(ArrayFromJSON(utf8(), "[\"hello\"]"), options);
+  }
+}
+
+TEST(Cast, ExtensionDictToExtension) {
+  auto extension_array = SmallintArrayFromJSON("[1, 2, 1]");
+  auto indices_array = ArrayFromJSON(int32(), "[0, 1, 0]");
+
+  ASSERT_OK_AND_ASSIGN(auto dict_array,
+                       DictionaryArray::FromArrays(indices_array, extension_array));
+
+  CastOptions options;
+  options.to_type = smallint();
+  CheckCast(dict_array, extension_array, options);
+}
+
+TEST(Cast, IntToExtensionTypeDowncast) {
+  CheckCast(ArrayFromJSON(uint8(), "[0, 100, 200, 1, 2]"),
+            SmallintArrayFromJSON("[0, 100, 200, 1, 2]"));
+
+  // int32 to Smallint(int16), with overflow
+  {
+    CastOptions options;
+    options.to_type = smallint();
+    CheckCastFails(ArrayFromJSON(int32(), "[0, null, 32768, 1, 3]"), options);
+
+    options.allow_int_overflow = true;
+    CheckCast(ArrayFromJSON(int32(), "[0, null, 32768, 1, 3]"),
+              SmallintArrayFromJSON("[0, null, -32768, 1, 3]"), options);
+  }
+
+  // int32 to Smallint(int16), with underflow
+  {
+    CastOptions options;
+    options.to_type = smallint();
+    CheckCastFails(ArrayFromJSON(int32(), "[0, null, -32769, 1, 3]"), options);
+
+    options.allow_int_overflow = true;
+    CheckCast(ArrayFromJSON(int32(), "[0, null, -32769, 1, 3]"),
+              SmallintArrayFromJSON("[0, null, 32767, 1, 3]"), options);
+  }
+
+  // Cannot cast between extension types when storage types differ
+  {
+    CastOptions options;
+    options.to_type = smallint();
+    auto tiny_array = TinyintArrayFromJSON("[0, 1, 3]");
+    ASSERT_NOT_OK(Cast(tiny_array, smallint(), options));
+  }
+}
+
 TEST(Cast, DictTypeToAnotherDict) {
   auto check_cast = [&](const std::shared_ptr<DataType>& in_type,
                         const std::shared_ptr<DataType>& out_type,

cpp/src/arrow/testing/extension_type.h

@@ -54,6 +54,11 @@ class ARROW_TESTING_EXPORT SmallintArray : public ExtensionArray {
   using ExtensionArray::ExtensionArray;
 };
 
+class ARROW_TESTING_EXPORT TinyintArray : public ExtensionArray {
+ public:
+  using ExtensionArray::ExtensionArray;
+};
+
 class ARROW_TESTING_EXPORT ListExtensionArray : public ExtensionArray {
  public:
   using ExtensionArray::ExtensionArray;
@@ -76,6 +81,23 @@ class ARROW_TESTING_EXPORT SmallintType : public ExtensionType {
   std::string Serialize() const override { return "smallint"; }
 };
 
+class ARROW_TESTING_EXPORT TinyintType : public ExtensionType {
+ public:
+  TinyintType() : ExtensionType(int8()) {}
+
+  std::string extension_name() const override { return "tinyint"; }
+
+  bool ExtensionEquals(const ExtensionType& other) const override;
+
+  std::shared_ptr<Array> MakeArray(std::shared_ptr<ArrayData> data) const override;
+
+  Result<std::shared_ptr<DataType>> Deserialize(
+      std::shared_ptr<DataType> storage_type,
+      const std::string& serialized) const override;
+
+  std::string Serialize() const override { return "tinyint"; }
+};
+
 class ARROW_TESTING_EXPORT ListExtensionType : public ExtensionType {
  public:
   ListExtensionType() : ExtensionType(list(int32())) {}
@@ -140,6 +162,9 @@ std::shared_ptr<DataType> uuid();
 ARROW_TESTING_EXPORT
 std::shared_ptr<DataType> smallint();
 
+ARROW_TESTING_EXPORT
+std::shared_ptr<DataType> tinyint();
+
 ARROW_TESTING_EXPORT
 std::shared_ptr<DataType> list_extension_type();
 
@@ -155,6 +180,9 @@ std::shared_ptr<Array> ExampleUuid();
 ARROW_TESTING_EXPORT
 std::shared_ptr<Array> ExampleSmallint();
 
+ARROW_TESTING_EXPORT
+std::shared_ptr<Array> ExampleTinyint();
+
 ARROW_TESTING_EXPORT
 std::shared_ptr<Array> ExampleDictExtension();
 

cpp/src/arrow/testing/gtest_util.cc

@@ -64,47 +64,6 @@ namespace arrow {
 using internal::checked_cast;
 using internal::checked_pointer_cast;
 
-std::vector<Type::type> AllTypeIds() {
-  return {Type::NA,
-          Type::BOOL,
-          Type::INT8,
-          Type::INT16,
-          Type::INT32,
-          Type::INT64,
-          Type::UINT8,
-          Type::UINT16,
-          Type::UINT32,
-          Type::UINT64,
-          Type::HALF_FLOAT,
-          Type::FLOAT,
-          Type::DOUBLE,
-          Type::DECIMAL128,
-          Type::DECIMAL256,
-          Type::DATE32,
-          Type::DATE64,
-          Type::TIME32,
-          Type::TIME64,
-          Type::TIMESTAMP,
-          Type::INTERVAL_DAY_TIME,
-          Type::INTERVAL_MONTHS,
-          Type::DURATION,
-          Type::STRING,
-          Type::BINARY,
-          Type::LARGE_STRING,
-          Type::LARGE_BINARY,
-          Type::FIXED_SIZE_BINARY,
-          Type::STRUCT,
-          Type::LIST,
-          Type::LARGE_LIST,
-          Type::FIXED_SIZE_LIST,
-          Type::MAP,
-          Type::DENSE_UNION,
-          Type::SPARSE_UNION,
-          Type::DICTIONARY,
-          Type::EXTENSION,
-          Type::INTERVAL_MONTH_DAY_NANO};
-}
-
 template <typename T, typename CompareFunctor>
 void AssertTsSame(const T& expected, const T& actual, CompareFunctor&& compare) {
   if (!compare(actual, expected)) {
@@ -832,6 +791,28 @@ Result<std::shared_ptr<DataType>> SmallintType::Deserialize(
   return std::make_shared<SmallintType>();
 }
 
+bool TinyintType::ExtensionEquals(const ExtensionType& other) const {
+  return (other.extension_name() == this->extension_name());
+}
+
+std::shared_ptr<Array> TinyintType::MakeArray(std::shared_ptr<ArrayData> data) const {
+  DCHECK_EQ(data->type->id(), Type::EXTENSION);
+  DCHECK_EQ("tinyint", static_cast<const ExtensionType&>(*data->type).extension_name());
+  return std::make_shared<TinyintArray>(data);
+}
+
+Result<std::shared_ptr<DataType>> TinyintType::Deserialize(
+    std::shared_ptr<DataType> storage_type, const std::string& serialized) const {
+  if (serialized != "tinyint") {
+    return Status::Invalid("Type identifier did not match: '", serialized, "'");
+  }
+  if (!storage_type->Equals(*int16())) {
+    return Status::Invalid("Invalid storage type for TinyintType: ",
+                           storage_type->ToString());
+  }
+  return std::make_shared<TinyintType>();
+}
+
 bool ListExtensionType::ExtensionEquals(const ExtensionType& other) const {
   return (other.extension_name() == this->extension_name());
 }
@@ -905,6 +886,8 @@ std::shared_ptr<DataType> uuid() { return std::make_shared<UuidType>(); }
 
 std::shared_ptr<DataType> smallint() { return std::make_shared<SmallintType>(); }
 
+std::shared_ptr<DataType> tinyint() { return std::make_shared<TinyintType>(); }
+
 std::shared_ptr<DataType> list_extension_type() {
   return std::make_shared<ListExtensionType>();
 }
@@ -936,6 +919,11 @@ std::shared_ptr<Array> ExampleSmallint() {
   return ExtensionType::WrapArray(smallint(), arr);
 }
 
+std::shared_ptr<Array> ExampleTinyint() {
+  auto arr = ArrayFromJSON(int8(), "[-128, null, 1, 2, 3, 4, 127]");
+  return ExtensionType::WrapArray(tinyint(), arr);
+}
+
 std::shared_ptr<Array> ExampleDictExtension() {
   auto arr = DictArrayFromJSON(dictionary(int8(), utf8()), "[0, 1, null, 1]",
                                R"(["foo", "bar"])");

cpp/src/arrow/testing/gtest_util.h

@@ -190,9 +190,6 @@ class RecordBatch;
 class Table;
 struct Datum;
 
-ARROW_TESTING_EXPORT
-std::vector<Type::type> AllTypeIds();
-
 #define ASSERT_ARRAYS_EQUAL(lhs, rhs) AssertArraysEqual((lhs), (rhs))
 #define ASSERT_BATCHES_EQUAL(lhs, rhs) AssertBatchesEqual((lhs), (rhs))
 #define ASSERT_BATCHES_APPROX_EQUAL(lhs, rhs) AssertBatchesApproxEqual((lhs), (rhs))