Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 1 addition & 3 deletions cpp/src/arrow/compute/kernels/scalar_cast_internal.cc
Original file line number Diff line number Diff line change
Expand Up @@ -280,10 +280,8 @@ void AddZeroCopyCast(Type::type in_type_id, InputType in_type, OutputType out_ty
}

static bool CanCastFromDictionary(Type::type type_id) {
/// TODO(GH-43010): add is_binary_view_like() here once array_take
/// can handle string-views
return (is_primitive(type_id) || is_base_binary_like(type_id) ||
is_fixed_size_binary(type_id));
is_binary_view_like(type_id) || is_fixed_size_binary(type_id));
}

void AddCommonCasts(Type::type out_type_id, OutputType out_ty, CastFunction* func) {
Expand Down
47 changes: 47 additions & 0 deletions cpp/src/arrow/compute/kernels/scalar_cast_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -4368,6 +4368,53 @@ TEST(Cast, FromDictionary) {
}
}

TEST(Cast, DictionaryDecodeFromViewDictionary) {
for (const auto& value_type : {binary_view(), utf8_view()}) {
ARROW_SCOPED_TRACE(value_type->ToString());
auto dict_values = ArrayFromJSON(
value_type, R"(["alpha", "long-value-over-inline-limit", "omega"])");
auto indices = ArrayFromJSON(int8(), "[0, 1, null, 2, 1]");
ASSERT_OK_AND_ASSIGN(auto dict_arr,
DictionaryArray::FromArrays(dictionary(int8(), value_type),
indices, dict_values));
auto expected = ArrayFromJSON(
value_type,
R"(["alpha", "long-value-over-inline-limit", null, "omega", "long-value-over-inline-limit"])");

ASSERT_OK_AND_ASSIGN(Datum decoded, CallFunction("dictionary_decode", {dict_arr}));
ValidateOutput(decoded);
AssertArraysEqual(*expected, *decoded.make_array(), /*verbose=*/true);
CheckCast(dict_arr, expected);

auto chunked_dict = std::make_shared<ChunkedArray>(
ArrayVector{dict_arr->Slice(0, 2), dict_arr->Slice(2, 3)});
ASSERT_OK_AND_ASSIGN(Datum decoded_chunked,
CallFunction("dictionary_decode", {chunked_dict}));
ValidateOutput(decoded_chunked);
AssertChunkedEqual(
*ChunkedArrayFromJSON(value_type,
{R"(["alpha", "long-value-over-inline-limit"])",
R"([null, "omega", "long-value-over-inline-limit"])"}),
*decoded_chunked.chunked_array());

auto dict_values_with_null = ArrayFromJSON(value_type, R"(["alpha", null, "omega"])");
auto indices_with_null_source = ArrayFromJSON(int8(), "[0, 1, null, 2]");
ASSERT_OK_AND_ASSIGN(
auto dict_arr_with_null,
DictionaryArray::FromArrays(dictionary(int8(), value_type),
indices_with_null_source, dict_values_with_null));
auto expected_with_null_source =
ArrayFromJSON(value_type, R"(["alpha", null, null, "omega"])");

ASSERT_OK_AND_ASSIGN(Datum decoded_with_null_source,
CallFunction("dictionary_decode", {dict_arr_with_null}));
ValidateOutput(decoded_with_null_source);
AssertArraysEqual(*expected_with_null_source, *decoded_with_null_source.make_array(),
/*verbose=*/true);
CheckCast(dict_arr_with_null, expected_with_null_source);
}
}

