diff --git a/cpp/src/arrow/array/builder_base.cc b/cpp/src/arrow/array/builder_base.cc index eea394977f4..c70ecec80cf 100644 --- a/cpp/src/arrow/array/builder_base.cc +++ b/cpp/src/arrow/array/builder_base.cc @@ -73,271 +73,23 @@ Status ArrayBuilder::Resize(int64_t capacity) { return null_bitmap_builder_.Resize(capacity); } -namespace { - -template -struct AppendScalarImpl { - template - Status HandleFixedWidth(const T&) { - auto builder = checked_cast::BuilderType*>(builder_); - RETURN_NOT_OK(builder->Reserve(n_repeats_ * (scalars_end_ - scalars_begin_))); - - for (int64_t i = 0; i < n_repeats_; i++) { - for (auto it = scalars_begin_; it != scalars_end_; ++it) { - const auto& scalar = checked_cast::ScalarType&>(*it); - if (scalar.is_valid) { - builder->UnsafeAppend(scalar.value); - } else { - builder->UnsafeAppendNull(); - } - } - } - return Status::OK(); - } - - template - enable_if_t::value, Status> Visit(const T& t) { - return HandleFixedWidth(t); - } - - Status Visit(const FixedSizeBinaryType& t) { return HandleFixedWidth(t); } - Status Visit(const Decimal32Type& t) { return HandleFixedWidth(t); } - Status Visit(const Decimal64Type& t) { return HandleFixedWidth(t); } - Status Visit(const Decimal128Type& t) { return HandleFixedWidth(t); } - Status Visit(const Decimal256Type& t) { return HandleFixedWidth(t); } - - template - enable_if_has_string_view Visit(const T&) { - int64_t data_size = 0; - for (auto it = scalars_begin_; it != scalars_end_; ++it) { - const auto& scalar = checked_cast::ScalarType&>(*it); - if (scalar.is_valid) { - data_size += scalar.value->size(); - } - } - - auto builder = checked_cast::BuilderType*>(builder_); - RETURN_NOT_OK(builder->Reserve(n_repeats_ * (scalars_end_ - scalars_begin_))); - RETURN_NOT_OK(builder->ReserveData(n_repeats_ * data_size)); - - for (int64_t i = 0; i < n_repeats_; i++) { - for (auto it = scalars_begin_; it != scalars_end_; ++it) { - const auto& scalar = checked_cast::ScalarType&>(*it); - if (scalar.is_valid) { - builder->UnsafeAppend(std::string_view{*scalar.value}); - } else { - builder->UnsafeAppendNull(); - } - } - } - return Status::OK(); - } - - template - enable_if_t::value || is_list_like_type::value, Status> Visit( - const T&) { - auto builder = checked_cast::BuilderType*>(builder_); - int64_t num_children = 0; - for (auto it = scalars_begin_; it != scalars_end_; ++it) { - if (!it->is_valid) continue; - num_children += checked_cast(*it).value->length(); - } - RETURN_NOT_OK(builder->value_builder()->Reserve(num_children * n_repeats_)); - - for (int64_t i = 0; i < n_repeats_; i++) { - for (auto it = scalars_begin_; it != scalars_end_; ++it) { - if (it->is_valid) { - const Array& list = *checked_cast(*it).value; - if constexpr (T::type_id == Type::MAP || T::type_id == Type::FIXED_SIZE_LIST) { - RETURN_NOT_OK(builder->Append()); - } else { - RETURN_NOT_OK(builder->Append(/*is_valid=*/true, list.length())); - } - for (int64_t i = 0; i < list.length(); i++) { - ARROW_ASSIGN_OR_RAISE(auto scalar, list.GetScalar(i)); - RETURN_NOT_OK(builder->value_builder()->AppendScalar(*scalar)); - } - } else { - RETURN_NOT_OK(builder_->AppendNull()); - } - } - } - return Status::OK(); - } - - Status Visit(const StructType& type) { - auto* builder = checked_cast(builder_); - auto count = n_repeats_ * (scalars_end_ - scalars_begin_); - RETURN_NOT_OK(builder->Reserve(count)); - for (int field_index = 0; field_index < type.num_fields(); ++field_index) { - RETURN_NOT_OK(builder->field_builder(field_index)->Reserve(count)); - } - for (int64_t i = 0; i < n_repeats_; i++) { - for (auto it = scalars_begin_; it != scalars_end_; ++it) { - const auto& scalar = checked_cast(*it); - for (int field_index = 0; field_index < type.num_fields(); ++field_index) { - if (!scalar.is_valid || !scalar.value[field_index]) { - RETURN_NOT_OK(builder->field_builder(field_index)->AppendNull()); - } else { - RETURN_NOT_OK(builder->field_builder(field_index) - ->AppendScalar(*scalar.value[field_index])); - } - } - RETURN_NOT_OK(builder->Append(scalar.is_valid)); - } - } - return Status::OK(); - } - - Status Visit(const SparseUnionType& type) { return MakeUnionArray(type); } - - Status Visit(const DenseUnionType& type) { return MakeUnionArray(type); } - - Status AppendUnionScalar(const DenseUnionType& type, const Scalar& s, - DenseUnionBuilder* builder) { - const auto& scalar = checked_cast(s); - const auto scalar_field_index = type.child_ids()[scalar.type_code]; - RETURN_NOT_OK(builder->Append(scalar.type_code)); - - for (int field_index = 0; field_index < type.num_fields(); ++field_index) { - auto* child_builder = builder->child_builder(field_index).get(); - if (field_index == scalar_field_index) { - if (scalar.is_valid) { - RETURN_NOT_OK(child_builder->AppendScalar(*scalar.value)); - } else { - RETURN_NOT_OK(child_builder->AppendNull()); - } - } - } - return Status::OK(); - } - - Status AppendUnionScalar(const SparseUnionType& type, const Scalar& s, - SparseUnionBuilder* builder) { - // For each scalar, - // 1. append the type code, - // 2. append the value to the corresponding child, - // 3. append null to the other children. - const auto& scalar = checked_cast(s); - RETURN_NOT_OK(builder->Append(scalar.type_code)); - - for (int field_index = 0; field_index < type.num_fields(); ++field_index) { - auto* child_builder = builder->child_builder(field_index).get(); - if (field_index == scalar.child_id) { - if (scalar.is_valid) { - RETURN_NOT_OK(child_builder->AppendScalar(*scalar.value[field_index])); - } else { - RETURN_NOT_OK(child_builder->AppendNull()); - } - } else { - RETURN_NOT_OK(child_builder->AppendNull()); - } - } - return Status::OK(); - } - - template - Status MakeUnionArray(const T& type) { - using BuilderType = typename TypeTraits::BuilderType; - - auto* builder = checked_cast(builder_); - const auto count = n_repeats_ * (scalars_end_ - scalars_begin_); - - RETURN_NOT_OK(builder->Reserve(count)); - - DCHECK_EQ(type.num_fields(), builder->num_children()); - for (int field_index = 0; field_index < type.num_fields(); ++field_index) { - RETURN_NOT_OK(builder->child_builder(field_index)->Reserve(count)); - } - - for (int64_t i = 0; i < n_repeats_; i++) { - for (auto it = scalars_begin_; it != scalars_end_; ++it) { - RETURN_NOT_OK(AppendUnionScalar(type, *it, builder)); - } - } - return Status::OK(); - } - - Status Visit(const RunEndEncodedType&) { - auto builder = checked_cast(builder_); - - RETURN_NOT_OK(builder->Reserve(n_repeats_ * (scalars_end_ - scalars_begin_))); - - for (int64_t i = 0; i < n_repeats_; i++) { - for (auto it = scalars_begin_; it != scalars_end_; ++it) { - if (it->is_valid) { - const auto& scalar_value = *checked_cast(*it).value; - RETURN_NOT_OK(builder->AppendScalar(scalar_value, 1)); - } else { - RETURN_NOT_OK(builder_->AppendNull()); - } - } - } - return Status::OK(); - } - - Status Visit(const DataType& type) { - return Status::NotImplemented("AppendScalar for type ", type); - } - - Status Convert() { return VisitTypeInline(*scalars_begin_->type, this); } - - ConstIterator scalars_begin_; - ConstIterator scalars_end_; - int64_t n_repeats_; - ArrayBuilder* builder_; -}; - -// Wraps a const_iterator that has a pointer (or pointer-like) to Scalar as the -// value_type and turns it into an iterator with Scalar as value_type. -template -struct DerefConstIterator { - ConstIterator it; - - using value_type = Scalar; - using pointer = const Scalar*; - using difference_type = typename ConstIterator::difference_type; - - const value_type& operator*() const { return *(*it); } - - DerefConstIterator& operator++() { - ++it; - return *this; - } - - difference_type operator-(const DerefConstIterator& other) const { - return it - other.it; - } - - bool operator!=(const DerefConstIterator& other) const { return it != other.it; } - - pointer operator->() const { return &(**it); } -}; - -} // namespace - Status ArrayBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { - if (!scalar.type->Equals(type())) { + if (scalar.type->id() == Type::NA) { + return AppendNulls(n_repeats); + } + if (scalar.type->id() != type()->id()) { return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), " to builder for type ", type()->ToString()); } - return AppendScalarImpl{&scalar, &scalar + 1, n_repeats, this}.Convert(); + return Status::NotImplemented("AppendScalar for builder for ", *type()); } Status ArrayBuilder::AppendScalars(const ScalarVector& scalars) { if (scalars.empty()) return Status::OK(); - const auto ty = type(); for (const auto& scalar : scalars) { - if (!scalar->type->Equals(ty)) { - return Status::Invalid("Cannot append scalar of type ", scalar->type->ToString(), - " to builder for type ", type()->ToString()); - } + RETURN_NOT_OK(AppendScalar(*scalar, 1)); } - - using Iterator = DerefConstIterator; - return AppendScalarImpl{Iterator{scalars.begin()}, Iterator{scalars.end()}, - /*n_repeats=*/1, this} - .Convert(); + return Status::OK(); } Status ArrayBuilder::Finish(std::shared_ptr* out) { diff --git a/cpp/src/arrow/array/builder_binary.cc b/cpp/src/arrow/array/builder_binary.cc index 19c4c2d523f..2936aa1851a 100644 --- a/cpp/src/arrow/array/builder_binary.cc +++ b/cpp/src/arrow/array/builder_binary.cc @@ -102,6 +102,102 @@ void BinaryViewBuilder::Reset() { data_heap_builder_.Reset(); } +Status BinaryViewBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { + if (scalar.type->id() == Type::NA) { + return AppendNulls(n_repeats); + } + if (scalar.type->id() != type()->id()) { + return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), + " to builder for type ", type()->ToString()); + } + const auto& s = checked_cast(scalar); + int64_t data_size = s.is_valid ? s.value->size() : 0; + ARROW_RETURN_NOT_OK(Reserve(n_repeats)); + ARROW_RETURN_NOT_OK(ReserveData(n_repeats * data_size)); + if (s.is_valid) { + std::string_view sv{*s.value}; + for (int64_t i = 0; i < n_repeats; ++i) { + UnsafeAppend(sv); + } + } else { + for (int64_t i = 0; i < n_repeats; ++i) { + UnsafeAppendNull(); + } + } + return Status::OK(); +} + +Status BinaryViewBuilder::AppendScalars(const ScalarVector& scalars) { + if (scalars.empty()) return Status::OK(); + const auto ty = type(); + int64_t data_size = 0; + for (const auto& scalar : scalars) { + if (!scalar->type->Equals(ty)) { + return Status::Invalid("Cannot append scalar of type ", scalar->type->ToString(), + " to builder for type ", type()->ToString()); + } + const auto& s = checked_cast(*scalar); + if (s.is_valid) { + data_size += s.value->size(); + } + } + ARROW_RETURN_NOT_OK(Reserve(static_cast(scalars.size()))); + ARROW_RETURN_NOT_OK(ReserveData(data_size)); + for (const auto& scalar : scalars) { + const auto& s = checked_cast(*scalar); + if (s.is_valid) { + UnsafeAppend(std::string_view{*s.value}); + } else { + UnsafeAppendNull(); + } + } + return Status::OK(); +} + +Status FixedSizeBinaryBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { + if (scalar.type->id() == Type::NA) { + return AppendNulls(n_repeats); + } + if (scalar.type->id() != type()->id()) { + return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), + " to builder for type ", type()->ToString()); + } + const auto& s = checked_cast(scalar); + ARROW_RETURN_NOT_OK(Reserve(n_repeats)); + if (s.is_valid) { + std::string_view sv{*s.value}; + for (int64_t i = 0; i < n_repeats; ++i) { + UnsafeAppend(sv); + } + } else { + for (int64_t i = 0; i < n_repeats; ++i) { + UnsafeAppendNull(); + } + } + return Status::OK(); +} + +Status FixedSizeBinaryBuilder::AppendScalars(const ScalarVector& scalars) { + if (scalars.empty()) return Status::OK(); + const auto ty = type(); + for (const auto& scalar : scalars) { + if (!scalar->type->Equals(ty)) { + return Status::Invalid("Cannot append scalar of type ", scalar->type->ToString(), + " to builder for type ", type()->ToString()); + } + } + ARROW_RETURN_NOT_OK(Reserve(static_cast(scalars.size()))); + for (const auto& scalar : scalars) { + const auto& s = checked_cast(*scalar); + if (s.is_valid) { + UnsafeAppend(std::string_view{*s.value}); + } else { + UnsafeAppendNull(); + } + } + return Status::OK(); +} + // ---------------------------------------------------------------------- // Fixed width binary diff --git a/cpp/src/arrow/array/builder_binary.h b/cpp/src/arrow/array/builder_binary.h index d0e761ae968..c68e952db8c 100644 --- a/cpp/src/arrow/array/builder_binary.h +++ b/cpp/src/arrow/array/builder_binary.h @@ -34,9 +34,11 @@ #include "arrow/array/data.h" #include "arrow/buffer.h" #include "arrow/buffer_builder.h" +#include "arrow/scalar.h" #include "arrow/status.h" #include "arrow/type.h" #include "arrow/util/binary_view_util.h" +#include "arrow/util/checked_cast.h" #include "arrow/util/macros.h" #include "arrow/util/visibility.h" @@ -138,6 +140,63 @@ class BaseBinaryBuilder return Status::OK(); } + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override { + if (scalar.type->id() == Type::NA) { + return AppendNulls(n_repeats); + } + if (scalar.type->id() != type()->id()) { + return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), + " to builder for type ", type()->ToString()); + } + const auto& s = + internal::checked_cast::ScalarType&>(scalar); + int64_t data_size = s.is_valid ? s.value->size() : 0; + ARROW_RETURN_NOT_OK(Reserve(n_repeats)); + ARROW_RETURN_NOT_OK(ReserveData(n_repeats * data_size)); + if (s.is_valid) { + std::string_view sv{*s.value}; + for (int64_t i = 0; i < n_repeats; ++i) { + UnsafeAppend(sv); + } + } else { + for (int64_t i = 0; i < n_repeats; ++i) { + UnsafeAppendNull(); + } + } + return Status::OK(); + } + + Status AppendScalars(const ScalarVector& scalars) override { + if (scalars.empty()) return Status::OK(); + const auto ty_id = type()->id(); + int64_t data_size = 0; + for (const auto& scalar : scalars) { + if (scalar->type->id() != Type::NA && scalar->type->id() != ty_id) { + return Status::Invalid("Cannot append scalar of type ", scalar->type->ToString(), + " to builder for type ", type()->ToString()); + } + if (scalar->type->id() != Type::NA) { + const auto& s = + internal::checked_cast::ScalarType&>(*scalar); + if (s.is_valid) { + data_size += s.value->size(); + } + } + } + ARROW_RETURN_NOT_OK(Reserve(static_cast(scalars.size()))); + ARROW_RETURN_NOT_OK(ReserveData(data_size)); + for (const auto& scalar : scalars) { + if (scalar->type->id() == Type::NA || !scalar->is_valid) { + UnsafeAppendNull(); + } else { + const auto& s = + internal::checked_cast::ScalarType&>(*scalar); + UnsafeAppend(std::string_view{*s.value}); + } + } + return Status::OK(); + } + /// \brief Append without checking capacity /// /// Offsets and data should have been presized using Reserve() and @@ -665,6 +724,9 @@ class ARROW_EXPORT BinaryViewBuilder : public ArrayBuilder { return Status::OK(); } + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; + Status AppendScalars(const ScalarVector& scalars) override; + /// \brief Append a empty element (length-0 inline string) Status AppendEmptyValue() final { ARROW_RETURN_NOT_OK(Reserve(1)); @@ -781,6 +843,9 @@ class ARROW_EXPORT FixedSizeBinaryBuilder : public ArrayBuilder { Status AppendNull() final; Status AppendNulls(int64_t length) final; + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; + Status AppendScalars(const ScalarVector& scalars) override; + Status AppendEmptyValue() final; Status AppendEmptyValues(int64_t length) final; diff --git a/cpp/src/arrow/array/builder_decimal.cc b/cpp/src/arrow/array/builder_decimal.cc index 868183768c1..febdd7e66f8 100644 --- a/cpp/src/arrow/array/builder_decimal.cc +++ b/cpp/src/arrow/array/builder_decimal.cc @@ -23,6 +23,7 @@ #include "arrow/array/data.h" #include "arrow/buffer.h" #include "arrow/buffer_builder.h" +#include "arrow/scalar.h" #include "arrow/status.h" #include "arrow/util/checked_cast.h" #include "arrow/util/decimal.h" @@ -32,6 +33,56 @@ namespace arrow { class Buffer; class MemoryPool; +namespace { + +template +Status DecimalAppendScalar(BuilderType* builder, const Scalar& scalar, + int64_t n_repeats) { + if (scalar.type->id() == Type::NA) { + return builder->AppendNulls(n_repeats); + } + if (scalar.type->id() != builder->type()->id()) { + return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), + " to builder for type ", builder->type()->ToString()); + } + const auto& s = internal::checked_cast(scalar); + ARROW_RETURN_NOT_OK(builder->Reserve(n_repeats)); + if (s.is_valid) { + for (int64_t i = 0; i < n_repeats; ++i) { + builder->UnsafeAppend(s.value); + } + } else { + for (int64_t i = 0; i < n_repeats; ++i) { + builder->UnsafeAppendNull(); + } + } + return Status::OK(); +} + +template +Status DecimalAppendScalars(BuilderType* builder, const ScalarVector& scalars) { + if (scalars.empty()) return Status::OK(); + const auto ty_id = builder->type()->id(); + for (const auto& scalar : scalars) { + if (scalar->type->id() != Type::NA && scalar->type->id() != ty_id) { + return Status::Invalid("Cannot append scalar of type ", scalar->type->ToString(), + " to builder for type ", builder->type()->ToString()); + } + } + ARROW_RETURN_NOT_OK(builder->Reserve(static_cast(scalars.size()))); + for (const auto& scalar : scalars) { + if (scalar->type->id() == Type::NA || !scalar->is_valid) { + builder->UnsafeAppendNull(); + } else { + const auto& s = internal::checked_cast(*scalar); + builder->UnsafeAppend(s.value); + } + } + return Status::OK(); +} + +} // namespace + // ---------------------------------------------------------------------- // Decimal32Builder @@ -172,4 +223,38 @@ Status Decimal256Builder::FinishInternal(std::shared_ptr* out) { return Status::OK(); } +Status Decimal32Builder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { + return DecimalAppendScalar(this, scalar, n_repeats); +} + +Status Decimal32Builder::AppendScalars(const ScalarVector& scalars) { + return DecimalAppendScalars(this, scalars); +} + +Status Decimal64Builder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { + return DecimalAppendScalar(this, scalar, n_repeats); +} + +Status Decimal64Builder::AppendScalars(const ScalarVector& scalars) { + return DecimalAppendScalars(this, scalars); +} + +Status Decimal128Builder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { + return DecimalAppendScalar(this, scalar, + n_repeats); +} + +Status Decimal128Builder::AppendScalars(const ScalarVector& scalars) { + return DecimalAppendScalars(this, scalars); +} + +Status Decimal256Builder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { + return DecimalAppendScalar(this, scalar, + n_repeats); +} + +Status Decimal256Builder::AppendScalars(const ScalarVector& scalars) { + return DecimalAppendScalars(this, scalars); +} + } // namespace arrow diff --git a/cpp/src/arrow/array/builder_decimal.h b/cpp/src/arrow/array/builder_decimal.h index a0bf0a04220..a6dc7d944a3 100644 --- a/cpp/src/arrow/array/builder_decimal.h +++ b/cpp/src/arrow/array/builder_decimal.h @@ -50,6 +50,9 @@ class ARROW_EXPORT Decimal32Builder : public FixedSizeBinaryBuilder { void UnsafeAppend(Decimal32 val); void UnsafeAppend(std::string_view val); + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; + Status AppendScalars(const ScalarVector& scalars) override; + Status FinishInternal(std::shared_ptr* out) override; /// \cond FALSE @@ -81,6 +84,9 @@ class ARROW_EXPORT Decimal64Builder : public FixedSizeBinaryBuilder { void UnsafeAppend(Decimal64 val); void UnsafeAppend(std::string_view val); + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; + Status AppendScalars(const ScalarVector& scalars) override; + Status FinishInternal(std::shared_ptr* out) override; /// \cond FALSE @@ -112,6 +118,9 @@ class ARROW_EXPORT Decimal128Builder : public FixedSizeBinaryBuilder { void UnsafeAppend(Decimal128 val); void UnsafeAppend(std::string_view val); + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; + Status AppendScalars(const ScalarVector& scalars) override; + Status FinishInternal(std::shared_ptr* out) override; /// \cond FALSE @@ -143,6 +152,9 @@ class ARROW_EXPORT Decimal256Builder : public FixedSizeBinaryBuilder { void UnsafeAppend(const Decimal256& val); void UnsafeAppend(std::string_view val); + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; + Status AppendScalars(const ScalarVector& scalars) override; + Status FinishInternal(std::shared_ptr* out) override; /// \cond FALSE diff --git a/cpp/src/arrow/array/builder_nested.cc b/cpp/src/arrow/array/builder_nested.cc index 915fbfbf895..ba59eeaa0c6 100644 --- a/cpp/src/arrow/array/builder_nested.cc +++ b/cpp/src/arrow/array/builder_nested.cc @@ -22,7 +22,9 @@ #include #include +#include "arrow/array.h" #include "arrow/buffer.h" +#include "arrow/scalar.h" #include "arrow/status.h" #include "arrow/type.h" #include "arrow/util/checked_cast.h" @@ -30,20 +32,6 @@ namespace arrow { -// ---------------------------------------------------------------------- -// VarLengthListLikeBuilder / BaseListBuilder / BaseListViewBuilder - -template class VarLengthListLikeBuilder; -template class VarLengthListLikeBuilder; -template class VarLengthListLikeBuilder; -template class VarLengthListLikeBuilder; - -template class BaseListBuilder; -template class BaseListBuilder; - -template class BaseListViewBuilder; -template class BaseListViewBuilder; - // ---------------------------------------------------------------------- // MapBuilder @@ -313,4 +301,131 @@ std::shared_ptr StructBuilder::type() const { return struct_(std::move(fields)); } +template +Status VarLengthListLikeBuilder::AppendScalar(const Scalar& scalar, + int64_t n_repeats) { + if (scalar.type->id() == Type::NA) { + return AppendNulls(n_repeats); + } + if (scalar.type->id() != type()->id()) { + return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), + " to builder for type ", type()->ToString()); + } + const auto& s = internal::checked_cast(scalar); + if (s.is_valid) { + const Array& list = *s.value; + RETURN_NOT_OK(value_builder_->Reserve(list.length() * n_repeats)); + for (int64_t r = 0; r < n_repeats; ++r) { + if constexpr (TYPE::type_id == Type::MAP || + TYPE::type_id == Type::FIXED_SIZE_LIST) { + RETURN_NOT_OK(Append()); + } else { + RETURN_NOT_OK(Append(/*is_valid=*/true, list.length())); + } + for (int64_t i = 0; i < list.length(); ++i) { + ARROW_ASSIGN_OR_RAISE(auto child_scalar, list.GetScalar(i)); + RETURN_NOT_OK(value_builder_->AppendScalar(*child_scalar)); + } + } + } else { + for (int64_t r = 0; r < n_repeats; ++r) { + RETURN_NOT_OK(AppendNull()); + } + } + return Status::OK(); +} + +Status MapBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { + if (scalar.type->id() == Type::NA) { + return AppendNulls(n_repeats); + } + if (scalar.type->id() != type()->id()) { + return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), + " to builder for type ", type()->ToString()); + } + const auto& s = internal::checked_cast(scalar); + if (s.is_valid) { + const Array& list = *s.value; + RETURN_NOT_OK(value_builder()->Reserve(list.length() * n_repeats)); + for (int64_t r = 0; r < n_repeats; ++r) { + RETURN_NOT_OK(Append()); + for (int64_t i = 0; i < list.length(); ++i) { + ARROW_ASSIGN_OR_RAISE(auto child_scalar, list.GetScalar(i)); + RETURN_NOT_OK(value_builder()->AppendScalar(*child_scalar)); + } + } + } else { + for (int64_t r = 0; r < n_repeats; ++r) { + RETURN_NOT_OK(AppendNull()); + } + } + return Status::OK(); +} + +Status FixedSizeListBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { + if (scalar.type->id() == Type::NA) { + return AppendNulls(n_repeats); + } + if (scalar.type->id() != type()->id()) { + return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), + " to builder for type ", type()->ToString()); + } + const auto& s = internal::checked_cast(scalar); + if (s.is_valid) { + const Array& list = *s.value; + RETURN_NOT_OK(value_builder_->Reserve(list.length() * n_repeats)); + for (int64_t r = 0; r < n_repeats; ++r) { + RETURN_NOT_OK(Append()); + for (int64_t i = 0; i < list.length(); ++i) { + ARROW_ASSIGN_OR_RAISE(auto child_scalar, list.GetScalar(i)); + RETURN_NOT_OK(value_builder_->AppendScalar(*child_scalar)); + } + } + } else { + for (int64_t r = 0; r < n_repeats; ++r) { + RETURN_NOT_OK(AppendNull()); + } + } + return Status::OK(); +} + +Status StructBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { + if (scalar.type->id() == Type::NA) { + return AppendNulls(n_repeats); + } + if (scalar.type->id() != type()->id()) { + return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), + " to builder for type ", type()->ToString()); + } + const auto& s = internal::checked_cast(scalar); + const int num_fields_count = static_cast(children_.size()); + RETURN_NOT_OK(Reserve(n_repeats)); + for (int field_index = 0; field_index < num_fields_count; ++field_index) { + RETURN_NOT_OK(field_builder(field_index)->Reserve(n_repeats)); + } + for (int64_t r = 0; r < n_repeats; ++r) { + for (int field_index = 0; field_index < num_fields_count; ++field_index) { + if (!s.is_valid || field_index >= static_cast(s.value.size()) || + !s.value[field_index]) { + RETURN_NOT_OK(field_builder(field_index)->AppendNull()); + } else { + RETURN_NOT_OK(field_builder(field_index)->AppendScalar(*s.value[field_index])); + } + } + RETURN_NOT_OK(Append(s.is_valid)); + } + return Status::OK(); +} + +template class VarLengthListLikeBuilder; +template class VarLengthListLikeBuilder; +template class VarLengthListLikeBuilder; +template class VarLengthListLikeBuilder; + +template class BaseListBuilder; +template class BaseListBuilder; + +template class BaseListViewBuilder; +template class BaseListViewBuilder; + } // namespace arrow diff --git a/cpp/src/arrow/array/builder_nested.h b/cpp/src/arrow/array/builder_nested.h index fdbeb0cd7d1..d161d096560 100644 --- a/cpp/src/arrow/array/builder_nested.h +++ b/cpp/src/arrow/array/builder_nested.h @@ -159,6 +159,8 @@ class VarLengthListLikeBuilder : public ArrayBuilder { return Status::OK(); } + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; + /// \brief Vector append /// /// For list-array builders, the sizes are inferred from the offsets. @@ -561,6 +563,8 @@ class ARROW_EXPORT MapBuilder : public ArrayBuilder { Status AppendNulls(int64_t length) final; + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; + Status AppendEmptyValue() final; Status AppendEmptyValues(int64_t length) final; @@ -694,6 +698,8 @@ class ARROW_EXPORT FixedSizeListBuilder : public ArrayBuilder { /// automatically. Status AppendNulls(int64_t length) final; + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; + Status ValidateOverflow(int64_t new_elements); Status AppendEmptyValue() final; @@ -807,6 +813,8 @@ class ARROW_EXPORT StructBuilder : public ArrayBuilder { return Status::OK(); } + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; + Status AppendArraySlice(const ArraySpan& array, int64_t offset, int64_t length) override { for (int i = 0; static_cast(i) < children_.size(); i++) { diff --git a/cpp/src/arrow/array/builder_primitive.h b/cpp/src/arrow/array/builder_primitive.h index 6d79d6e9649..785bd7c0a40 100644 --- a/cpp/src/arrow/array/builder_primitive.h +++ b/cpp/src/arrow/array/builder_primitive.h @@ -24,8 +24,10 @@ #include "arrow/array/builder_base.h" #include "arrow/array/data.h" #include "arrow/result.h" +#include "arrow/scalar.h" #include "arrow/type.h" #include "arrow/type_traits.h" +#include "arrow/util/checked_cast.h" #include "arrow/util/float16.h" namespace arrow { @@ -58,6 +60,25 @@ class ARROW_EXPORT NullBuilder : public ArrayBuilder { Status Append(std::nullptr_t) { return AppendNull(); } + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override { + if (!scalar.type->Equals(type())) { + return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), + " to builder for type ", type()->ToString()); + } + return AppendNulls(n_repeats); + } + + Status AppendScalars(const ScalarVector& scalars) override { + if (scalars.empty()) return Status::OK(); + for (const auto& scalar : scalars) { + if (!scalar->type->Equals(type())) { + return Status::Invalid("Cannot append scalar of type ", scalar->type->ToString(), + " to builder for type ", type()->ToString()); + } + } + return AppendNulls(static_cast(scalars.size())); + } + Status AppendArraySlice(const ArraySpan&, int64_t, int64_t length) override { return AppendNulls(length); } @@ -140,6 +161,51 @@ class NumericBuilder return Status::OK(); } + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override { + if (scalar.type->id() == Type::NA) { + return AppendNulls(n_repeats); + } + if (scalar.type->id() != type()->id()) { + return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), + " to builder for type ", type()->ToString()); + } + const auto& s = + internal::checked_cast::ScalarType&>(scalar); + ARROW_RETURN_NOT_OK(Reserve(n_repeats)); + if (s.is_valid) { + for (int64_t i = 0; i < n_repeats; ++i) { + UnsafeAppend(s.value); + } + } else { + for (int64_t i = 0; i < n_repeats; ++i) { + UnsafeAppendNull(); + } + } + return Status::OK(); + } + + Status AppendScalars(const ScalarVector& scalars) override { + if (scalars.empty()) return Status::OK(); + const auto ty_id = type()->id(); + for (const auto& scalar : scalars) { + if (scalar->type->id() != Type::NA && scalar->type->id() != ty_id) { + return Status::Invalid("Cannot append scalar of type ", scalar->type->ToString(), + " to builder for type ", type()->ToString()); + } + } + ARROW_RETURN_NOT_OK(Reserve(static_cast(scalars.size()))); + for (const auto& scalar : scalars) { + if (scalar->type->id() == Type::NA || !scalar->is_valid) { + UnsafeAppendNull(); + } else { + const auto& s = + internal::checked_cast::ScalarType&>(*scalar); + UnsafeAppend(s.value); + } + } + return Status::OK(); + } + value_type GetValue(int64_t index) const { return data_builder_.data()[index]; } value_type* GetMutableValue(int64_t index) { @@ -534,6 +600,49 @@ class ARROW_EXPORT BooleanBuilder return Status::OK(); } + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override { + if (scalar.type->id() == Type::NA) { + return AppendNulls(n_repeats); + } + if (scalar.type->id() != type()->id()) { + return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), + " to builder for type ", type()->ToString()); + } + const auto& s = internal::checked_cast(scalar); + ARROW_RETURN_NOT_OK(Reserve(n_repeats)); + if (s.is_valid) { + for (int64_t i = 0; i < n_repeats; ++i) { + UnsafeAppend(s.value); + } + } else { + for (int64_t i = 0; i < n_repeats; ++i) { + UnsafeAppendNull(); + } + } + return Status::OK(); + } + + Status AppendScalars(const ScalarVector& scalars) override { + if (scalars.empty()) return Status::OK(); + const auto ty_id = type()->id(); + for (const auto& scalar : scalars) { + if (scalar->type->id() != Type::NA && scalar->type->id() != ty_id) { + return Status::Invalid("Cannot append scalar of type ", scalar->type->ToString(), + " to builder for type ", type()->ToString()); + } + } + ARROW_RETURN_NOT_OK(Reserve(static_cast(scalars.size()))); + for (const auto& scalar : scalars) { + if (scalar->type->id() == Type::NA || !scalar->is_valid) { + UnsafeAppendNull(); + } else { + const auto& s = internal::checked_cast(*scalar); + UnsafeAppend(s.value); + } + } + return Status::OK(); + } + Status Append(const uint8_t val) { return Append(val != 0); } /// Scalar append, without checking for capacity diff --git a/cpp/src/arrow/array/builder_union.cc b/cpp/src/arrow/array/builder_union.cc index 1ced2a21bee..6977903c823 100644 --- a/cpp/src/arrow/array/builder_union.cc +++ b/cpp/src/arrow/array/builder_union.cc @@ -21,6 +21,7 @@ #include #include "arrow/buffer.h" +#include "arrow/scalar.h" #include "arrow/util/checked_cast.h" #include "arrow/util/logging_internal.h" @@ -151,4 +152,75 @@ Status SparseUnionBuilder::AppendArraySlice(const ArraySpan& array, const int64_ return Status::OK(); } +Status DenseUnionBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { + if (scalar.type->id() == Type::NA) { + return AppendNulls(n_repeats); + } + if (scalar.type->id() != type()->id()) { + return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), + " to builder for type ", type()->ToString()); + } + const auto& union_type = checked_cast(*type()); + const auto& s = checked_cast(scalar); + const auto scalar_field_index = union_type.child_ids()[s.type_code]; + + RETURN_NOT_OK(Reserve(n_repeats)); + for (int field_index = 0; field_index < union_type.num_fields(); ++field_index) { + if (field_index == scalar_field_index) { + RETURN_NOT_OK(child_builder(field_index)->Reserve(n_repeats)); + } + } + + for (int64_t r = 0; r < n_repeats; ++r) { + RETURN_NOT_OK(Append(s.type_code)); + for (int field_index = 0; field_index < union_type.num_fields(); ++field_index) { + auto* cb = child_builder(field_index).get(); + if (field_index == scalar_field_index) { + if (s.is_valid) { + RETURN_NOT_OK(cb->AppendScalar(*s.value)); + } else { + RETURN_NOT_OK(cb->AppendNull()); + } + } + } + } + return Status::OK(); +} + +Status SparseUnionBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { + if (scalar.type->id() == Type::NA) { + return AppendNulls(n_repeats); + } + if (scalar.type->id() != type()->id()) { + return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), + " to builder for type ", type()->ToString()); + } + const auto& union_type = checked_cast(*type()); + const auto& s = checked_cast(scalar); + const auto scalar_field_index = union_type.child_ids()[s.type_code]; + + RETURN_NOT_OK(Reserve(n_repeats)); + for (int field_index = 0; field_index < union_type.num_fields(); ++field_index) { + RETURN_NOT_OK(child_builder(field_index)->Reserve(n_repeats)); + } + + for (int64_t r = 0; r < n_repeats; ++r) { + RETURN_NOT_OK(Append(s.type_code)); + for (int field_index = 0; field_index < union_type.num_fields(); ++field_index) { + auto* cb = child_builder(field_index).get(); + if (field_index == scalar_field_index) { + if (s.is_valid && field_index < static_cast(s.value.size()) && + s.value[field_index]) { + RETURN_NOT_OK(cb->AppendScalar(*s.value[field_index])); + } else { + RETURN_NOT_OK(cb->AppendNull()); + } + } else { + RETURN_NOT_OK(cb->AppendNull()); + } + } + } + return Status::OK(); +} + } // namespace arrow diff --git a/cpp/src/arrow/array/builder_union.h b/cpp/src/arrow/array/builder_union.h index 718ef4c32ce..d99e1838ac8 100644 --- a/cpp/src/arrow/array/builder_union.h +++ b/cpp/src/arrow/array/builder_union.h @@ -167,6 +167,8 @@ class ARROW_EXPORT DenseUnionBuilder : public BasicUnionBuilder { Status AppendArraySlice(const ArraySpan& array, int64_t offset, int64_t length) override; + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; + Status FinishInternal(std::shared_ptr* out) override; private: @@ -247,6 +249,8 @@ class ARROW_EXPORT SparseUnionBuilder : public BasicUnionBuilder { Status AppendArraySlice(const ArraySpan& array, int64_t offset, int64_t length) override; + + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; }; /// @}