diff --git a/cpp/src/arrow/array/array_test.cc b/cpp/src/arrow/array/array_test.cc index dcfe1c76c301..c11d23d886e2 100644 --- a/cpp/src/arrow/array/array_test.cc +++ b/cpp/src/arrow/array/array_test.cc @@ -767,6 +767,62 @@ TEST_F(TestArray, TestMakeArrayFromScalar) { } } +TEST_F(TestArray, TestAppendScalarValidation) { + // Untyped NullScalar (Type::NA) rejected by Int32Builder + Int32Builder int_builder(pool_); + ASSERT_RAISES(Invalid, int_builder.AppendScalar(NullScalar(), 1)); + + // Typed null scalar accepted + auto typed_null = MakeNullScalar(int32()); + ASSERT_OK(int_builder.AppendScalar(*typed_null, 1)); + ASSERT_EQ(int_builder.length(), 1); + ASSERT_EQ(int_builder.null_count(), 1); + + // Parameterized type mismatches + // FixedSizeBinary byte width mismatch + FixedSizeBinaryBuilder fsb_builder(fixed_size_binary(3), pool_); + FixedSizeBinaryScalar fsb_scalar(Buffer::FromString("12345"), fixed_size_binary(5)); + ASSERT_RAISES(Invalid, fsb_builder.AppendScalar(fsb_scalar, 1)); + + // Decimal scale mismatch + Decimal128Builder dec_builder(decimal128(10, 2), pool_); + Decimal128Scalar dec_scalar(Decimal128(12345), decimal128(10, 4)); + ASSERT_RAISES(Invalid, dec_builder.AppendScalar(dec_scalar, 1)); + + // Timestamp unit mismatch + TimestampBuilder ts_builder(timestamp(TimeUnit::MICRO), pool_); + TimestampScalar ts_scalar(100, timestamp(TimeUnit::SECOND)); + ASSERT_RAISES(Invalid, ts_builder.AppendScalar(ts_scalar, 1)); + + // FixedSizeList size mismatch + auto fsl_type = fixed_size_list(int32(), 3); + std::unique_ptr fsl_builder_ptr; + ASSERT_OK(MakeBuilder(pool_, fsl_type, &fsl_builder_ptr)); + Int32Builder child_builder(pool_); + ASSERT_OK(child_builder.AppendValues({1, 2, 3, 4, 5})); + ASSERT_OK_AND_ASSIGN(auto child_array, child_builder.Finish()); + FixedSizeListScalar fsl_scalar(child_array, fixed_size_list(int32(), 5)); + ASSERT_RAISES(Invalid, fsl_builder_ptr->AppendScalar(fsl_scalar, 1)); + + // 1-argument AppendScalar call directly on concrete builders (verifying no method hiding) + ASSERT_OK(int_builder.AppendScalar(*MakeScalar(42))); + FixedSizeBinaryScalar valid_fsb_scalar(Buffer::FromString("123"), fixed_size_binary(3)); + ASSERT_OK(fsb_builder.AppendScalar(valid_fsb_scalar)); + Decimal128Scalar valid_dec_scalar(Decimal128(12345), decimal128(10, 2)); + ASSERT_OK(dec_builder.AppendScalar(valid_dec_scalar)); + // ValidateAppendScalars with temporary builder type + ScalarVector fsb_scalars = {std::make_shared(valid_fsb_scalar)}; + ASSERT_OK(fsb_builder.AppendScalars(fsb_scalars)); +} + +TEST_F(TestArray, TestAppendScalarsAtomicPreValidation) { + Int32Builder builder(pool_); + ScalarVector scalars = {MakeScalar(10), MakeScalar(20), + MakeScalar("string_type_mismatch")}; + ASSERT_RAISES(Invalid, builder.AppendScalars(scalars)); + ASSERT_EQ(builder.length(), 0); +} + TEST_F(TestArray, TestMakeArrayFromScalarSliced) { // Regression test for ARROW-13437 auto scalars = GetScalars(); diff --git a/cpp/src/arrow/array/builder_base.cc b/cpp/src/arrow/array/builder_base.cc index eea394977f45..25862ff8a07c 100644 --- a/cpp/src/arrow/array/builder_base.cc +++ b/cpp/src/arrow/array/builder_base.cc @@ -73,271 +73,43 @@ Status ArrayBuilder::Resize(int64_t capacity) { return null_bitmap_builder_.Resize(capacity); } -namespace { +namespace internal { -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 ValidateAppendScalar(const ArrayBuilder& builder, const Scalar& scalar) { + const auto type = builder.type(); + if (!scalar.type->Equals(*type)) { + return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), + " to builder for type ", type->ToString()); } + 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()); - } - } +Status ValidateAppendScalars(const ArrayBuilder& builder, const ScalarVector& scalars) { + if (scalars.empty()) return Status::OK(); + const auto type = builder.type(); + 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 ", ty.ToString()); } - return Status::OK(); - } - - Status Visit(const DataType& type) { - return Status::NotImplemented("AppendScalar for type ", type); } + return Status::OK(); +} - 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 +} // namespace internal Status ArrayBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { - if (!scalar.type->Equals(type())) { - 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(); + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalar(*this, scalar)); + return Status::NotImplemented("AppendScalar for builder for ", *type()); } Status ArrayBuilder::AppendScalars(const ScalarVector& scalars) { - if (scalars.empty()) return Status::OK(); - const auto ty = type(); + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalars(*this, scalars)); 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_base.h b/cpp/src/arrow/array/builder_base.h index 0cd57ff04dac..50f52a07e968 100644 --- a/cpp/src/arrow/array/builder_base.h +++ b/cpp/src/arrow/array/builder_base.h @@ -38,6 +38,12 @@ namespace arrow { namespace internal { +ARROW_EXPORT +Status ValidateAppendScalar(const ArrayBuilder& builder, const Scalar& scalar); + +ARROW_EXPORT +Status ValidateAppendScalars(const ArrayBuilder& builder, const ScalarVector& scalars); + template class ArrayBuilderExtraOps { public: diff --git a/cpp/src/arrow/array/builder_binary.cc b/cpp/src/arrow/array/builder_binary.cc index 19c4c2d523fe..ba966842d118 100644 --- a/cpp/src/arrow/array/builder_binary.cc +++ b/cpp/src/arrow/array/builder_binary.cc @@ -102,6 +102,78 @@ void BinaryViewBuilder::Reset() { data_heap_builder_.Reset(); } +Status BinaryViewBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalar(*this, scalar)); + 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) { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalars(*this, scalars)); + int64_t data_size = 0; + for (const auto& scalar : scalars) { + 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) { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalar(*this, scalar)); + 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) { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalars(*this, scalars)); + 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 d0e761ae9684..75bb40f5c6de 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,64 @@ class BaseBinaryBuilder return Status::OK(); } + using ArrayBuilder::AppendScalar; + 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 +725,10 @@ class ARROW_EXPORT BinaryViewBuilder : public ArrayBuilder { return Status::OK(); } + using ArrayBuilder::AppendScalar; + 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 +845,10 @@ class ARROW_EXPORT FixedSizeBinaryBuilder : public ArrayBuilder { Status AppendNull() final; Status AppendNulls(int64_t length) final; + using ArrayBuilder::AppendScalar; + 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 868183768c1d..3a343a489281 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,43 @@ namespace arrow { class Buffer; class MemoryPool; +namespace { + +template +Status DecimalAppendScalar(BuilderType* builder, const Scalar& scalar, + int64_t n_repeats) { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalar(*builder, scalar)); + 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) { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalars(*builder, scalars)); + ARROW_RETURN_NOT_OK(builder->Reserve(static_cast(scalars.size()))); + for (const auto& scalar : scalars) { + if (!scalar->is_valid) { + builder->UnsafeAppendNull(); + } else { + const auto& s = internal::checked_cast(*scalar); + builder->UnsafeAppend(s.value); + } + } + return Status::OK(); +} + +} // namespace + // ---------------------------------------------------------------------- // Decimal32Builder @@ -172,4 +210,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 a0bf0a042208..986c5f2a69c7 100644 --- a/cpp/src/arrow/array/builder_decimal.h +++ b/cpp/src/arrow/array/builder_decimal.h @@ -50,6 +50,10 @@ class ARROW_EXPORT Decimal32Builder : public FixedSizeBinaryBuilder { void UnsafeAppend(Decimal32 val); void UnsafeAppend(std::string_view val); + using ArrayBuilder::AppendScalar; + 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 +85,10 @@ class ARROW_EXPORT Decimal64Builder : public FixedSizeBinaryBuilder { void UnsafeAppend(Decimal64 val); void UnsafeAppend(std::string_view val); + using ArrayBuilder::AppendScalar; + 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 +120,10 @@ class ARROW_EXPORT Decimal128Builder : public FixedSizeBinaryBuilder { void UnsafeAppend(Decimal128 val); void UnsafeAppend(std::string_view val); + using ArrayBuilder::AppendScalar; + 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 +155,10 @@ class ARROW_EXPORT Decimal256Builder : public FixedSizeBinaryBuilder { void UnsafeAppend(const Decimal256& val); void UnsafeAppend(std::string_view val); + using ArrayBuilder::AppendScalar; + 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 915fbfbf895d..99d94b0bff87 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,80 @@ std::shared_ptr StructBuilder::type() const { return struct_(std::move(fields)); } + +Status MapBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalar(*this, scalar)); + 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) { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalar(*this, scalar)); + 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) { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalar(*this, scalar)); + 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 fdbeb0cd7d17..35aee0cdd6a7 100644 --- a/cpp/src/arrow/array/builder_nested.h +++ b/cpp/src/arrow/array/builder_nested.h @@ -159,6 +159,33 @@ class VarLengthListLikeBuilder : public ArrayBuilder { return Status::OK(); } + using ArrayBuilder::AppendScalar; + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalar(*this, scalar)); + 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(); + } + /// \brief Vector append /// /// For list-array builders, the sizes are inferred from the offsets. @@ -561,6 +588,9 @@ class ARROW_EXPORT MapBuilder : public ArrayBuilder { Status AppendNulls(int64_t length) final; + using ArrayBuilder::AppendScalar; + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; + Status AppendEmptyValue() final; Status AppendEmptyValues(int64_t length) final; @@ -694,6 +724,9 @@ class ARROW_EXPORT FixedSizeListBuilder : public ArrayBuilder { /// automatically. Status AppendNulls(int64_t length) final; + using ArrayBuilder::AppendScalar; + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; + Status ValidateOverflow(int64_t new_elements); Status AppendEmptyValue() final; @@ -807,6 +840,9 @@ class ARROW_EXPORT StructBuilder : public ArrayBuilder { return Status::OK(); } + using ArrayBuilder::AppendScalar; + 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 6d79d6e96499..c9b307e524ca 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,26 @@ class ARROW_EXPORT NullBuilder : public ArrayBuilder { Status Append(std::nullptr_t) { return AppendNull(); } + using ArrayBuilder::AppendScalar; + 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 +162,39 @@ class NumericBuilder return Status::OK(); } + using ArrayBuilder::AppendScalar; + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalar(*this, scalar)); + 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 { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalars(*this, scalars)); + ARROW_RETURN_NOT_OK(Reserve(static_cast(scalars.size()))); + for (const auto& scalar : scalars) { + if (!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 +589,37 @@ class ARROW_EXPORT BooleanBuilder return Status::OK(); } + using ArrayBuilder::AppendScalar; + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalar(*this, scalar)); + 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 { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalars(*this, scalars)); + ARROW_RETURN_NOT_OK(Reserve(static_cast(scalars.size()))); + for (const auto& scalar : scalars) { + if (!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_run_end.cc b/cpp/src/arrow/array/builder_run_end.cc index 2edeaff504d2..f6e4b9022e86 100644 --- a/cpp/src/arrow/array/builder_run_end.cc +++ b/cpp/src/arrow/array/builder_run_end.cc @@ -94,6 +94,7 @@ Status RunCompressorBuilder::AppendEmptyValues(int64_t length) { } Status RunCompressorBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalar(*this, scalar)); if (ARROW_PREDICT_FALSE(n_repeats == 0)) { return Status::OK(); } @@ -122,6 +123,7 @@ Status RunCompressorBuilder::AppendScalars(const ScalarVector& scalars) { if (scalars.empty()) { return Status::OK(); } + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalars(*this, scalars)); RETURN_NOT_OK(ArrayBuilder::AppendScalars(scalars)); UpdateDimensions(); return Status::OK(); @@ -204,6 +206,10 @@ Status RunEndEncodedBuilder::AppendEmptyValues(int64_t length) { Status RunEndEncodedBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { if (scalar.type->id() == Type::RUN_END_ENCODED) { + if (!scalar.type->Equals(*type())) { + return Status::Invalid("Cannot append scalar of type ", scalar.type->ToString(), + " to builder for type ", type()->ToString()); + } return AppendScalar(*internal::checked_cast(scalar).value, n_repeats); } @@ -214,6 +220,17 @@ Status RunEndEncodedBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeat Status RunEndEncodedBuilder::AppendScalars(const ScalarVector& scalars) { if (scalars.empty()) return Status::OK(); + for (const auto& scalar : scalars) { + if (scalar->type->id() == Type::RUN_END_ENCODED) { + if (!scalar->type->Equals(*type())) { + return Status::Invalid("Cannot append scalar of type ", scalar->type->ToString(), + " to builder for type ", type()->ToString()); + } + } else if (!scalar->type->Equals(*type_->value_type())) { + return Status::Invalid("Cannot append scalar of type ", scalar->type->ToString(), + " to builder for type ", type()->ToString()); + } + } for (const auto& scalar : scalars) { RETURN_NOT_OK(AppendScalar(*scalar, 1)); } diff --git a/cpp/src/arrow/array/builder_run_end.h b/cpp/src/arrow/array/builder_run_end.h index ac92efbd0dbe..a52e66db8373 100644 --- a/cpp/src/arrow/array/builder_run_end.h +++ b/cpp/src/arrow/array/builder_run_end.h @@ -115,6 +115,7 @@ class RunCompressorBuilder : public ArrayBuilder { Status AppendEmptyValue() final { return AppendEmptyValues(1); } Status AppendEmptyValues(int64_t length) override; + using ArrayBuilder::AppendScalar; Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; Status AppendScalars(const ScalarVector& scalars) override; @@ -236,6 +237,7 @@ class ARROW_EXPORT RunEndEncodedBuilder : public ArrayBuilder { Status AppendEmptyValue() final { return AppendEmptyValues(1); } Status AppendEmptyValues(int64_t length) override; + using ArrayBuilder::AppendScalar; Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; Status AppendScalars(const ScalarVector& scalars) override; Status AppendArraySlice(const ArraySpan& array, int64_t offset, diff --git a/cpp/src/arrow/array/builder_union.cc b/cpp/src/arrow/array/builder_union.cc index 1ced2a21bee1..16bcbfb05372 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,63 @@ Status SparseUnionBuilder::AppendArraySlice(const ArraySpan& array, const int64_ return Status::OK(); } +Status DenseUnionBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalar(*this, scalar)); + 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) { + ARROW_RETURN_NOT_OK(internal::ValidateAppendScalar(*this, scalar)); + 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 718ef4c32ceb..004b9aef0e48 100644 --- a/cpp/src/arrow/array/builder_union.h +++ b/cpp/src/arrow/array/builder_union.h @@ -167,6 +167,9 @@ class ARROW_EXPORT DenseUnionBuilder : public BasicUnionBuilder { Status AppendArraySlice(const ArraySpan& array, int64_t offset, int64_t length) override; + using ArrayBuilder::AppendScalar; + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; + Status FinishInternal(std::shared_ptr* out) override; private: @@ -247,6 +250,9 @@ class ARROW_EXPORT SparseUnionBuilder : public BasicUnionBuilder { Status AppendArraySlice(const ArraySpan& array, int64_t offset, int64_t length) override; + + using ArrayBuilder::AppendScalar; + Status AppendScalar(const Scalar& scalar, int64_t n_repeats) override; }; /// @} diff --git a/python/pyarrow/parquet/core.py b/python/pyarrow/parquet/core.py index ff880fdcf52c..3d8a5d010dc7 100644 --- a/python/pyarrow/parquet/core.py +++ b/python/pyarrow/parquet/core.py @@ -201,115 +201,143 @@ def convert_single_predicate(col, op, val): filters_to_expression, "10.0.0", DeprecationWarning) +_read_docstring_common = """\ +read_dictionary : list, default None + List of names or column paths (for nested types) to read directly + as DictionaryArray. Only supported for BYTE_ARRAY storage. To read + a flat column as dictionary-encoded pass the column name. For + nested types, you must pass the full column "path", which could be + something like level1.level2.list.item. Refer to the Parquet + file's schema to obtain the paths. +binary_type : pyarrow.DataType, default None + If given, Parquet binary columns will be read as this datatype. + This setting is ignored if a serialized Arrow schema is found in + the Parquet metadata. +list_type : subclass of pyarrow.DataType, default None + If given, non-MAP repeated columns will be read as an instance of + this datatype (either pyarrow.ListType or pyarrow.LargeListType). + This setting is ignored if a serialized Arrow schema is found in + the Parquet metadata. +memory_map : bool, default False + If the source is a file path, use a memory map to read file, which can + improve performance in some environments. +buffer_size : int, default 0 + If positive, perform read buffering when deserializing individual + column chunks. Otherwise IO calls are unbuffered. +pre_buffer : bool, default True + Coalesce and issue file reads in parallel to improve performance on + high-latency filesystems (e.g. S3, GCS). If True, Arrow will use a + background I/O thread pool. If using a filesystem layer that itself + performs readahead (e.g. fsspec's S3FS), disable readahead for best + results. Set to False if you want to prioritize minimal memory usage + over maximum speed. +coerce_int96_timestamp_unit : str, default None + Cast timestamps that are stored in INT96 format to a particular resolution + (e.g. 'ms'). Setting to None is equivalent to 'ns' and therefore INT96 + timestamps will be inferred as timestamps in nanoseconds. +decryption_properties : FileDecryptionProperties or None, default None + File decryption properties for Parquet Modular Encryption. +thrift_string_size_limit : int, default None + If not None, override the maximum total string size allocated + when decoding Thrift structures. The default limit should be + sufficient for most Parquet files. +thrift_container_size_limit : int, default None + If not None, override the maximum total size of containers allocated + when decoding Thrift structures. The default limit should be + sufficient for most Parquet files. +page_checksum_verification : bool, default False + If True, verify the page checksum for each page read from the file. +arrow_extensions_enabled : bool, default True + If True, read Parquet logical types as Arrow extension types where possible, + (e.g., read JSON as the canonical `arrow.json` extension type or UUID as + the canonical `arrow.uuid` extension type).""" + + +_read_docstring_dataset = """\ +partitioning : pyarrow.dataset.Partitioning or str or list of str, \ +default "hive" + The partitioning scheme for a partitioned dataset. The default of "hive" + assumes directory names with key=value pairs like "/year=2009/month=11". + In addition, a scheme like "/2009/11" is also supported, in which case + you need to specify the field names or a full schema. See the + ``pyarrow.dataset.partitioning()`` function for more details. +ignore_prefixes : list, optional + Files matching any of these prefixes will be ignored by the + discovery process. + This is matched to the basename of a path. + By default this is ['.', '_']. + Note that discovery happens only if a directory is passed as source.""" + + # ---------------------------------------------------------------------- # Reading a single Parquet file -class ParquetFile: - """ - Reader interface for a single Parquet file. +_parquet_file_example = """\ +Generate an example PyArrow Table and write it to Parquet file: - Parameters - ---------- - source : str, pathlib.Path, pyarrow.NativeFile, or file-like object - Readable source. For passing bytes or buffer-like file containing a - Parquet file, use pyarrow.BufferReader. - metadata : FileMetaData, default None - Use existing metadata object, rather than reading from file. - common_metadata : FileMetaData, default None - Will be used in reads for pandas schema metadata if not found in the - main file's metadata, no other uses at the moment. - read_dictionary : list - List of column names to read directly as DictionaryArray. - binary_type : pyarrow.DataType, default None - If given, Parquet binary columns will be read as this datatype. - This setting is ignored if a serialized Arrow schema is found in - the Parquet metadata. - list_type : subclass of pyarrow.DataType, default None - If given, non-MAP repeated columns will be read as an instance of - this datatype (either pyarrow.ListType or pyarrow.LargeListType). - This setting is ignored if a serialized Arrow schema is found in - the Parquet metadata. - memory_map : bool, default False - If the source is a file path, use a memory map to read file, which can - improve performance in some environments. - buffer_size : int, default 0 - If positive, perform read buffering when deserializing individual - column chunks. Otherwise IO calls are unbuffered. - pre_buffer : bool, default True - Coalesce and issue file reads in parallel to improve performance on - high-latency filesystems (e.g. S3, GCS). If True, Arrow will use a - background I/O thread pool. If using a filesystem layer that itself - performs readahead (e.g. fsspec's S3FS), disable readahead for best - results. Set to False if you want to prioritize minimal memory usage - over maximum speed. - coerce_int96_timestamp_unit : str, default None - Cast timestamps that are stored in INT96 format to a particular - resolution (e.g. 'ms'). Setting to None is equivalent to 'ns' - and therefore INT96 timestamps will be inferred as timestamps - in nanoseconds. - decryption_properties : FileDecryptionProperties, default None - File decryption properties for Parquet Modular Encryption. - thrift_string_size_limit : int, default None - If not None, override the maximum total string size allocated - when decoding Thrift structures. The default limit should be - sufficient for most Parquet files. - thrift_container_size_limit : int, default None - If not None, override the maximum total size of containers allocated - when decoding Thrift structures. The default limit should be - sufficient for most Parquet files. - filesystem : FileSystem, default None - If nothing passed, will be inferred based on path. - Path will try to be found in the local on-disk filesystem otherwise - it will be parsed as an URI to determine the filesystem. - page_checksum_verification : bool, default False - If True, verify the checksum for each page read from the file. - arrow_extensions_enabled : bool, default True - If True, read Parquet logical types as Arrow extension types where - possible (e.g., read JSON as the canonical `arrow.json` extension type - or UUID as the canonical `arrow.uuid` extension type). +>>> import pyarrow as pa +>>> table = pa.table({'n_legs': [2, 2, 4, 4, 5, 100], +... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", +... "Brittle stars", "Centipede"]}) - Examples - -------- +>>> import pyarrow.parquet as pq +>>> pq.write_table(table, 'example.parquet') - Generate an example PyArrow Table and write it to Parquet file: +Create a ``ParquetFile`` object from the Parquet file: - >>> import pyarrow as pa - >>> table = pa.table({'n_legs': [2, 2, 4, 4, 5, 100], - ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", - ... "Brittle stars", "Centipede"]}) +>>> parquet_file = pq.ParquetFile('example.parquet') - >>> import pyarrow.parquet as pq - >>> pq.write_table(table, 'example.parquet') +Read the data: + +>>> parquet_file.read() +pyarrow.Table +n_legs: int64 +animal: string +---- +n_legs: [[2,2,4,4,5,100]] +animal: [["Flamingo","Parrot","Dog","Horse","Brittle stars","Centipede"]] - Create a ``ParquetFile`` object from the Parquet file: +Create a ParquetFile object with "animal" column as DictionaryArray: - >>> parquet_file = pq.ParquetFile('example.parquet') +>>> parquet_file = pq.ParquetFile('example.parquet', +... read_dictionary=["animal"]) +>>> parquet_file.read() +pyarrow.Table +n_legs: int64 +animal: dictionary +---- +n_legs: [[2,2,4,4,5,100]] +animal: [ -- dictionary: +["Flamingo","Parrot",...,"Brittle stars","Centipede"] -- indices: +[0,1,2,3,4,5]] +""" - Read the data: - >>> parquet_file.read() - pyarrow.Table - n_legs: int64 - animal: string - ---- - n_legs: [[2,2,4,4,5,100]] - animal: [["Flamingo","Parrot","Dog","Horse","Brittle stars","Centipede"]] +class ParquetFile: + __doc__ = f""" +Reader interface for a single Parquet file. - Create a ParquetFile object with "animal" column as DictionaryArray: +Parameters +---------- +source : str, pathlib.Path, pyarrow.NativeFile, or file-like object + Readable source. For passing bytes or buffer-like file containing a + Parquet file, use pyarrow.BufferReader. +metadata : FileMetaData, default None + Use existing metadata object, rather than reading from file. +common_metadata : FileMetaData, default None + Will be used in reads for pandas schema metadata if not found in the + main file's metadata, no other uses at the moment. +filesystem : FileSystem, default None + If nothing passed, will be inferred based on path. + Path will try to be found in the local on-disk filesystem otherwise + it will be parsed as an URI to determine the filesystem. +{_read_docstring_common} - >>> parquet_file = pq.ParquetFile('example.parquet', - ... read_dictionary=["animal"]) - >>> parquet_file.read() - pyarrow.Table - n_legs: int64 - animal: dictionary - ---- - n_legs: [[2,2,4,4,5,100]] - animal: [ -- dictionary: - ["Flamingo","Parrot",...,"Brittle stars","Centipede"] -- indices: - [0,1,2,3,4,5]] - """ +Examples +-------- +{_parquet_file_example} +""" def __init__(self, source, *, metadata=None, common_metadata=None, read_dictionary=None, binary_type=None, list_type=None, @@ -1249,38 +1277,6 @@ def _get_pandas_index_columns(keyvalues): EXCLUDED_PARQUET_PATHS = {'_SUCCESS'} -_read_docstring_common = """\ -read_dictionary : list, default None - List of names or column paths (for nested types) to read directly - as DictionaryArray. Only supported for BYTE_ARRAY storage. To read - a flat column as dictionary-encoded pass the column name. For - nested types, you must pass the full column "path", which could be - something like level1.level2.list.item. Refer to the Parquet - file's schema to obtain the paths. -binary_type : pyarrow.DataType, default None - If given, Parquet binary columns will be read as this datatype. - This setting is ignored if a serialized Arrow schema is found in - the Parquet metadata. -list_type : subclass of pyarrow.DataType, default None - If given, non-MAP repeated columns will be read as an instance of - this datatype (either pyarrow.ListType or pyarrow.LargeListType). - This setting is ignored if a serialized Arrow schema is found in - the Parquet metadata. -memory_map : bool, default False - If the source is a file path, use a memory map to read file, which can - improve performance in some environments. -buffer_size : int, default 0 - If positive, perform read buffering when deserializing individual - column chunks. Otherwise IO calls are unbuffered. -partitioning : pyarrow.dataset.Partitioning or str or list of str, \ -default "hive" - The partitioning scheme for a partitioned dataset. The default of "hive" - assumes directory names with key=value pairs like "/year=2009/month=11". - In addition, a scheme like "/2009/11" is also supported, in which case - you need to specify the field names or a full schema. See the - ``pyarrow.dataset.partitioning()`` function for more details.""" - - _parquet_dataset_example = """\ Generate an example PyArrow Table and write it to a partitioned dataset: @@ -1343,41 +1339,7 @@ class ParquetDataset: {_DNF_filter_doc} {_read_docstring_common} -ignore_prefixes : list, optional - Files matching any of these prefixes will be ignored by the - discovery process. - This is matched to the basename of a path. - By default this is ['.', '_']. - Note that discovery happens only if a directory is passed as source. -pre_buffer : bool, default True - Coalesce and issue file reads in parallel to improve performance on - high-latency filesystems (e.g. S3, GCS). If True, Arrow will use a - background I/O thread pool. If using a filesystem layer that itself - performs readahead (e.g. fsspec's S3FS), disable readahead for best - results. Set to False if you want to prioritize minimal memory usage - over maximum speed. -coerce_int96_timestamp_unit : str, default None - Cast timestamps that are stored in INT96 format to a particular resolution - (e.g. 'ms'). Setting to None is equivalent to 'ns' and therefore INT96 - timestamps will be inferred as timestamps in nanoseconds. -decryption_properties : FileDecryptionProperties or None - File-level decryption properties. - The decryption properties can be created using - ``CryptoFactory.file_decryption_properties()``. -thrift_string_size_limit : int, default None - If not None, override the maximum total string size allocated - when decoding Thrift structures. The default limit should be - sufficient for most Parquet files. -thrift_container_size_limit : int, default None - If not None, override the maximum total size of containers allocated - when decoding Thrift structures. The default limit should be - sufficient for most Parquet files. -page_checksum_verification : bool, default False - If True, verify the page checksum for each page read from the file. -arrow_extensions_enabled : bool, default True - If True, read Parquet logical types as Arrow extension types where possible, - (e.g., read JSON as the canonical `arrow.json` extension type or UUID as - the canonical `arrow.uuid` extension type). +{_read_docstring_dataset} Examples -------- @@ -1726,8 +1688,8 @@ def partitioning(self): return self._dataset.partitioning -_read_table_docstring = """ -{0} +_read_table_docstring = f""" +{{0}} Parameters ---------- @@ -1747,7 +1709,7 @@ def partitioning(self): schema : Schema, optional Optionally provide the Schema for the parquet dataset, in which case it will not be inferred from the source. -{1} +{{1}} filesystem : FileSystem, default None If nothing passed, will be inferred based on path. Path will try to be found in the local on-disk filesystem otherwise @@ -1758,48 +1720,14 @@ def partitioning(self): exploited to avoid loading files at all if they contain no matching rows. Within-file level filtering and different partitioning schemes are supported. - {3} -ignore_prefixes : list, optional - Files matching any of these prefixes will be ignored by the - discovery process. - This is matched to the basename of a path. - By default this is ['.', '_']. - Note that discovery happens only if a directory is passed as source. -pre_buffer : bool, default True - Coalesce and issue file reads in parallel to improve performance on - high-latency filesystems (e.g. S3). If True, Arrow will use a - background I/O thread pool. If using a filesystem layer that itself - performs readahead (e.g. fsspec's S3FS), disable readahead for best - results. -coerce_int96_timestamp_unit : str, default None - Cast timestamps that are stored in INT96 format to a particular - resolution (e.g. 'ms'). Setting to None is equivalent to 'ns' - and therefore INT96 timestamps will be inferred as timestamps - in nanoseconds. -decryption_properties : FileDecryptionProperties or None - File-level decryption properties. - The decryption properties can be created using - ``CryptoFactory.file_decryption_properties()``. -thrift_string_size_limit : int, default None - If not None, override the maximum total string size allocated - when decoding Thrift structures. The default limit should be - sufficient for most Parquet files. -thrift_container_size_limit : int, default None - If not None, override the maximum total size of containers allocated - when decoding Thrift structures. The default limit should be - sufficient for most Parquet files. -page_checksum_verification : bool, default False - If True, verify the checksum for each page read from the file. -arrow_extensions_enabled : bool, default True - If True, read Parquet logical types as Arrow extension types where possible, - (e.g., read JSON as the canonical `arrow.json` extension type or UUID as - the canonical `arrow.uuid` extension type). + {{3}} +{_read_docstring_dataset} Returns ------- -{2} +{{2}} -{4} +{{4}} """ _read_table_example = """\