wesm commented on code in PR #13364:
URL: https://github.com/apache/arrow/pull/13364#discussion_r894725732
##########
cpp/src/arrow/array/array_nested.h:
##########
@@ -378,7 +378,7 @@ class ARROW_EXPORT StructArray : public Array {
// Return a shared pointer in case the requestor desires to share ownership
// with this array. The returned array has its offset, length and null
// count adjusted.
- std::shared_ptr<Array> field(int pos) const;
+ const std::shared_ptr<Array>& field(int pos) const;
Review Comment:
This makes a shared_ptr copy that is often not needed
##########
cpp/src/arrow/array/data.cc:
##########
@@ -128,6 +131,143 @@ int64_t ArrayData::GetNullCount() const {
return precomputed;
}
+// ----------------------------------------------------------------------
+// Methods for ArraySpan
+
+void ArraySpan::SetMembers(const ArrayData& data) {
+ this->type = data.type.get();
+ this->length = data.length;
+ this->null_count = data.null_count.load();
+ this->offset = data.offset;
+
+ for (size_t i = 0; i < data.buffers.size(); ++i) {
+ const std::shared_ptr<Buffer>& buffer = data.buffers[i];
+ // It is the invoker-of-kernels's responsibility to ensure that
+ // const buffers are not written to accidentally.
+ if (buffer) {
+ SetBuffer(i, buffer);
+ } else {
+ ClearBuffer(i);
+ }
+ }
+
+ // Makes sure any other buffers are seen as null / non-existent
+ for (size_t i = data.buffers.size(); i < 3; ++i) {
+ ClearBuffer(i);
+ }
+
+ // TODO(wesm): what about extension arrays?
Review Comment:
Nothing needed, will correct comment
##########
cpp/src/arrow/array/data.cc:
##########
@@ -128,6 +131,143 @@ int64_t ArrayData::GetNullCount() const {
return precomputed;
}
+// ----------------------------------------------------------------------
+// Methods for ArraySpan
+
+void ArraySpan::SetMembers(const ArrayData& data) {
+ this->type = data.type.get();
+ this->length = data.length;
+ this->null_count = data.null_count.load();
+ this->offset = data.offset;
+
+ for (size_t i = 0; i < data.buffers.size(); ++i) {
+ const std::shared_ptr<Buffer>& buffer = data.buffers[i];
+ // It is the invoker-of-kernels's responsibility to ensure that
+ // const buffers are not written to accidentally.
+ if (buffer) {
+ SetBuffer(i, buffer);
+ } else {
+ ClearBuffer(i);
+ }
+ }
+
+ // Makes sure any other buffers are seen as null / non-existent
+ for (size_t i = data.buffers.size(); i < 3; ++i) {
+ ClearBuffer(i);
+ }
+
+ // TODO(wesm): what about extension arrays?
+
+ if (this->type->id() == Type::DICTIONARY) {
+ this->child_data.resize(1);
+ this->child_data[0].SetMembers(*data.dictionary);
+ } else {
+ this->child_data.resize(data.child_data.size());
+ for (size_t child_index = 0; child_index < data.child_data.size();
++child_index) {
+ this->child_data[child_index].SetMembers(*data.child_data[child_index]);
+ }
+ }
+}
+
+void ArraySpan::FillFromScalar(const Scalar& value) {
+ static const uint8_t kValidByte = 0x01;
+ static const uint8_t kNullByte = 0x00;
+
+ this->type = value.type.get();
+ this->length = 1;
+
+ // Populate null count and validity bitmap
+ this->null_count = value.is_valid ? 0 : 1;
+ this->buffers[0].data = const_cast<uint8_t*>(value.is_valid ? &kValidByte :
&kNullByte);
+ this->buffers[0].size = 1;
+
+ if (is_primitive(value.type->id())) {
+ const auto& scalar =
+ internal::checked_cast<const internal::PrimitiveScalarBase&>(value);
+ const uint8_t* scalar_data = reinterpret_cast<const
uint8_t*>(scalar.data());
+ this->buffers[1].data = const_cast<uint8_t*>(scalar_data);
+ this->buffers[1].size = scalar.type->byte_width();
+ } else {
+ // TODO(wesm): implement for other types
+ DCHECK(false) << "need to implement for other types";
Review Comment:
Open Jira for follow up work
##########
cpp/src/arrow/compute/kernels/scalar_if_else.cc:
##########
@@ -77,170 +64,239 @@ Status CheckIdenticalTypes(const Datum* begin, size_t
count) {
return Status::OK();
}
+constexpr uint64_t kAllNull = 0;
+constexpr uint64_t kAllValid = ~kAllNull;
+
+util::optional<uint64_t> GetConstantValidityWord(const ExecValue& data) {
+ if (data.is_scalar()) {
+ return data.scalar->is_valid ? kAllValid : kAllNull;
+ }
+
+ if (data.array.null_count == data.array.length) return kAllNull;
+ if (!data.array.MayHaveNulls()) return kAllValid;
+
+ // no constant validity word available
+ return {};
+}
+
// if the condition is null then output is null otherwise we take validity
from the
// selected argument
// ie. cond.valid & (cond.data & left.valid | ~cond.data & right.valid)
-template <typename AllocateNullBitmap>
-Status PromoteNullsVisitor(KernelContext* ctx, const Datum& cond_d, const
Datum& left_d,
- const Datum& right_d, ArrayData* output) {
- auto cond_const = GetConstantValidityWord(cond_d);
- auto left_const = GetConstantValidityWord(left_d);
- auto right_const = GetConstantValidityWord(right_d);
+struct IfElseNullPromoter {
+ KernelContext* ctx;
+ const ArraySpan& cond;
+ const ExecValue& left_d;
+ const ExecValue& right_d;
+ ExecResult* output;
enum { COND_CONST = 1, LEFT_CONST = 2, RIGHT_CONST = 4 };
- auto flag = COND_CONST * cond_const.has_value() | LEFT_CONST *
left_const.has_value() |
- RIGHT_CONST * right_const.has_value();
-
- const ArrayData& cond = *cond_d.array();
- // cond.data will always be available
- Bitmap cond_data{cond.buffers[1], cond.offset, cond.length};
- Bitmap cond_valid{cond.buffers[0], cond.offset, cond.length};
- Bitmap left_valid = GetBitmap(left_d, 0);
- Bitmap right_valid = GetBitmap(right_d, 0);
-
- // cond.valid & (cond.data & left.valid | ~cond.data & right.valid)
- // In the following cases, we dont need to allocate out_valid bitmap
-
- // if cond & left & right all ones, then output is all valid.
- // if output validity buffer is already allocated (NullHandling::
- // COMPUTED_PREALLOCATE) -> set all bits
- // else, return nullptr
- if (cond_const == kAllValid && left_const == kAllValid && right_const ==
kAllValid) {
- if (AllocateNullBitmap::value) { // NullHandling::COMPUTED_NO_PREALLOCATE
- output->buffers[0] = nullptr;
- } else { // NullHandling::COMPUTED_PREALLOCATE
- bit_util::SetBitmap(output->buffers[0]->mutable_data(), output->offset,
- output->length);
+ int64_t constant_validity_flag;
+ util::optional<uint64_t> cond_const, left_const, right_const;
+ Bitmap cond_data, cond_valid, left_valid, right_valid;
+
+ IfElseNullPromoter(KernelContext* ctx, const ExecValue& cond_d, const
ExecValue& left_d,
+ const ExecValue& right_d, ExecResult* output)
+ : ctx(ctx), cond(cond_d.array), left_d(left_d), right_d(right_d),
output(output) {
+ cond_const = GetConstantValidityWord(cond_d);
+ left_const = GetConstantValidityWord(left_d);
+ right_const = GetConstantValidityWord(right_d);
+
+ // Encodes whether each of the arguments is respectively all-null or
+ // all-valid
+ constant_validity_flag =
+ (COND_CONST * cond_const.has_value() | LEFT_CONST *
left_const.has_value() |
+ RIGHT_CONST * right_const.has_value());
+
+ // cond.data will always be available / is always an array
+ cond_data = Bitmap{cond.buffers[1].data, cond.offset, cond.length};
+ cond_valid = Bitmap{cond.buffers[0].data, cond.offset, cond.length};
+ left_valid = GetBitmap(left_d, 0);
+ right_valid = GetBitmap(right_d, 0);
+ }
+
+ Status ExecIntoArrayData(bool need_to_allocate) {
+ ArrayData* out_arr = output->array_data().get();
+
+ // cond.valid & (cond.data & left.valid | ~cond.data & right.valid)
+ // In the following cases, we dont need to allocate out_valid bitmap
+
+ // if cond & left & right all ones, then output is all valid.
+ // if output validity buffer is already allocated (NullHandling::
+ // COMPUTED_PREALLOCATE) -> set all bits
+ // else, return nullptr
+ if (cond_const == kAllValid && left_const == kAllValid && right_const ==
kAllValid) {
+ if (need_to_allocate) { // NullHandling::COMPUTED_NO_PREALLOCATE
+ out_arr->buffers[0] = nullptr;
+ } else { // NullHandling::COMPUTED_PREALLOCATE
+ bit_util::SetBitmap(out_arr->buffers[0]->mutable_data(),
out_arr->offset,
+ out_arr->length);
+ }
+ } else if (left_const == kAllValid && right_const == kAllValid) {
+ // if both left and right are valid, no need to calculate out_valid
bitmap. Copy
+ // cond validity buffer
+ if (need_to_allocate) { // NullHandling::COMPUTED_NO_PREALLOCATE
+ // if there's an offset, copy bitmap (cannot slice a bitmap)
+ if (cond.offset) {
+ ARROW_ASSIGN_OR_RAISE(
+ out_arr->buffers[0],
+ arrow::internal::CopyBitmap(ctx->memory_pool(),
cond.buffers[0].data,
+ cond.offset, cond.length));
+ } else {
+ // just copy assign cond validity buffer
+ out_arr->buffers[0] = cond.GetBuffer(0);
+ }
+ } else { // NullHandling::COMPUTED_PREALLOCATE
+ arrow::internal::CopyBitmap(cond.buffers[0].data, cond.offset,
cond.length,
+ out_arr->buffers[0]->mutable_data(),
out_arr->offset);
+ }
+ } else {
+ if (need_to_allocate) {
+ // following cases requires a separate out_valid buffer.
COMPUTED_NO_PREALLOCATE
+ // would not have allocated buffers for it.
+ ARROW_ASSIGN_OR_RAISE(out_arr->buffers[0],
ctx->AllocateBitmap(cond.length));
+ }
+ WriteOutput(
+ Bitmap{out_arr->buffers[0]->mutable_data(), out_arr->offset,
out_arr->length});
}
return Status::OK();
}
- if (left_const == kAllValid && right_const == kAllValid) {
- // if both left and right are valid, no need to calculate out_valid
bitmap. Copy
- // cond validity buffer
- if (AllocateNullBitmap::value) { // NullHandling::COMPUTED_NO_PREALLOCATE
- // if there's an offset, copy bitmap (cannot slice a bitmap)
- if (cond.offset) {
- ARROW_ASSIGN_OR_RAISE(
- output->buffers[0],
- arrow::internal::CopyBitmap(ctx->memory_pool(),
cond.buffers[0]->data(),
- cond.offset, cond.length));
- } else { // just copy assign cond validity buffer
- output->buffers[0] = cond.buffers[0];
- }
- } else { // NullHandling::COMPUTED_PREALLOCATE
- arrow::internal::CopyBitmap(cond.buffers[0]->data(), cond.offset,
cond.length,
- output->buffers[0]->mutable_data(),
output->offset);
+ Status ExecIntoArraySpan() {
+ ArraySpan* out_span = output->array_span();
+
+ // cond.valid & (cond.data & left.valid | ~cond.data & right.valid)
+ // In the following cases, we dont need to allocate out_valid bitmap
+
+ // if cond & left & right all ones, then output is all valid, so set all
+ // bits
+ if (cond_const == kAllValid && left_const == kAllValid && right_const ==
kAllValid) {
+ bit_util::SetBitmap(out_span->buffers[0].data, out_span->offset,
out_span->length);
+ } else if (left_const == kAllValid && right_const == kAllValid) {
+ // if both left and right are valid, no need to calculate out_valid
+ // bitmap. Copy cond validity buffer
+ arrow::internal::CopyBitmap(cond.buffers[0].data, cond.offset,
cond.length,
+ out_span->buffers[0].data, out_span->offset);
+ } else {
+ WriteOutput(Bitmap{out_span->buffers[0].data, out_span->offset,
out_span->length});
}
return Status::OK();
}
- // lambda function that will be used inside the visitor
- auto apply = [&](uint64_t c_valid, uint64_t c_data, uint64_t l_valid,
- uint64_t r_valid) {
- return c_valid & ((c_data & l_valid) | (~c_data & r_valid));
- };
-
- if (AllocateNullBitmap::value) {
- // following cases requires a separate out_valid buffer.
COMPUTED_NO_PREALLOCATE
- // would not have allocated buffers for it.
- ARROW_ASSIGN_OR_RAISE(output->buffers[0],
ctx->AllocateBitmap(cond.length));
+ Status Exec(bool need_to_allocate) {
+ if (output->is_array_data()) {
+ return ExecIntoArrayData(need_to_allocate);
+ } else {
+ if (need_to_allocate) {
+ // TODO: turn this into a DCHECK, but have this strong error to be
+ // helpful for now
+ return Status::Invalid(
+ "Conditional kernel writing into array span must "
+ "preallocate validity bitmap");
+ }
Review Comment:
Turn into dcheck
##########
cpp/src/arrow/compute/kernels/scalar_random.cc:
##########
@@ -66,23 +66,23 @@ Status ExecRandom(KernelContext* ctx, const ExecBatch&
batch, Datum* out) {
return Status::Invalid("Negative number of elements");
}
- auto out_data = ArrayData::Make(float64(), options.length, 0);
- out_data->buffers.resize(2, nullptr);
-
- ARROW_ASSIGN_OR_RAISE(out_data->buffers[1],
- ctx->Allocate(options.length * sizeof(double)));
- double* out_buffer = out_data->template GetMutableValues<double>(1);
-
if (options.initializer == RandomOptions::Seed) {
gen.seed(options.seed);
} else {
std::lock_guard<std::mutex> seed_gen_lock(seed_gen_mutex);
gen.seed(seed_gen());
}
+ // TODO(wesm): refactor to use batch length instead of passing length in
+ // options
+ auto out_data = ArrayData::Make(float64(), options.length, 0);
Review Comment:
Note: I found this aspect of `random()` implementation to be a huge eyesore
-- will create a Jira to fix this to write into preallocated memory and not
pass the length through the options
##########
cpp/src/arrow/chunked_array.h:
##########
@@ -111,7 +111,7 @@ class ARROW_EXPORT ChunkedArray {
int num_chunks() const { return static_cast<int>(chunks_.size()); }
/// \return chunk a particular chunk from the chunked array
- std::shared_ptr<Array> chunk(int i) const { return chunks_[i]; }
+ const std::shared_ptr<Array>& chunk(int i) const { return chunks_[i]; }
Review Comment:
another place where shared_ptr being copied needlessly
##########
cpp/src/arrow/type.h:
##########
@@ -174,6 +175,25 @@ class ARROW_EXPORT DataType : public
detail::Fingerprintable {
/// \brief Return the type category of the storage type
virtual Type::type storage_id() const { return id_; }
+ /// \brief Returns the type's fixed byte width, if any. Returns -1
+ /// for non-fixed-width types, and should only be used for
+ /// subclasses of FixedWidthType
+ virtual int32_t byte_width() const {
Review Comment:
I added this because getting the byte width from a fixed width type is quite
tedious at the moment (there's many places where we do the cast-- I think
having the bit_width()/byte_width() methods available on non-fixed-width types
(returning -1) is pretty harmless, but if others don't like it we can remove
these and add a helper function or something)
##########
cpp/src/arrow/util/bit_block_counter.h:
##########
@@ -490,26 +493,34 @@ static void VisitBitBlocksVoid(const
std::shared_ptr<Buffer>& bitmap_buf, int64_
}
template <typename VisitNotNull, typename VisitNull>
-static Status VisitTwoBitBlocks(const std::shared_ptr<Buffer>& left_bitmap_buf,
- int64_t left_offset,
- const std::shared_ptr<Buffer>&
right_bitmap_buf,
- int64_t right_offset, int64_t length,
- VisitNotNull&& visit_not_null, VisitNull&&
visit_null) {
- if (left_bitmap_buf == NULLPTR || right_bitmap_buf == NULLPTR) {
+static void VisitBitBlocksVoid(const std::shared_ptr<Buffer>& bitmap_buf,
int64_t offset,
Review Comment:
It's sorta tedious to have two versions of each of these `VisitBitBlocks`
functions, one for `shared_ptr<Buffer>` and one for raw pointers. we should
probably clean this up later
##########
cpp/src/arrow/compute/exec.cc:
##########
@@ -229,18 +229,45 @@ Status CheckAllValues(const std::vector<Datum>& values) {
return Status::OK();
}
-ExecBatchIterator::ExecBatchIterator(std::vector<Datum> args, int64_t length,
+Status GetBatchLength(const std::vector<Datum>& values, int64_t* out_length) {
+ for (const auto& arg : values) {
+ if (!(arg.is_arraylike() || arg.is_scalar())) {
+ return Status::Invalid(
+ "Batch iteration only works with Scalar, Array, and "
+ "ChunkedArray arguments");
+ }
+ }
+
+ // If the arguments are all scalars, then the length is 1
+ int64_t length = 1;
+
+ bool length_set = false;
+ for (auto& arg : values) {
+ if (arg.is_scalar()) {
+ continue;
+ }
+ if (!length_set) {
+ length = arg.length();
+ length_set = true;
+ } else {
+ if (arg.length() != length) {
+ return Status::Invalid("Array arguments must all be the same length");
+ }
+ }
+ }
+ *out_length = length;
+ return Status::OK();
+}
+
+ExecBatchIterator::ExecBatchIterator(const std::vector<Datum>& args, int64_t
length,
int64_t max_chunksize)
- : args_(std::move(args)),
- position_(0),
- length_(length),
- max_chunksize_(max_chunksize) {
+ : args_(args), position_(0), length_(length),
max_chunksize_(max_chunksize) {
chunk_indexes_.resize(args_.size(), 0);
chunk_positions_.resize(args_.size(), 0);
}
Result<std::unique_ptr<ExecBatchIterator>> ExecBatchIterator::Make(
- std::vector<Datum> args, int64_t max_chunksize) {
+ const std::vector<Datum>& args, int64_t max_chunksize) {
Review Comment:
Aside: I noticed a couple places where very expensive `std::vector<Datum>`
copies were happening
##########
cpp/src/arrow/compute/exec_internal.h:
##########
@@ -138,7 +188,10 @@ class ARROW_EXPORT KernelExecutor {
/// \param[in] batch the data batch
/// \param[in] out the output ArrayData, must not be null
ARROW_EXPORT
-Status PropagateNulls(KernelContext* ctx, const ExecBatch& batch, ArrayData*
out);
+Status PropagateNulls(KernelContext* ctx, const ExecSpan& batch, ArrayData*
out);
+
+ARROW_EXPORT
+void PropagateNullsSpans(const ExecSpan& batch, ArraySpan* out);
Review Comment:
Only needed for testing
##########
cpp/src/arrow/compute/kernels/scalar_cast_numeric.cc:
##########
@@ -249,10 +259,11 @@ Status CheckForIntegerToFloatingTruncation(const Datum&
input, Type::type out_ty
return Status::OK();
}
-Status CastIntegerToFloating(KernelContext* ctx, const ExecBatch& batch,
Datum* out) {
+Status CastIntegerToFloating(KernelContext* ctx, const ExecSpan& batch,
ExecResult* out) {
const auto& options = checked_cast<const CastState*>(ctx->state())->options;
Type::type out_type = out->type()->id();
if (!options.allow_float_truncate) {
+ /// XXX: refactor to not use Datum
Review Comment:
comment outdated
##########
cpp/src/arrow/compute/kernels/scalar_temporal_unary.cc:
##########
@@ -1311,10 +1316,18 @@ struct Strptime {
}
out_writer.Next();
};
- VisitArrayDataInline<InType>(in, visit_value, visit_null);
+ VisitArraySpanInline<InType>(in, visit_value, visit_null);
out_writer.Finish();
- out->null_count = null_count;
+ out_span->null_count = null_count;
} else {
+ if (in.buffers[0].data != nullptr) {
+ ::arrow::internal::CopyBitmap(in.buffers[0].data, in.offset, in.length,
+ out_span->buffers[0].data,
out_span->offset);
+ } else {
+ // Input is all non-null
+ bit_util::SetBitmap(out_span->buffers[0].data, out_span->offset,
+ out_span->length);
+ }
Review Comment:
I think Strptime should probably be split into two kernels, one for
errors-as-null and the other not, that way we don't have to do this manual null
propagation business for the non-errors-as-null case
##########
cpp/src/arrow/compute/kernels/scalar_string_ascii.cc:
##########
@@ -156,22 +153,27 @@ struct FixedSizeBinaryTransformExecWithState
const std::vector<ValueDescr>& descrs) {
DCHECK_EQ(1, descrs.size());
const auto& options = State::Get(ctx);
- const int32_t input_width =
- checked_cast<const FixedSizeBinaryType&>(*descrs[0].type).byte_width();
+ const int32_t input_width = descrs[0].type->byte_width();
const int32_t output_width = StringTransform::FixedOutputSize(options,
input_width);
return ValueDescr(fixed_size_binary(output_width), descrs[0].shape);
}
};
+template <typename T>
Review Comment:
nit: `offset_type`
##########
cpp/src/arrow/compute/kernels/scalar_validity.cc:
##########
@@ -337,16 +333,15 @@ const FunctionDoc is_nan_doc("Return true if NaN",
void RegisterScalarValidity(FunctionRegistry* registry) {
static auto kNullOptions = NullOptions::Defaults();
MakeFunction("is_valid", is_valid_doc, {ValueDescr::ANY}, boolean(),
IsValidExec,
- registry, MemAllocation::NO_PREALLOCATE,
NullHandling::OUTPUT_NOT_NULL,
Review Comment:
Note: is_valid will no longer do zero copy optimizations. I think these
buffer-reuse (whether slicing the input buffer or passing it on to the output)
optimizations have relatively little value on the whole because they prevent
parallel execution by input-splitting (the ideal scenario for parallel
execution is to have each parallel task writing into a slice of a larger memory
preallocation)
##########
cpp/src/arrow/compute/kernels/codegen_internal.h:
##########
@@ -397,14 +400,198 @@ struct BoxScalar<Decimal256Type> {
static void Box(T val, Scalar* out) { checked_cast<ScalarType*>(out)->value
= val; }
};
-// A VisitArrayDataInline variant that calls its visitor function with logical
+// ----------------------------------------------------------------------
+// Like VisitArrayDataInline, but for ArraySpans
Review Comment:
TODO: move this to visit_data_inline.h
Question for the reviewers: do you think it's worth maintaining the
VisitArrayDataInline code path if we can migrate everything to use ArraySpans?
There's a lot of other code in the codebase that could probably be simplified
and also probably even be faster to compile
##########
cpp/src/arrow/compute/kernels/util_internal.cc:
##########
@@ -57,22 +57,67 @@ PrimitiveArg GetPrimitiveArg(const ArrayData& arr) {
return arg;
}
-ArrayKernelExec TrivialScalarUnaryAsArraysExec(ArrayKernelExec exec,
- NullHandling::type
null_handling) {
- return [=](KernelContext* ctx, const ExecBatch& batch, Datum* out) -> Status
{
- if (out->is_array()) {
- return exec(ctx, batch, out);
+// TODO(wesm): ARROW-16577: this will be unneeded later
Review Comment:
This function is a huge mess but after ARROW-16577 it will not be needed at
all, so during that patch (eliminating the `ValueDescr::SCALAR` output path and
frankly the whole ValueDescr struct altogether) is the right time to clean this
up and move the scalar-to-array-of-length-1 promotion into the kernel execution
layer
##########
cpp/src/arrow/compute/kernels/scalar_validity_test.cc:
##########
@@ -67,12 +67,6 @@ TEST_F(TestBooleanValidityKernels, IsValidIsNullNullType) {
ArrayFromJSON(boolean(), "[null, null, null, null, null]"));
}
-TEST_F(TestBooleanValidityKernels, ArrayIsValidBufferPassthruOptimization) {
- Datum arg = ArrayFromJSON(boolean(), "[null, 1, 0, null]");
- ASSERT_OK_AND_ASSIGN(auto validity, arrow::compute::IsValid(arg));
- ASSERT_EQ(validity.array()->buffers[1], arg.array()->buffers[0]);
-}
Review Comment:
Deleted test case because this optimization is no longer performed. Copying
a bitmap is cheap
##########
cpp/src/arrow/compute/kernels/scalar_string_ascii.cc:
##########
@@ -2542,9 +2540,11 @@ void AddAsciiStringReplaceSlice(FunctionRegistry*
registry) {
}
using TransformExec =
FixedSizeBinaryTransformExecWithState<BinaryReplaceSliceTransform>;
- DCHECK_OK(func->AddKernel({InputType(Type::FIXED_SIZE_BINARY)},
- OutputType(TransformExec::OutputType),
TransformExec::Exec,
- ReplaceStringSliceTransformBase::State::Init));
+ ScalarKernel fsb_kernel({InputType(Type::FIXED_SIZE_BINARY)},
+ OutputType(TransformExec::OutputType),
TransformExec::Exec,
+ ReplaceStringSliceTransformBase::State::Init);
+ fsb_kernel.mem_allocation = MemAllocation::NO_PREALLOCATE;
Review Comment:
The prior kernel implementation here turned out to be preallocating only to
overwrite the preallocated buffer...
##########
cpp/src/arrow/scalar.h:
##########
@@ -49,7 +49,8 @@ class Array;
/// Scalars are useful for passing single value inputs to compute functions,
/// or for representing individual array elements (with a non-trivial
/// wrapping cost, though).
-struct ARROW_EXPORT Scalar : public util::EqualityComparable<Scalar> {
+struct ARROW_EXPORT Scalar : public std::enable_shared_from_this<Scalar>,
Review Comment:
IMHO this isn't wonderful. I would rather introduce `virtual
std::shared_ptr<DataType> Scalar::Copy()`
##########
cpp/src/arrow/compute/kernels/scalar_nested.cc:
##########
@@ -472,78 +526,95 @@ struct MapLookupFunctor {
return Status::OK();
}
- static Status ExecMapArray(KernelContext* ctx, const ExecBatch& batch,
Datum* out) {
+ static Status ExecMapArray(KernelContext* ctx, const ExecSpan& batch,
ExecResult* out) {
const auto& options = OptionsWrapper<MapLookupOptions>::Get(ctx);
- const auto& query_key = options.query_key;
- const auto& occurrence = options.occurrence;
- const MapArray map_array(batch[0].array());
+ const UnboxedKey query_key =
UnboxScalar<KeyType>::Unbox(*options.query_key);
+
+ const ArraySpan& map = batch[0].array;
+ const int32_t* offsets = map.GetValues<int32_t>(1);
+
+ // We create a copy of the keys array because we will adjust the
+ // offset and length for the map probes below
+ ArraySpan map_keys = map.child_data[0].child_data[0];
+ const ArraySpan& map_items = map.child_data[0].child_data[1];
+
+ std::shared_ptr<DataType> item_type =
+ checked_cast<const MapType*>(map.type)->item_type();
std::unique_ptr<ArrayBuilder> builder;
- if (occurrence == MapLookupOptions::Occurrence::ALL) {
- RETURN_NOT_OK(MakeBuilder(ctx->memory_pool(),
- list(map_array.map_type()->item_type()),
&builder));
+ if (options.occurrence == MapLookupOptions::Occurrence::ALL) {
+ RETURN_NOT_OK(MakeBuilder(ctx->memory_pool(), list(item_type),
&builder));
auto list_builder = checked_cast<ListBuilder*>(builder.get());
auto value_builder = list_builder->value_builder();
- for (int64_t map_array_idx = 0; map_array_idx < map_array.length();
- ++map_array_idx) {
- if (!map_array.IsValid(map_array_idx)) {
+ for (int64_t map_index = 0; map_index < map.length; ++map_index) {
+ if (!map.IsValid(map_index)) {
RETURN_NOT_OK(list_builder->AppendNull());
continue;
}
- auto map = map_array.value_slice(map_array_idx);
- auto keys = checked_cast<const StructArray&>(*map).field(0);
- auto items = checked_cast<const StructArray&>(*map).field(1);
Review Comment:
Note: avoiding all this value slicing and array boxing is a huge efficiency
bonus
##########
cpp/src/arrow/type.h:
##########
@@ -126,7 +126,8 @@ struct ARROW_EXPORT DataTypeLayout {
///
/// Simple datatypes may be entirely described by their Type::type id, but
/// complex datatypes are usually parametric.
-class ARROW_EXPORT DataType : public detail::Fingerprintable {
+class ARROW_EXPORT DataType : public std::enable_shared_from_this<DataType>,
Review Comment:
Also not wonderful. I would rather add `DataType::Copy`, what do others
think?
##########
cpp/src/arrow/util/bitmap.h:
##########
@@ -453,7 +459,8 @@ class ARROW_EXPORT Bitmap : public
util::ToStringOstreamable<Bitmap>,
return bitmaps[0].length();
}
- std::shared_ptr<Buffer> buffer_;
Review Comment:
In general, we should watch out for classes that take copies of
`shared_ptr<T>` without a good reason to; this was a good example
##########
cpp/src/arrow/compute/exec.h:
##########
@@ -252,6 +254,192 @@ struct ARROW_EXPORT ExecBatch {
inline bool operator==(const ExecBatch& l, const ExecBatch& r) { return
l.Equals(r); }
inline bool operator!=(const ExecBatch& l, const ExecBatch& r) { return
!l.Equals(r); }
+struct ExecValue {
+ enum Kind { ARRAY, SCALAR };
+ Kind kind = ARRAY;
+ ArraySpan array;
+ const Scalar* scalar;
+
+ ExecValue(Scalar* scalar) // NOLINT implicit conversion
+ : kind(SCALAR), scalar(scalar) {}
+
+ ExecValue(ArraySpan array) // NOLINT implicit conversion
+ : kind(ARRAY), array(std::move(array)) {}
+
+ ExecValue(const ArrayData& array) // NOLINT implicit conversion
+ : kind(ARRAY) {
+ this->array.SetMembers(array);
+ }
+
+ ExecValue() = default;
+ ExecValue(const ExecValue& other) = default;
+ ExecValue& operator=(const ExecValue& other) = default;
+ ExecValue(ExecValue&& other) = default;
+ ExecValue& operator=(ExecValue&& other) = default;
+
+ int64_t length() const { return this->is_array() ? this->array.length : 1; }
+
+ bool is_array() const { return this->kind == ARRAY; }
+ bool is_scalar() const { return this->kind == SCALAR; }
+
+ void SetArray(const ArrayData& array) {
+ this->kind = ARRAY;
+ this->array.SetMembers(array);
+ }
+
+ void SetScalar(const Scalar* scalar) {
+ this->kind = SCALAR;
+ this->scalar = scalar;
+ }
+
+ template <typename ExactType>
+ const ExactType& scalar_as() const {
+ return ::arrow::internal::checked_cast<const ExactType&>(*this->scalar);
+ }
+
+ /// XXX: here only temporarily until type resolution can be cleaned
+ /// up to not use ValueDescr
+ ValueDescr descr() const {
+ ValueDescr::Shape shape = this->is_array() ? ValueDescr::ARRAY :
ValueDescr::SCALAR;
+ return ValueDescr(const_cast<DataType*>(this->type())->shared_from_this(),
shape);
+ }
+
+ /// XXX: here temporarily for compatibility with datum, see
+ /// e.g. MakeStructExec in scalar_nested.cc
+ int64_t null_count() const {
+ if (this->is_array()) {
+ return this->array.GetNullCount();
+ } else {
+ return this->scalar->is_valid ? 0 : 1;
+ }
+ }
+
+ const DataType* type() const {
+ if (this->kind == ARRAY) {
+ return array.type;
+ } else {
+ return scalar->type.get();
+ }
+ }
+};
+
+struct ARROW_EXPORT ExecResult {
+ // The default value of the variant is ArraySpan
+ // TODO(wesm): remove Scalar output modality in ARROW-16577
+ util::Variant<ArraySpan, std::shared_ptr<ArrayData>,
std::shared_ptr<Scalar>> value;
+
+ int64_t length() const {
+ if (this->is_array_span()) {
+ return this->array_span()->length;
+ } else if (this->is_array_data()) {
+ return this->array_data()->length;
+ } else {
+ // Should not reach here
+ return 1;
+ }
+ }
+
+ const DataType* type() const {
+ switch (this->value.index()) {
+ case 0:
+ return this->array_span()->type;
+ case 1:
+ return this->array_data()->type.get();
+ default:
+ // scalar
+ return this->scalar()->type.get();
+ };
+ }
+
+ ArraySpan* array_span() const {
+ return const_cast<ArraySpan*>(&util::get<ArraySpan>(this->value));
+ }
+ bool is_array_span() const { return this->value.index() == 0; }
+
+ const std::shared_ptr<ArrayData>& array_data() const {
+ return util::get<std::shared_ptr<ArrayData>>(this->value);
+ }
+
+ bool is_array_data() const { return this->value.index() == 1; }
+
+ const std::shared_ptr<Scalar>& scalar() const {
+ return util::get<std::shared_ptr<Scalar>>(this->value);
+ }
+
+ bool is_scalar() const { return this->value.index() == 2; }
+};
+
+/// \brief A "lightweight" column batch object which contains no
+/// std::shared_ptr objects and does not have any memory ownership
+/// semantics. Can represent a view onto an "owning" ExecBatch.
+struct ARROW_EXPORT ExecSpan {
+ ExecSpan() = default;
+ ExecSpan(const ExecSpan& other) = default;
+ ExecSpan& operator=(const ExecSpan& other) = default;
+ ExecSpan(ExecSpan&& other) = default;
+ ExecSpan& operator=(ExecSpan&& other) = default;
+
+ explicit ExecSpan(std::vector<ExecValue> values, int64_t length)
+ : length(length), values(std::move(values)) {}
+
+ explicit ExecSpan(const ExecBatch& batch) {
+ this->length = batch.length;
+ this->values.resize(batch.values.size());
+ for (size_t i = 0; i < batch.values.size(); ++i) {
+ const Datum& in_value = batch[i];
+ ExecValue* out_value = &this->values[i];
+ if (in_value.is_array()) {
+ out_value->SetArray(*in_value.array());
+ } else {
+ out_value->SetScalar(in_value.scalar().get());
+ }
+ }
+ }
+
+ bool is_all_scalar() const {
+ return std::all_of(this->values.begin(), this->values.end(),
+ [](const ExecValue& v) { return v.is_scalar(); });
+ }
+
+ /// \brief Return the value at the i-th index
+ template <typename index_type>
+ inline const ExecValue& operator[](index_type i) const {
+ return values[i];
+ }
+
+ void AddOffset(int64_t offset) {
+ for (ExecValue& value : values) {
+ if (value.kind == ExecValue::ARRAY) {
+ value.array.AddOffset(offset);
+ }
+ }
+ }
+
+ void SetOffset(int64_t offset) {
+ for (ExecValue& value : values) {
+ if (value.kind == ExecValue::ARRAY) {
+ value.array.SetOffset(offset);
+ }
+ }
+ }
+
+ /// \brief A convenience for the number of values / arguments.
+ int num_values() const { return static_cast<int>(values.size()); }
+
+ // XXX: eliminate the need for ValueDescr; copied temporarily from
+ // ExecBatch
+ std::vector<ValueDescr> GetDescriptors() const {
+ std::vector<ValueDescr> result;
+ for (const auto& value : this->values) {
+ result.emplace_back(value.descr());
+ }
+ return result;
+ }
+
+ int64_t length;
+ std::vector<ExecValue> values;
Review Comment:
I agree -- potentially we could use SmallVector when there are a small
number of fields and then switch to `std::vector` only for "wide" spans, either
way it's something we can optimize and at minimum limit the exposure of this
`std::vector` to users of `ExecSpan`
##########
cpp/src/arrow/scalar.h:
##########
@@ -129,6 +135,9 @@ namespace internal {
struct ARROW_EXPORT PrimitiveScalarBase : public Scalar {
using Scalar::Scalar;
+ /// \brief Get an immutable pointer to the value of this scalar. May be null.
+ virtual const void* data() const = 0;
Review Comment:
will look
##########
cpp/src/arrow/util/int_util.h:
##########
@@ -99,20 +98,24 @@ Status TransposeInts(const DataType& src_type, const
DataType& dest_type,
/// indices must be nonnegative and strictly less than the passed upper
/// limit (which is usually the length of an array that is being indexed-into).
ARROW_EXPORT
-Status CheckIndexBounds(const ArrayData& indices, uint64_t upper_limit);
+Status CheckIndexBounds(const ArraySpan& values, uint64_t upper_limit);
/// \brief Boundscheck integer values to determine if they are all between the
/// passed upper and lower limits (inclusive). Upper and lower bounds must be
/// the same type as the data and are not currently casted.
ARROW_EXPORT
-Status CheckIntegersInRange(const Datum& datum, const Scalar& bound_lower,
+Status CheckIntegersInRange(const ArraySpan& values, const Scalar& bound_lower,
Review Comment:
There's plenty of other places where we should nix the use of Datum because
of its bloat. We might consider promoting `ExecValue` to the Datum level (like
`ValueSpan` or something)
##########
cpp/src/arrow/util/bitmap.h:
##########
@@ -53,27 +53,36 @@ class ARROW_EXPORT Bitmap : public
util::ToStringOstreamable<Bitmap>,
Bitmap() = default;
- Bitmap(std::shared_ptr<Buffer> buffer, int64_t offset, int64_t length)
- : buffer_(std::move(buffer)), offset_(offset), length_(length) {}
+ Bitmap(const std::shared_ptr<Buffer>& buffer, int64_t offset, int64_t length)
Review Comment:
I agree that a rename would be good. also agree with consistent naming for
view/span/reference types
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