std::shared_ptr<Array> SmallintArrayFromJSON(const std::string& json_data) {
auto arr = ArrayFromJSON(int16(), json_data);
auto ext_data = arr->data()->Copy();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -919,6 +919,11 @@ Status SparseUnionFilterExec(KernelContext* ctx, const ExecSpan& batch, ExecResu
return FilterWithTakeExec(SparseUnionTakeExec, ctx, batch, out);
}

Status VarBinaryViewFilterExec(KernelContext* ctx, const ExecSpan& batch,
ExecResult* out) {
return FilterWithTakeExec(VarBinaryViewTakeExec, ctx, batch, out);
}

// ----------------------------------------------------------------------
// Implement Filter metafunction

Expand Down Expand Up @@ -1094,6 +1099,8 @@ void PopulateFilterKernels(std::vector<SelectionKernelData>* out) {
{InputType(match::Primitive()), plain_filter, PrimitiveFilterExec},
{InputType(match::BinaryLike()), plain_filter, BinaryFilterExec},
{InputType(match::LargeBinaryLike()), plain_filter, BinaryFilterExec},
{InputType(Type::BINARY_VIEW), plain_filter, VarBinaryViewFilterExec},
{InputType(Type::STRING_VIEW), plain_filter, VarBinaryViewFilterExec},
{InputType(null()), plain_filter, NullFilterExec},
{InputType(Type::FIXED_SIZE_BINARY), plain_filter, PrimitiveFilterExec},
{InputType(Type::DECIMAL32), plain_filter, PrimitiveFilterExec},
Expand All @@ -1116,6 +1123,8 @@ void PopulateFilterKernels(std::vector<SelectionKernelData>* out) {
{InputType(match::Primitive()), ree_filter, PrimitiveFilterExec},
{InputType(match::BinaryLike()), ree_filter, BinaryFilterExec},
{InputType(match::LargeBinaryLike()), ree_filter, BinaryFilterExec},
{InputType(Type::BINARY_VIEW), ree_filter, VarBinaryViewFilterExec},
{InputType(Type::STRING_VIEW), ree_filter, VarBinaryViewFilterExec},
{InputType(null()), ree_filter, NullFilterExec},
{InputType(Type::FIXED_SIZE_BINARY), ree_filter, PrimitiveFilterExec},
{InputType(Type::DECIMAL32), ree_filter, PrimitiveFilterExec},
Expand Down
1 change: 1 addition & 0 deletions cpp/src/arrow/compute/kernels/vector_selection_internal.h
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@ Status MapFilterExec(KernelContext*, const ExecSpan&, ExecResult*);

Status VarBinaryTakeExec(KernelContext*, const ExecSpan&, ExecResult*);
Status LargeVarBinaryTakeExec(KernelContext*, const ExecSpan&, ExecResult*);
Status VarBinaryViewTakeExec(KernelContext*, const ExecSpan&, ExecResult*);
Status FixedWidthTakeExec(KernelContext*, const ExecSpan&, ExecResult*);
Status ListTakeExec(KernelContext*, const ExecSpan&, ExecResult*);
Status LargeListTakeExec(KernelContext*, const ExecSpan&, ExecResult*);
Expand Down
120 changes: 120 additions & 0 deletions cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@

#include "arrow/array/builder_primitive.h"
#include "arrow/array/concatenate.h"
#include "arrow/buffer.h"
#include "arrow/buffer_builder.h"
#include "arrow/chunked_array.h"
#include "arrow/compute/api_vector.h"
Expand Down Expand Up @@ -488,6 +489,123 @@ Status FixedWidthTakeExec(KernelContext* ctx, const ExecSpan& batch, ExecResult*

namespace {

template <typename IndexCType>
Status VarBinaryViewTakeTyped(const ArraySpan& values, const ArraySpan& indices,
BinaryViewType::c_type* out_views, uint8_t* out_validity,
int64_t* valid_count) {
const auto* source_views = values.GetValues<BinaryViewType::c_type>(1);
const auto* index_values = indices.GetValues<IndexCType>(1);

const bool values_may_have_nulls = values.MayHaveNulls();
const bool indices_may_have_nulls = indices.MayHaveNulls();

if (!values_may_have_nulls && !indices_may_have_nulls) {
for (int64_t out_i = 0; out_i < indices.length; ++out_i) {
out_views[out_i] = source_views[static_cast<int64_t>(index_values[out_i])];
}
*valid_count = indices.length;
return Status::OK();
}

for (int64_t out_i = 0; out_i < indices.length; ++out_i) {
if (indices_may_have_nulls &&
!bit_util::GetBit(indices.buffers[0].data, indices.offset + out_i)) {
continue;
}

const int64_t source_i = static_cast<int64_t>(index_values[out_i]);
const bool source_valid =
!values_may_have_nulls ||
bit_util::GetBit(values.buffers[0].data, values.offset + source_i);
if (!source_valid) {
continue;
}

out_views[out_i] = source_views[source_i];
if (out_validity != nullptr) {
bit_util::SetBit(out_validity, out_i);
}
++(*valid_count);
}

return Status::OK();
}

Status VarBinaryViewTakeDispatch(const ArraySpan& values, const ArraySpan& indices,
BinaryViewType::c_type* out_views, uint8_t* out_validity,
int64_t* valid_count) {
switch (indices.type->byte_width()) {
case 1:
return VarBinaryViewTakeTyped<uint8_t>(values, indices, out_views, out_validity,
valid_count);
case 2:
return VarBinaryViewTakeTyped<uint16_t>(values, indices, out_views, out_validity,
valid_count);
case 4:
return VarBinaryViewTakeTyped<uint32_t>(values, indices, out_views, out_validity,
valid_count);
default:
DCHECK_EQ(indices.type->byte_width(), 8);
return VarBinaryViewTakeTyped<uint64_t>(values, indices, out_views, out_validity,
valid_count);
}
}

} // namespace

Status VarBinaryViewTakeExec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
const ArraySpan& values = batch[0].array;
const ArraySpan& indices = batch[1].array;

if (TakeState::Get(ctx).boundscheck) {
RETURN_NOT_OK(CheckIndexBounds(indices, values.length));
}

const int64_t out_length = indices.length;
const bool may_have_nulls = values.MayHaveNulls() || indices.MayHaveNulls();
const auto data_buffers = values.GetVariadicBuffers();

ARROW_ASSIGN_OR_RAISE(
auto views_buf,
AllocateBuffer(out_length * static_cast<int64_t>(sizeof(BinaryViewType::c_type)),
ctx->memory_pool()));
auto* out_views = reinterpret_cast<BinaryViewType::c_type*>(views_buf->mutable_data());
if (may_have_nulls && views_buf->size() > 0) {
std::memset(out_views, 0, views_buf->size());
}

std::shared_ptr<Buffer> validity_buf;
uint8_t* out_validity = nullptr;
if (may_have_nulls) {
ARROW_ASSIGN_OR_RAISE(validity_buf,
AllocateEmptyBitmap(out_length, ctx->memory_pool()));
if (validity_buf->size() > 0) {
std::memset(validity_buf->mutable_data(), 0, validity_buf->size());
}
out_validity = validity_buf->mutable_data();
}

int64_t valid_count = 0;
RETURN_NOT_OK(
VarBinaryViewTakeDispatch(values, indices, out_views, out_validity, &valid_count));

const int64_t null_count = out_length - valid_count;
BufferVector buffers;
buffers.reserve(2 + data_buffers.size());
buffers.push_back(null_count == 0 ? nullptr : std::move(validity_buf));
buffers.push_back(std::move(views_buf));

for (const auto& data_buffer : data_buffers) {
buffers.push_back(data_buffer);
}

out->value = ArrayData::Make(values.type->GetSharedPtr(), out_length,
std::move(buffers), null_count, /*offset=*/0);
return Status::OK();
}

namespace {

// ----------------------------------------------------------------------
// Null take

Expand Down Expand Up @@ -740,6 +858,8 @@ void PopulateTakeKernels(std::vector<SelectionKernelData>* out) {
{InputType(match::Primitive()), take_indices, FixedWidthTakeExec},
{InputType(match::BinaryLike()), take_indices, VarBinaryTakeExec},
{InputType(match::LargeBinaryLike()), take_indices, LargeVarBinaryTakeExec},
{InputType(Type::BINARY_VIEW), take_indices, VarBinaryViewTakeExec},
{InputType(Type::STRING_VIEW), take_indices, VarBinaryViewTakeExec},
{InputType(match::FixedSizeBinaryLike()), take_indices, FixedWidthTakeExec},
{InputType(null()), take_indices, NullTakeExec},
{InputType(Type::DICTIONARY), take_indices, DictionaryTake},
Expand Down
Loading
Loading