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))