Fix Take on LargeBinary / LargeString data

cpp/src/arrow/array/builder_binary.h

@@ -387,6 +387,8 @@ class ARROW_EXPORT LargeStringBuilder : public LargeBinaryBuilder {
 
 class ARROW_EXPORT FixedSizeBinaryBuilder : public ArrayBuilder {
  public:
+  using TypeClass = FixedSizeBinaryType;
+
   FixedSizeBinaryBuilder(const std::shared_ptr<DataType>& type,
                          MemoryPool* pool ARROW_MEMORY_POOL_DEFAULT);
 

cpp/src/arrow/array/builder_decimal.h

@@ -28,6 +28,8 @@ class Decimal128;
 
 class ARROW_EXPORT Decimal128Builder : public FixedSizeBinaryBuilder {
  public:
+  using TypeClass = Decimal128Type;
+
   explicit Decimal128Builder(const std::shared_ptr<DataType>& type,
                              MemoryPool* pool ARROW_MEMORY_POOL_DEFAULT);
 

cpp/src/arrow/array/builder_primitive.h

@@ -58,6 +58,7 @@ class ARROW_EXPORT NullBuilder : public ArrayBuilder {
 template <typename T>
 class NumericBuilder : public ArrayBuilder {
  public:
+  using TypeClass = T;
   using value_type = typename T::c_type;
   using ArrayType = typename TypeTraits<T>::ArrayType;
   using ArrayBuilder::ArrayBuilder;
@@ -265,7 +266,9 @@ using DoubleBuilder = NumericBuilder<DoubleType>;
 
 class ARROW_EXPORT BooleanBuilder : public ArrayBuilder {
  public:
+  using TypeClass = BooleanType;
   using value_type = bool;
+
   explicit BooleanBuilder(MemoryPool* pool ARROW_MEMORY_POOL_DEFAULT);
 
   explicit BooleanBuilder(const std::shared_ptr<DataType>& type, MemoryPool* pool);

cpp/src/arrow/array/builder_time.h

@@ -34,6 +34,7 @@ namespace arrow {
 
 class ARROW_EXPORT DayTimeIntervalBuilder : public ArrayBuilder {
  public:
+  using TypeClass = DayTimeIntervalType;
   using DayMilliseconds = DayTimeIntervalType::DayMilliseconds;
 
   explicit DayTimeIntervalBuilder(MemoryPool* pool ARROW_MEMORY_POOL_DEFAULT)

cpp/src/arrow/compute/kernels/filter-test.cc

@@ -285,18 +285,26 @@ TYPED_TEST(TestFilterKernelWithNumeric, ScalarInRangeAndFilterRandomNumeric) {
   }
 }
 
-class TestFilterKernelWithString : public TestFilterKernel<StringType> {
+using StringTypes =
+    ::testing::Types<BinaryType, StringType, LargeBinaryType, LargeStringType>;
+
+template <typename TypeClass>
+class TestFilterKernelWithString : public TestFilterKernel<TypeClass> {
  protected:
+  std::shared_ptr<DataType> value_type() {
+    return TypeTraits<TypeClass>::type_singleton();
+  }
+
   void AssertFilter(const std::string& values, const std::string& filter,
                     const std::string& expected) {
-    TestFilterKernel<StringType>::AssertFilter(utf8(), values, filter, expected);
+    TestFilterKernel<TypeClass>::AssertFilter(value_type(), values, filter, expected);
   }
   void AssertFilterDictionary(const std::string& dictionary_values,
                               const std::string& dictionary_filter,
                               const std::string& filter,
                               const std::string& expected_filter) {
-    auto dict = ArrayFromJSON(utf8(), dictionary_values);
-    auto type = dictionary(int8(), utf8());
+    auto dict = ArrayFromJSON(value_type(), dictionary_values);
+    auto type = dictionary(int8(), value_type());
     std::shared_ptr<Array> values, actual, expected;
     ASSERT_OK(DictionaryArray::FromArrays(type, ArrayFromJSON(int8(), dictionary_filter),
                                           dict, &values));
@@ -307,13 +315,15 @@ class TestFilterKernelWithString : public TestFilterKernel<StringType> {
   }
 };
 
-TEST_F(TestFilterKernelWithString, FilterString) {
+TYPED_TEST_CASE(TestFilterKernelWithString, StringTypes);
+
+TYPED_TEST(TestFilterKernelWithString, FilterString) {
   this->AssertFilter(R"(["a", "b", "c"])", "[0, 1, 0]", R"(["b"])");
   this->AssertFilter(R"([null, "b", "c"])", "[0, 1, 0]", R"(["b"])");
   this->AssertFilter(R"(["a", "b", "c"])", "[null, 1, 0]", R"([null, "b"])");
 }
 
-TEST_F(TestFilterKernelWithString, FilterDictionary) {
+TYPED_TEST(TestFilterKernelWithString, FilterDictionary) {
   auto dict = R"(["a", "b", "c", "d", "e"])";
   this->AssertFilterDictionary(dict, "[3, 4, 2]", "[0, 1, 0]", "[4]");
   this->AssertFilterDictionary(dict, "[null, 4, 2]", "[0, 1, 0]", "[4]");

cpp/src/arrow/compute/kernels/take-internal.h

@@ -20,11 +20,13 @@
 #include <algorithm>
 #include <limits>
 #include <memory>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
 #include "arrow/builder.h"
 #include "arrow/compute/context.h"
+#include "arrow/type_traits.h"
 #include "arrow/util/bit-util.h"
 #include "arrow/util/checked_cast.h"
 #include "arrow/util/logging.h"
@@ -37,21 +39,19 @@ namespace compute {
 using internal::checked_cast;
 using internal::checked_pointer_cast;
 
+// For non-binary builders, use regular value append
 template <typename Builder, typename Scalar>
-static Status UnsafeAppend(Builder* builder, Scalar&& value) {
+static typename std::enable_if<
+    !std::is_base_of<BaseBinaryType, typename Builder::TypeClass>::value, Status>::type
+UnsafeAppend(Builder* builder, Scalar&& value) {
   builder->UnsafeAppend(std::forward<Scalar>(value));
   return Status::OK();
 }
 
-// Use BinaryBuilder::UnsafeAppend, but reserve byte storage first
-static Status UnsafeAppend(BinaryBuilder* builder, util::string_view value) {
-  RETURN_NOT_OK(builder->ReserveData(static_cast<int64_t>(value.size())));
-  builder->UnsafeAppend(value);
-  return Status::OK();
-}
-
-// Use StringBuilder::UnsafeAppend, but reserve character storage first
-static Status UnsafeAppend(StringBuilder* builder, util::string_view value) {
+// For binary builders, need to reserve byte storage first
+template <typename Builder>
+static enable_if_base_binary<typename Builder::TypeClass, Status> UnsafeAppend(
+    Builder* builder, util::string_view value) {
   RETURN_NOT_OK(builder->ReserveData(static_cast<int64_t>(value.size())));
   builder->UnsafeAppend(value);
   return Status::OK();

cpp/src/arrow/compute/kernels/take-test.cc

@@ -179,18 +179,26 @@ TYPED_TEST(TestTakeKernelWithNumeric, TakeRandomNumeric) {
   }
 }
 
-class TestTakeKernelWithString : public TestTakeKernel<StringType> {
- protected:
+using StringTypes =
+    ::testing::Types<BinaryType, StringType, LargeBinaryType, LargeStringType>;
+
+template <typename TypeClass>
+class TestTakeKernelWithString : public TestTakeKernel<TypeClass> {
+ public:
+  std::shared_ptr<DataType> value_type() {
+    return TypeTraits<TypeClass>::type_singleton();
+  }
+
   void AssertTake(const std::string& values, const std::string& indices,
                   const std::string& expected) {
-    TestTakeKernel<StringType>::AssertTake(utf8(), values, indices, expected);
+    TestTakeKernel<TypeClass>::AssertTake(value_type(), values, indices, expected);
   }
   void AssertTakeDictionary(const std::string& dictionary_values,
                             const std::string& dictionary_indices,
                             const std::string& indices,
                             const std::string& expected_indices) {
-    auto dict = ArrayFromJSON(utf8(), dictionary_values);
-    auto type = dictionary(int8(), utf8());
+    auto dict = ArrayFromJSON(value_type(), dictionary_values);
+    auto type = dictionary(int8(), value_type());
     std::shared_ptr<Array> values, actual, expected;
     ASSERT_OK(DictionaryArray::FromArrays(type, ArrayFromJSON(int8(), dictionary_indices),
                                           dict, &values));
@@ -201,19 +209,22 @@ class TestTakeKernelWithString : public TestTakeKernel<StringType> {
   }
 };
 
-TEST_F(TestTakeKernelWithString, TakeString) {
+TYPED_TEST_CASE(TestTakeKernelWithString, StringTypes);
+
+TYPED_TEST(TestTakeKernelWithString, TakeString) {
   this->AssertTake(R"(["a", "b", "c"])", "[0, 1, 0]", R"(["a", "b", "a"])");
   this->AssertTake(R"([null, "b", "c"])", "[0, 1, 0]", "[null, \"b\", null]");
   this->AssertTake(R"(["a", "b", "c"])", "[null, 1, 0]", R"([null, "b", "a"])");
 
+  std::shared_ptr<DataType> type = this->value_type();
   std::shared_ptr<Array> arr;
   ASSERT_RAISES(IndexError,
-                this->Take(utf8(), R"(["a", "b", "c"])", int8(), "[0, 9, 0]", &arr));
-  ASSERT_RAISES(IndexError, this->Take(utf8(), R"(["a", "b", null, "ddd", "ee"])",
-                                       int64(), "[2, 5]", &arr));
+                this->Take(type, R"(["a", "b", "c"])", int8(), "[0, 9, 0]", &arr));
+  ASSERT_RAISES(IndexError, this->Take(type, R"(["a", "b", null, "ddd", "ee"])", int64(),
+                                       "[2, 5]", &arr));
 }
 
-TEST_F(TestTakeKernelWithString, TakeDictionary) {
+TYPED_TEST(TestTakeKernelWithString, TakeDictionary) {
   auto dict = R"(["a", "b", "c", "d", "e"])";
   this->AssertTakeDictionary(dict, "[3, 4, 2]", "[0, 1, 0]", "[3, 4, 3]");
   this->AssertTakeDictionary(dict, "[null, 4, 2]", "[0, 1, 0]", "[null, 4, null]");