Skip to content

Commit 2c707d4

Browse files
bkietzpitrou
bkietz
authored and
pitrou
committed
ARROW-9196: [C++][Compute] All casts accept scalar and sliced inputs
This is primarily a refactoring of the tests for "cast" to ensure that every case is also verified with scalar and sliced inputs through `CheckScalarUnary`. Caveat: ExtensionScalar is not implemented, so those checks aren't enabled. It should be straightforward for them to be enabled in a follow up, however. NB: also resolves ARROW-9198 Closes apache#9490 from bkietz/9196-Make-temporal-casts-work- Authored-by: Benjamin Kietzman <bengilgit@gmail.com> Signed-off-by: Antoine Pitrou <antoine@python.org>
1 parent 848c803 commit 2c707d4

11 files changed

Lines changed: 1219 additions & 1608 deletions

File tree

‎cpp/src/arrow/array/array_base.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ class ARROW_EXPORT Array {
9191
///
9292
/// Note that for `null_count == 0` or for null type, this will be null.
9393
/// This buffer does not account for any slice offset
94-
std::shared_ptr<Buffer> null_bitmap() const { return data_->buffers[0]; }
94+
const std::shared_ptr<Buffer>& null_bitmap() const { return data_->buffers[0]; }
9595

9696
/// Raw pointer to the null bitmap.
9797
///
@@ -160,7 +160,7 @@ class ARROW_EXPORT Array {
160160
/// Input-checking variant of Array::Slice
161161
Result<std::shared_ptr<Array>> SliceSafe(int64_t offset) const;
162162

163-
std::shared_ptr<ArrayData> data() const { return data_; }
163+
const std::shared_ptr<ArrayData>& data() const { return data_; }
164164

165165
int num_fields() const { return static_cast<int>(data_->child_data.size()); }
166166

‎cpp/src/arrow/compute/kernels/codegen_internal.h‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -663,14 +663,13 @@ struct ScalarUnaryNotNullStateful {
663663
static void Exec(const ThisType& functor, KernelContext* ctx, const ArrayData& arg0,
664664
Datum* out) {
665665
ArrayData* out_arr = out->mutable_array();
666-
auto out_data = out_arr->GetMutableValues<uint8_t>(1);
666+
auto out_data = out_arr->GetMutableValues<Decimal128>(1);
667667
VisitArrayValuesInline<Arg0Type>(
668668
arg0,
669669
[&](Arg0Value v) {
670-
functor.op.template Call<OutValue, Arg0Value>(ctx, v).ToBytes(out_data);
671-
out_data += 16;
670+
*out_data++ = functor.op.template Call<OutValue, Arg0Value>(ctx, v);
672671
},
673-
[&]() { out_data += 16; });
672+
[&]() { ++out_data; });
674673
}
675674
};
676675

‎cpp/src/arrow/compute/kernels/scalar_cast_nested.cc‎

Lines changed: 49 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -21,37 +21,71 @@
2121
#include <vector>
2222

2323
#include "arrow/array/builder_nested.h"
24+
#include "arrow/compute/api_scalar.h"
2425
#include "arrow/compute/cast.h"
2526
#include "arrow/compute/kernels/common.h"
2627
#include "arrow/compute/kernels/scalar_cast_internal.h"
28+
#include "arrow/util/bitmap_ops.h"
2729

2830
namespace arrow {
31+
32+
using internal::CopyBitmap;
33+
2934
namespace compute {
3035
namespace internal {
3136

3237
template <typename Type>
3338
void CastListExec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
34-
const CastOptions& options = checked_cast<const CastState&>(*ctx->state()).options;
39+
using offset_type = typename Type::offset_type;
40+
using ScalarType = typename TypeTraits<Type>::ScalarType;
41+
42+
const CastOptions& options = CastState::Get(ctx);
3543

36-
const ArrayData& input = *batch[0].array();
37-
ArrayData* result = out->mutable_array();
44+
auto child_type = checked_cast<const Type&>(*out->type()).value_type();
3845

39-
if (input.offset != 0) {
40-
ctx->SetStatus(Status::NotImplemented(
41-
"Casting sliced lists (non-zero offset) not yet implemented"));
46+
if (out->kind() == Datum::SCALAR) {
47+
const auto& in_scalar = checked_cast<const ScalarType&>(*batch[0].scalar());
48+
auto out_scalar = checked_cast<ScalarType*>(out->scalar().get());
49+
50+
DCHECK(!out_scalar->is_valid);
51+
if (in_scalar.is_valid) {
52+
KERNEL_ASSIGN_OR_RAISE(
53+
out_scalar->value, ctx,
54+
Cast(*in_scalar.value, child_type, options, ctx->exec_context()));
55+
56+
out_scalar->is_valid = true;
57+
}
4258
return;
4359
}
44-
// Copy buffers from parent
45-
result->buffers = input.buffers;
4660

47-
auto child_type = checked_cast<const Type&>(*result->type).value_type();
61+
const ArrayData& in_array = *batch[0].array();
62+
ArrayData* out_array = out->mutable_array();
63+
64+
// Copy from parent
65+
out_array->buffers = in_array.buffers;
66+
Datum values = in_array.child_data[0];
67+
68+
if (in_array.offset != 0) {
69+
KERNEL_ASSIGN_OR_RAISE(out_array->buffers[0], ctx,
70+
CopyBitmap(ctx->memory_pool(), in_array.buffers[0]->data(),
71+
in_array.offset, in_array.length));
72+
KERNEL_ASSIGN_OR_RAISE(out_array->buffers[1], ctx,
73+
ctx->Allocate(sizeof(offset_type) * (in_array.length + 1)));
74+
75+
auto offsets = in_array.GetValues<offset_type>(1);
76+
auto shifted_offsets = out_array->GetMutableValues<offset_type>(1);
77+
78+
for (int64_t i = 0; i < in_array.length + 1; ++i) {
79+
shifted_offsets[i] = offsets[i] - offsets[0];
80+
}
81+
values = in_array.child_data[0]->Slice(offsets[0], offsets[in_array.length]);
82+
}
83+
84+
KERNEL_ASSIGN_OR_RAISE(Datum cast_values, ctx,
85+
Cast(values, child_type, options, ctx->exec_context()));
4886

49-
Datum casted_child;
50-
KERNEL_RETURN_IF_ERROR(
51-
ctx, Cast(Datum(input.child_data[0]), child_type, options, ctx->exec_context())
52-
.Value(&casted_child));
53-
DCHECK_EQ(Datum::ARRAY, casted_child.kind());
54-
result->child_data.push_back(casted_child.array());
87+
DCHECK_EQ(Datum::ARRAY, cast_values.kind());
88+
out_array->child_data.push_back(cast_values.array());
5589
}
5690

5791
template <typename Type>

‎cpp/src/arrow/compute/kernels/scalar_cast_numeric.cc‎

Lines changed: 47 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
#include "arrow/array/builder_primitive.h"
2121
#include "arrow/compute/kernels/common.h"
2222
#include "arrow/compute/kernels/scalar_cast_internal.h"
23+
#include "arrow/compute/kernels/util_internal.h"
2324
#include "arrow/util/bit_block_counter.h"
2425
#include "arrow/util/int_util.h"
2526
#include "arrow/util/value_parsing.h"
@@ -361,8 +362,7 @@ struct CastFunctor<O, Decimal128Type, enable_if_t<is_integer_type<O>::value>> {
361362
static void Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
362363
const auto& options = checked_cast<const CastState*>(ctx->state())->options;
363364

364-
const ArrayData& input = *batch[0].array();
365-
const auto& in_type_inst = checked_cast<const Decimal128Type&>(*input.type);
365+
const auto& in_type_inst = checked_cast<const Decimal128Type&>(*batch[0].type());
366366
const auto in_scale = in_type_inst.scale();
367367

368368
if (options.allow_decimal_truncate) {
@@ -395,34 +395,34 @@ struct CastFunctor<O, Decimal128Type, enable_if_t<is_integer_type<O>::value>> {
395395
struct UnsafeUpscaleDecimal {
396396
template <typename... Unused>
397397
Decimal128 Call(KernelContext* ctx, Decimal128 val) const {
398-
return val.IncreaseScaleBy(out_scale_ - in_scale_);
398+
return val.IncreaseScaleBy(by_);
399399
}
400-
401-
int32_t out_scale_, in_scale_;
400+
int32_t by_;
402401
};
403402

404403
struct UnsafeDownscaleDecimal {
405404
template <typename... Unused>
406405
Decimal128 Call(KernelContext* ctx, Decimal128 val) const {
407-
return val.ReduceScaleBy(in_scale_ - out_scale_, false);
406+
return val.ReduceScaleBy(by_, false);
408407
}
409-
410-
int32_t out_scale_, in_scale_;
408+
int32_t by_;
411409
};
412410

413411
struct SafeRescaleDecimal {
414412
template <typename... Unused>
415413
Decimal128 Call(KernelContext* ctx, Decimal128 val) const {
416-
auto result = val.Rescale(in_scale_, out_scale_);
417-
if (ARROW_PREDICT_FALSE(!result.ok())) {
418-
ctx->SetStatus(result.status());
419-
return Decimal128(); // Zero
420-
} else if (ARROW_PREDICT_FALSE(!(*result).FitsInPrecision(out_precision_))) {
421-
ctx->SetStatus(Status::Invalid("Decimal value does not fit in precision"));
422-
return Decimal128(); // Zero
423-
} else {
424-
return *std::move(result);
414+
auto maybe_rescaled = val.Rescale(in_scale_, out_scale_);
415+
if (ARROW_PREDICT_FALSE(!maybe_rescaled.ok())) {
416+
ctx->SetStatus(maybe_rescaled.status());
417+
return {}; // Zero
425418
}
419+
420+
if (ARROW_PREDICT_TRUE(maybe_rescaled->FitsInPrecision(out_precision_))) {
421+
return maybe_rescaled.MoveValueUnsafe();
422+
}
423+
424+
ctx->SetStatus(Status::Invalid("Decimal value does not fit in precision"));
425+
return {}; // Zero
426426
}
427427

428428
int32_t out_scale_, out_precision_, in_scale_;
@@ -432,36 +432,33 @@ template <>
432432
struct CastFunctor<Decimal128Type, Decimal128Type> {
433433
static void Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
434434
const auto& options = checked_cast<const CastState*>(ctx->state())->options;
435-
const ArrayData& input = *batch[0].array();
436-
ArrayData* output = out->mutable_array();
437435

438-
const auto& in_type_inst = checked_cast<const Decimal128Type&>(*input.type);
439-
const auto& out_type_inst = checked_cast<const Decimal128Type&>(*output->type);
440-
const auto in_scale = in_type_inst.scale();
441-
const auto out_scale = out_type_inst.scale();
442-
const auto out_precision = out_type_inst.precision();
436+
const auto& in_type = checked_cast<const Decimal128Type&>(*batch[0].type());
437+
const auto& out_type = checked_cast<const Decimal128Type&>(*out->type());
438+
const auto in_scale = in_type.scale();
439+
const auto out_scale = out_type.scale();
443440

444441
if (options.allow_decimal_truncate) {
445442
if (in_scale < out_scale) {
446443
// Unsafe upscale
447444
applicator::ScalarUnaryNotNullStateful<Decimal128Type, Decimal128Type,
448445
UnsafeUpscaleDecimal>
449-
kernel(UnsafeUpscaleDecimal{out_scale, in_scale});
446+
kernel(UnsafeUpscaleDecimal{out_scale - in_scale});
450447
return kernel.Exec(ctx, batch, out);
451448
} else {
452449
// Unsafe downscale
453450
applicator::ScalarUnaryNotNullStateful<Decimal128Type, Decimal128Type,
454451
UnsafeDownscaleDecimal>
455-
kernel(UnsafeDownscaleDecimal{out_scale, in_scale});
452+
kernel(UnsafeDownscaleDecimal{in_scale - out_scale});
456453
return kernel.Exec(ctx, batch, out);
457454
}
458-
} else {
459-
// Safe rescale
460-
applicator::ScalarUnaryNotNullStateful<Decimal128Type, Decimal128Type,
461-
SafeRescaleDecimal>
462-
kernel(SafeRescaleDecimal{out_scale, out_precision, in_scale});
463-
return kernel.Exec(ctx, batch, out);
464455
}
456+
457+
// Safe rescale
458+
applicator::ScalarUnaryNotNullStateful<Decimal128Type, Decimal128Type,
459+
SafeRescaleDecimal>
460+
kernel(SafeRescaleDecimal{out_scale, out_type.precision(), in_scale});
461+
return kernel.Exec(ctx, batch, out);
465462
}
466463
};
467464

@@ -471,15 +468,16 @@ struct CastFunctor<Decimal128Type, Decimal128Type> {
471468
struct RealToDecimal {
472469
template <typename OutValue, typename RealType>
473470
Decimal128 Call(KernelContext* ctx, RealType val) const {
474-
auto result = Decimal128::FromReal(val, out_precision_, out_scale_);
475-
if (ARROW_PREDICT_FALSE(!result.ok())) {
476-
if (!allow_truncate_) {
477-
ctx->SetStatus(result.status());
478-
}
479-
return Decimal128(); // Zero
480-
} else {
481-
return *std::move(result);
471+
auto maybe_decimal = Decimal128::FromReal(val, out_precision_, out_scale_);
472+
473+
if (ARROW_PREDICT_TRUE(maybe_decimal.ok())) {
474+
return maybe_decimal.MoveValueUnsafe();
482475
}
476+
477+
if (!allow_truncate_) {
478+
ctx->SetStatus(maybe_decimal.status());
479+
}
480+
return {}; // Zero
483481
}
484482

485483
int32_t out_scale_, out_precision_;
@@ -490,10 +488,9 @@ template <typename I>
490488
struct CastFunctor<Decimal128Type, I, enable_if_t<is_floating_type<I>::value>> {
491489
static void Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
492490
const auto& options = checked_cast<const CastState*>(ctx->state())->options;
493-
ArrayData* output = out->mutable_array();
494-
const auto& out_type_inst = checked_cast<const Decimal128Type&>(*output->type);
495-
const auto out_scale = out_type_inst.scale();
496-
const auto out_precision = out_type_inst.precision();
491+
const auto& out_type = checked_cast<const Decimal128Type&>(*out->type());
492+
const auto out_scale = out_type.scale();
493+
const auto out_precision = out_type.precision();
497494

498495
applicator::ScalarUnaryNotNullStateful<Decimal128Type, I, RealToDecimal> kernel(
499496
RealToDecimal{out_scale, out_precision, options.allow_decimal_truncate});
@@ -516,9 +513,8 @@ struct DecimalToReal {
516513
template <typename O>
517514
struct CastFunctor<O, Decimal128Type, enable_if_t<is_floating_type<O>::value>> {
518515
static void Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
519-
const auto& in_type_inst =
520-
checked_cast<const Decimal128Type&>(*batch[0].array()->type);
521-
const auto in_scale = in_type_inst.scale();
516+
const auto& in_type = checked_cast<const Decimal128Type&>(*batch[0].type());
517+
const auto in_scale = in_type.scale();
522518

523519
applicator::ScalarUnaryNotNullStateful<O, Decimal128Type, DecimalToReal> kernel(
524520
DecimalToReal{in_scale});
@@ -564,7 +560,7 @@ std::shared_ptr<CastFunction> GetCastToInteger(std::string name) {
564560
AddCommonNumberCasts<OutType>(out_ty, func.get());
565561

566562
// From decimal to integer
567-
DCHECK_OK(func->AddKernel(Type::DECIMAL, {InputType::Array(Type::DECIMAL)}, out_ty,
563+
DCHECK_OK(func->AddKernel(Type::DECIMAL, {InputType(Type::DECIMAL)}, out_ty,
568564
CastFunctor<OutType, Decimal128Type>::Exec));
569565
return func;
570566
}
@@ -588,7 +584,7 @@ std::shared_ptr<CastFunction> GetCastToFloating(std::string name) {
588584
AddCommonNumberCasts<OutType>(out_ty, func.get());
589585

590586
// From decimal to floating point
591-
DCHECK_OK(func->AddKernel(Type::DECIMAL, {InputType::Array(Type::DECIMAL)}, out_ty,
587+
DCHECK_OK(func->AddKernel(Type::DECIMAL, {InputType(Type::DECIMAL)}, out_ty,
592588
CastFunctor<OutType, Decimal128Type>::Exec));
593589
return func;
594590
}
@@ -608,8 +604,8 @@ std::shared_ptr<CastFunction> GetCastToDecimal128() {
608604
// Cast from other decimal
609605
auto exec = CastFunctor<Decimal128Type, Decimal128Type>::Exec;
610606
// We resolve the output type of this kernel from the CastOptions
611-
DCHECK_OK(func->AddKernel(Type::DECIMAL128, {InputType::Array(Type::DECIMAL128)},
612-
sig_out_ty, exec));
607+
DCHECK_OK(
608+
func->AddKernel(Type::DECIMAL128, {InputType(Type::DECIMAL128)}, sig_out_ty, exec));
613609
return func;
614610
}
615611

0 commit comments

Comments
 (0)