apacheGH-31868: [C++] Support concatenating extension arrays (apache#14463)

(still needs a test in C++ and more elaborate tests (eg extension types with nested data))
* Closes: apache#31868

Authored-by: Joris Van den Bossche <jorisvandenbossche@gmail.com>
Signed-off-by: Joris Van den Bossche <jorisvandenbossche@gmail.com>

Author: jorisvandenbossche

cpp/src/arrow/array/concatenate.cc

@@ -459,8 +459,16 @@ class ConcatenateImpl {
   }
 
   Status Visit(const ExtensionType& e) {
-    // XXX can we just concatenate their storage?
-    return Status::NotImplemented("concatenation of ", e);
+    ArrayDataVector storage_data(in_.size());
+    for (size_t i = 0; i < in_.size(); ++i) {
+      storage_data[i] = in_[i]->Copy();
+      storage_data[i]->type = e.storage_type();
+    }
+    std::shared_ptr<ArrayData> out_storage;
+    RETURN_NOT_OK(ConcatenateImpl(storage_data, pool_).Concatenate(&out_storage));
+    out_storage->type = in_[0]->type;
+    out_ = std::move(out_storage);
+    return Status::OK();
   }
 
  private:

cpp/src/arrow/array/concatenate_test.cc

@@ -36,6 +36,7 @@
 #include "arrow/buffer.h"
 #include "arrow/status.h"
 #include "arrow/testing/builder.h"
+#include "arrow/testing/extension_type.h"
 #include "arrow/testing/random.h"
 #include "arrow/testing/util.h"
 #include "arrow/type.h"
@@ -529,6 +530,13 @@ TEST_F(ConcatenateTest, DenseUnionType) {
       *concat_array_type_codes);
 }
 
+TEST_F(ConcatenateTest, ExtensionType) {
+  Check([this](int32_t size, double null_probability, std::shared_ptr<Array>* out) {
+    auto storage = this->GeneratePrimitive<Int16Type>(size, null_probability);
+    *out = ExtensionType::WrapArray(smallint(), storage);
+  });
+}
+
 TEST_F(ConcatenateTest, OffsetOverflow) {
   auto fake_long = ArrayFromJSON(utf8(), "[\"\"]");
   fake_long->data()->GetMutableValues<int32_t>(1)[1] =

python/pyarrow/tests/test_extension_type.py

@@ -630,6 +630,22 @@ def test_casting_dict_array_to_extension_type():
                                UUID('30313233-3435-3637-3839-616263646566')]
 
 
+def test_concat():
+    arr1 = pa.array([1, 2, 3], IntegerType())
+    arr2 = pa.array([4, 5, 6], IntegerType())
+
+    result = pa.concat_arrays([arr1, arr2])
+    expected = pa.array([1, 2, 3, 4, 5, 6], IntegerType())
+    assert result.equals(expected)
+
+    # nested in a struct
+    struct_arr1 = pa.StructArray.from_arrays([arr1], names=["a"])
+    struct_arr2 = pa.StructArray.from_arrays([arr2], names=["a"])
+    result = pa.concat_arrays([struct_arr1, struct_arr2])
+    expected = pa.StructArray.from_arrays([expected], names=["a"])
+    assert result.equals(expected)
+
+
 def test_null_storage_type():
     ext_type = AnnotatedType(pa.null(), {"key": "value"})
     storage = pa.array([None] * 10, pa.null())