You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "lidavidm (via GitHub)" <gi...@apache.org> on 2023/02/16 20:49:10 UTC
[GitHub] [arrow] lidavidm commented on a diff in pull request #33641: GH-32104: [C++] Add support for Run-End encoded data to Arrow
lidavidm commented on code in PR #33641:
URL: https://github.com/apache/arrow/pull/33641#discussion_r1108960323
##########
cpp/src/arrow/type.h:
##########
@@ -1223,6 +1223,35 @@ class ARROW_EXPORT DenseUnionType : public UnionType {
std::string name() const override { return "dense_union"; }
};
+/// \brief Type class for run-end encoded data
+class ARROW_EXPORT RunEndEncodedType : public NestedType {
+ public:
+ static constexpr Type::type type_id = Type::RUN_END_ENCODED;
+
+ static constexpr const char* type_name() { return "run_end_encoded"; }
+
+ explicit RunEndEncodedType(std::shared_ptr<DataType> run_end_type,
+ std::shared_ptr<DataType> value_type);
+
+ DataTypeLayout layout() const override {
+ // NOTE(felipecrv): his can be removed once we ensure existing code does not
+ // assume the existence of at least one buffer.
Review Comment:
Should we file a task for it (and make this a `TODO(apache/arrow#NNN)`)?
##########
cpp/src/arrow/util/ree_util.h:
##########
@@ -0,0 +1,351 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <cstdint>
+
+#include "arrow/array/data.h"
+#include "arrow/type_traits.h"
+#include "arrow/util/macros.h"
+
+namespace arrow {
+namespace ree_util {
+
+/// \brief Get the child array holding the run ends from an REE array
+inline const ArraySpan& RunEndsArray(const ArraySpan& span) { return span.child_data[0]; }
+
+/// \brief Get the child array holding the data values from an REE array
+inline const ArraySpan& ValuesArray(const ArraySpan& span) { return span.child_data[1]; }
+
+/// \brief Get a pointer to run ends values of an REE array
+template <typename RunEndsType>
+const RunEndsType* RunEnds(const ArraySpan& span) {
+ assert(RunEndsArray(span).type->id() == CTypeTraits<RunEndsType>::ArrowType::type_id);
+ return RunEndsArray(span).GetValues<RunEndsType>(1);
+}
+
+namespace internal {
+
+/// \brief Uses binary-search to find the physical offset given a logical offset
+/// and run-end values
+///
+/// \return the physical offset or run_ends_size if the physical offset is not
+/// found in run_ends
+template <typename RunEndsType>
+int64_t FindPhysicalIndex(const RunEndsType* run_ends, int64_t run_ends_size, int64_t i,
+ int64_t absolute_offset) {
+ auto it = std::upper_bound(run_ends, run_ends + run_ends_size, absolute_offset + i);
+ int64_t result = std::distance(run_ends, it);
+ assert(result <= run_ends_size);
+ return result;
+}
+
+/// \brief Uses binary-search to calculate the number of physical values (and
+/// run-ends) necessary to represent the logical range of values from
+/// offset to length
+template <typename RunEndsType>
+int64_t FindPhysicalLength(int64_t length, int64_t offset, const RunEndsType* run_ends,
+ int64_t run_ends_size) {
Review Comment:
nit: it's a bit confusing that two similar functions have all their parameters in a different order
##########
cpp/src/arrow/scalar.cc:
##########
@@ -583,6 +608,17 @@ Result<std::shared_ptr<Scalar>> StructScalar::field(FieldRef ref) const {
}
}
+RunEndEncodedScalar::RunEndEncodedScalar(std::shared_ptr<Scalar> value,
+ std::shared_ptr<DataType> type)
+ : Scalar{std::move(type), value->is_valid}, value{std::move(value)} {}
Review Comment:
should we have a DCHECK here (e.g. that the value type is of the right type, that `type` is actually an REE type)?
##########
cpp/src/arrow/testing/json_internal.cc:
##########
@@ -437,6 +437,10 @@ class SchemaWriter {
Status Visit(const DictionaryType& type) { return VisitType(*type.value_type()); }
+ Status Visit(const RunEndEncodedType& type) {
+ return Status::NotImplemented(type.name());
Review Comment:
Would be good to put up a follow-up to implement this since this is used for other tests.
##########
cpp/src/arrow/array/array_base.cc:
##########
@@ -165,7 +175,7 @@ struct ScalarFromArraySlotImpl {
array_.length());
}
- if (array_.IsNull(index_)) {
+ if (array_.type()->id() != Type::RUN_END_ENCODED && array_.IsNull(index_)) {
Review Comment:
Hmm, I suppose this is unavoidable.
##########
cpp/src/arrow/array/array_test.cc:
##########
@@ -375,8 +375,9 @@ TEST_F(TestArray, TestMakeArrayOfNull) {
ASSERT_EQ(array->type(), type);
ASSERT_OK(array->ValidateFull());
ASSERT_EQ(array->length(), length);
- if (is_union(type->id())) {
- // For unions, MakeArrayOfNull places the nulls in the children
+ if (is_union(type->id()) || type->id() == Type::RUN_END_ENCODED) {
+ // For unions and run-end encoded, MakeArrayOfNull places the nulls in the
+ // children
Review Comment:
Hmm.
Since this pattern repeats a bit: factor out a helper predicate to test it? It seems we now have a class of 'no top level null' types (union, REE)
##########
cpp/src/arrow/util/ree_util_test.cc:
##########
@@ -0,0 +1,261 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gtest/gtest.h>
+
+#include "arrow/array.h"
+#include "arrow/builder.h"
+#include "arrow/compute/api_vector.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/type_traits.h"
+#include "arrow/util/logging.h"
+#include "arrow/util/ree_util.h"
+
+namespace arrow {
+namespace ree_util {
+
+template <typename RunEndsType>
+struct ReeUtilTest : public ::testing::Test {
+ // Re-implementation of FindPhysicalIndex that uses trivial linear-search
+ // instead of the more efficient implementation.
+ int64_t FindPhysicalIndexTestImpl(const RunEndsType* run_ends, int64_t run_ends_size,
+ int64_t i, int64_t absolute_offset = 0) {
+ for (int64_t j = 0; j < run_ends_size; j++) {
+ if (absolute_offset + i < run_ends[j]) {
+ return j;
+ }
+ }
+ return run_ends_size;
+ }
+};
+TYPED_TEST_SUITE_P(ReeUtilTest);
+
+TYPED_TEST_P(ReeUtilTest, PhysicalOffset) {
+ using RE = TypeParam; // Run-end type
+ const RE run_ends1[] = {1};
+ ASSERT_EQ(internal::FindPhysicalIndex(run_ends1, 1, 0, 0), 0);
+ const RE run_ends124[] = {1, 2, 4};
+ ASSERT_EQ(internal::FindPhysicalIndex(run_ends124, 3, 0, 0), 0);
+ ASSERT_EQ(internal::FindPhysicalIndex(run_ends124, 3, 1, 0), 1);
+ ASSERT_EQ(internal::FindPhysicalIndex(run_ends124, 3, 2, 0), 2);
+ const RE run_ends234[] = {2, 3, 4};
+ ASSERT_EQ(internal::FindPhysicalIndex(run_ends234, 3, 0, 0), 0);
+ ASSERT_EQ(internal::FindPhysicalIndex(run_ends234, 3, 1, 0), 0);
+ ASSERT_EQ(internal::FindPhysicalIndex(run_ends234, 3, 2, 0), 1);
+ ASSERT_EQ(internal::FindPhysicalIndex(run_ends234, 3, 3, 0), 2);
+ const RE run_ends246[] = {2, 4, 6};
+ ASSERT_EQ(internal::FindPhysicalIndex(run_ends246, 3, 3, 0), 1);
+
+ // Out-of-range logical offset should return run_ends size
Review Comment:
Should we assert that the offset is nonnegative?
##########
cpp/src/arrow/array/concatenate.cc:
##########
@@ -436,6 +438,28 @@ class ConcatenateImpl {
return Status::OK();
}
+ Status Visit(const RunEndEncodedType& type) {
+ int64_t physical_length = 0;
+ for (const auto& input : in_) {
+ if (internal::AddWithOverflow(physical_length,
+ ree_util::FindPhysicalLength(ArraySpan(*input)),
+ &physical_length) ||
+ // Make sure we can safely downcast to int32_t
Review Comment:
Hmm, why do we have to downcast to int32_t here? Shouldn't it be based on the offsets type?
##########
cpp/src/arrow/array/builder_run_end.cc:
##########
@@ -0,0 +1,331 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/array/builder_run_end.h"
+#include "arrow/array/builder_primitive.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "arrow/scalar.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/logging.h"
+#include "arrow/util/ree_util.h"
+
+namespace arrow {
+namespace internal {
+
+RunCompressorBuilder::RunCompressorBuilder(MemoryPool* pool,
+ std::shared_ptr<ArrayBuilder> inner_builder,
+ std::shared_ptr<DataType> type)
+ : ArrayBuilder(pool), inner_builder_(std::move(inner_builder)) {}
+
+RunCompressorBuilder::~RunCompressorBuilder() = default;
+
+void RunCompressorBuilder::Reset() {
+ current_run_length_ = 0;
+ current_value_.reset();
+ inner_builder_->Reset();
+ UpdateDimensions();
+}
+
+Status RunCompressorBuilder::ResizePhysical(int64_t capacity) {
+ RETURN_NOT_OK(inner_builder_->Resize(capacity));
+ UpdateDimensions();
+ return Status::OK();
+}
+
+Status RunCompressorBuilder::AppendNulls(int64_t length) {
+ if (ARROW_PREDICT_FALSE(length == 0)) {
+ return Status::OK();
+ }
+ if (ARROW_PREDICT_FALSE(current_run_length_ == 0)) {
+ // Open a new NULL run
+ DCHECK_EQ(current_value_, NULLPTR);
+ current_run_length_ = length;
+ } else if (current_value_ == NULLPTR) {
+ // Extend the currently open NULL run
+ current_run_length_ += length;
+ } else {
+ // Close then non-NULL run
+ ARROW_RETURN_NOT_OK(WillCloseRun(current_value_, current_run_length_));
+ ARROW_RETURN_NOT_OK(inner_builder_->AppendScalar(*current_value_));
+ UpdateDimensions();
+ // Open a new NULL run
+ current_value_.reset();
+ current_run_length_ = length;
+ }
+ return Status::OK();
+}
+
+Status RunCompressorBuilder::AppendEmptyValues(int64_t length) {
+ return Status::NotImplemented("Append empty values to a run-compressed array.");
+}
+
+Status RunCompressorBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) {
+ if (ARROW_PREDICT_FALSE(n_repeats == 0)) {
+ return Status::OK();
+ }
+ if (ARROW_PREDICT_FALSE(current_run_length_ == 0)) {
+ // Open a new run
+ current_value_ = scalar.is_valid ? scalar.shared_from_this() : NULLPTR;
+ current_run_length_ = n_repeats;
+ } else if ((current_value_ == NULLPTR && !scalar.is_valid) ||
+ (current_value_ != NULLPTR && current_value_->Equals(scalar))) {
+ // Extend the currently open run
+ current_run_length_ += n_repeats;
+ } else {
+ // Close the current run
+ ARROW_RETURN_NOT_OK(WillCloseRun(current_value_, current_run_length_));
+ ARROW_RETURN_NOT_OK(current_value_ ? inner_builder_->AppendScalar(*current_value_)
+ : inner_builder_->AppendNull());
+ UpdateDimensions();
+ // Open a new run
+ current_value_ = scalar.is_valid ? scalar.shared_from_this() : NULLPTR;
+ current_run_length_ = n_repeats;
+ }
+ return Status::OK();
+}
+
+Status RunCompressorBuilder::AppendScalars(const ScalarVector& scalars) {
+ if (scalars.empty()) {
+ return Status::OK();
+ }
+ RETURN_NOT_OK(ArrayBuilder::AppendScalars(scalars));
+ UpdateDimensions();
+ return Status::OK();
+}
+
+Status RunCompressorBuilder::AppendRunCompressedArraySlice(
+ const ArraySpan& run_compressed_array, int64_t offset, int64_t length) {
+ DCHECK(!has_open_run());
+ RETURN_NOT_OK(inner_builder_->AppendArraySlice(run_compressed_array, offset, length));
+ UpdateDimensions();
+ return Status::OK();
+}
+
+Status RunCompressorBuilder::FinishCurrentRun() {
+ if (current_run_length_ > 0) {
+ // Close the current run
+ ARROW_RETURN_NOT_OK(WillCloseRun(current_value_, current_run_length_));
+ ARROW_RETURN_NOT_OK(current_value_ ? inner_builder_->AppendScalar(*current_value_)
+ : inner_builder_->AppendNull());
+ UpdateDimensions();
+ // Clear the current run
+ current_value_.reset();
+ current_run_length_ = 0;
+ }
+ return Status::OK();
+}
+
+Status RunCompressorBuilder::FinishInternal(std::shared_ptr<ArrayData>* out) {
+ ARROW_RETURN_NOT_OK(FinishCurrentRun());
+ return inner_builder_->FinishInternal(out);
+}
+
+} // namespace internal
+
+// ----------------------------------------------------------------------
+// RunEndEncodedBuilder
+
+RunEndEncodedBuilder::ValueRunBuilder::ValueRunBuilder(
+ MemoryPool* pool, const std::shared_ptr<ArrayBuilder>& value_builder,
+ const std::shared_ptr<DataType>& value_type, RunEndEncodedBuilder& ree_builder)
+ : RunCompressorBuilder(pool, std::move(value_builder), std::move(value_type)),
+ ree_builder_(ree_builder) {}
+
+RunEndEncodedBuilder::RunEndEncodedBuilder(
+ MemoryPool* pool, const std::shared_ptr<ArrayBuilder>& run_end_builder,
+ const std::shared_ptr<ArrayBuilder>& value_builder, std::shared_ptr<DataType> type)
+ : ArrayBuilder(pool), type_(internal::checked_pointer_cast<RunEndEncodedType>(type)) {
+ auto value_run_builder =
+ std::make_shared<ValueRunBuilder>(pool, value_builder, type_->value_type(), *this);
+ value_run_builder_ = value_run_builder.get();
+ children_ = {run_end_builder, std::move(value_run_builder)};
+ UpdateDimensions(0, 0);
+ null_count_ = 0;
+}
+
+Status RunEndEncodedBuilder::ResizePhysical(int64_t capacity) {
+ RETURN_NOT_OK(value_run_builder_->ResizePhysical(capacity));
+ RETURN_NOT_OK(run_end_builder().Resize(capacity));
+ UpdateDimensions(committed_logical_length_, 0);
+ return Status::OK();
+}
+
+void RunEndEncodedBuilder::Reset() {
+ value_run_builder_->Reset();
+ run_end_builder().Reset();
+ UpdateDimensions(0, 0);
+}
+
+Status RunEndEncodedBuilder::AppendNulls(int64_t length) {
+ RETURN_NOT_OK(value_run_builder_->AppendNulls(length));
+ UpdateDimensions(committed_logical_length_, value_run_builder_->open_run_length());
+ return Status::OK();
+}
+
+Status RunEndEncodedBuilder::AppendEmptyValues(int64_t length) {
+ return Status::NotImplemented("Append empty values to run-end encoded array.");
+}
+
+Status RunEndEncodedBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) {
+ if (scalar.type->id() == Type::RUN_END_ENCODED) {
+ return AppendScalar(*internal::checked_cast<const RunEndEncodedScalar&>(scalar).value,
+ n_repeats);
+ }
+ RETURN_NOT_OK(value_run_builder_->AppendScalar(scalar, n_repeats));
+ UpdateDimensions(committed_logical_length_, value_run_builder_->open_run_length());
+ return Status::OK();
+}
+
+Status RunEndEncodedBuilder::AppendScalars(const ScalarVector& scalars) {
+ RETURN_NOT_OK(this->ArrayBuilder::AppendScalars(scalars));
+ UpdateDimensions(committed_logical_length_, value_run_builder_->open_run_length());
+ return Status::OK();
+}
+
+template <typename RunEndsType>
+Status RunEndEncodedBuilder::DoAppendArray(const ArraySpan& to_append) {
+ DCHECK_GT(to_append.length, 0);
+ DCHECK(!value_run_builder_->has_open_run());
+
+ ree_util::RunEndEncodedArraySpan<RunEndsType> ree_span(to_append);
+ const int64_t physical_offset = ree_span.PhysicalIndex(0);
+ const int64_t physical_length =
+ ree_span.PhysicalIndex(ree_span.length() - 1) + 1 - physical_offset;
+
+ RETURN_NOT_OK(ReservePhysical(physical_length));
+
+ // Append all the run ends from to_append
+ const auto end = ree_span.end();
+ for (auto it = ree_span.iterator(0, physical_offset); it != end; ++it) {
+ const int64_t run_end = committed_logical_length_ + it.run_length();
+ RETURN_NOT_OK(DoAppendRunEnd<RunEndsType>(run_end));
+ UpdateDimensions(run_end, 0);
+ }
+
+ // Append all the values directly
+ RETURN_NOT_OK(value_run_builder_->AppendRunCompressedArraySlice(
+ ree_util::ValuesArray(to_append), physical_offset, physical_length));
+
+ return Status::OK();
+}
+
+Status RunEndEncodedBuilder::AppendArraySlice(const ArraySpan& array, int64_t offset,
+ int64_t length) {
+ ARROW_DCHECK(offset + length <= array.length);
+ ARROW_DCHECK(array.type->Equals(type_));
+
+ // Ensure any open run is closed before appending the array slice.
+ RETURN_NOT_OK(value_run_builder_->FinishCurrentRun());
+
+ if (length == 0) {
+ return Status::OK();
+ }
+
+ ArraySpan to_append = array;
+ to_append.SetSlice(array.offset + offset, length);
+
+ switch (type_->run_end_type()->id()) {
+ case Type::INT16:
+ RETURN_NOT_OK(DoAppendArray<int16_t>(to_append));
+ break;
+ case Type::INT32:
+ RETURN_NOT_OK(DoAppendArray<int32_t>(to_append));
+ break;
+ case Type::INT64:
+ RETURN_NOT_OK(DoAppendArray<int64_t>(to_append));
+ break;
+ default:
+ return Status::Invalid("Invalid type for run ends array: ", type_->run_end_type());
+ }
+
+ return Status::OK();
+}
+
+std::shared_ptr<DataType> RunEndEncodedBuilder::type() const { return type_; }
+
+Status RunEndEncodedBuilder::FinishInternal(std::shared_ptr<ArrayData>* out) {
+ // Finish the values array before so we can close the current run and append
+ // the last run-end.
+ std::shared_ptr<ArrayData> values_data;
+ RETURN_NOT_OK(value_run_builder_->FinishInternal(&values_data));
+ auto values_array = MakeArray(values_data);
+
+ ARROW_ASSIGN_OR_RAISE(auto run_ends_array, run_end_builder().Finish());
+
+ ARROW_ASSIGN_OR_RAISE(auto ree_array,
+ RunEndEncodedArray::Make(length_, run_ends_array, values_array));
+ *out = std::move(ree_array->data());
+ return Status::OK();
+}
+
+Status RunEndEncodedBuilder::FinishCurrentRun() {
+ RETURN_NOT_OK(value_run_builder_->FinishCurrentRun());
+ UpdateDimensions(length_, 0);
+ return Status::OK();
+}
+
+template <typename RunEndsType>
+Status RunEndEncodedBuilder::DoAppendRunEnd(int64_t run_end) {
+ constexpr auto max = std::numeric_limits<RunEndsType>::max();
+ if (ARROW_PREDICT_FALSE(run_end > max)) {
+ return Status::Invalid("Run end value must fit on run ends type but ", run_end, " > ",
+ max, ".");
+ }
+ return internal::checked_cast<typename CTypeTraits<RunEndsType>::BuilderType*>(
+ children_[0].get())
+ ->Append(static_cast<RunEndsType>(run_end));
+}
+
+Status RunEndEncodedBuilder::AppendRunEnd(int64_t run_end) {
+ switch (type_->run_end_type()->id()) {
+ case Type::INT16:
+ RETURN_NOT_OK(DoAppendRunEnd<int16_t>(run_end));
+ break;
+ case Type::INT32:
+ RETURN_NOT_OK(DoAppendRunEnd<int32_t>(run_end));
+ break;
+ case Type::INT64:
+ RETURN_NOT_OK(DoAppendRunEnd<int64_t>(run_end));
+ break;
+ default:
+ return Status::Invalid("Invalid type for run ends array: ", type_->run_end_type());
+ }
+ return Status::OK();
+}
+
+Status RunEndEncodedBuilder::CloseRun(const std::shared_ptr<const Scalar>& value,
+ int64_t run_length) {
+ // TODO(felipecrv): gracefully fragment runs bigger than INT32_MAX
+ if (ARROW_PREDICT_FALSE(run_length > std::numeric_limits<int32_t>::max())) {
+ return Status::Invalid(
+ "Run-length of run-encoded arrays must fit in a 32-bit signed integer.");
+ }
+ const int64_t run_end = committed_logical_length_ + run_length;
Review Comment:
I suppose, check for overflow?
##########
cpp/src/arrow/array/builder_run_end.cc:
##########
@@ -0,0 +1,331 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/array/builder_run_end.h"
+#include "arrow/array/builder_primitive.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "arrow/scalar.h"
+#include "arrow/util/checked_cast.h"
+#include "arrow/util/logging.h"
+#include "arrow/util/ree_util.h"
+
+namespace arrow {
+namespace internal {
+
+RunCompressorBuilder::RunCompressorBuilder(MemoryPool* pool,
+ std::shared_ptr<ArrayBuilder> inner_builder,
+ std::shared_ptr<DataType> type)
+ : ArrayBuilder(pool), inner_builder_(std::move(inner_builder)) {}
+
+RunCompressorBuilder::~RunCompressorBuilder() = default;
+
+void RunCompressorBuilder::Reset() {
+ current_run_length_ = 0;
+ current_value_.reset();
+ inner_builder_->Reset();
+ UpdateDimensions();
+}
+
+Status RunCompressorBuilder::ResizePhysical(int64_t capacity) {
+ RETURN_NOT_OK(inner_builder_->Resize(capacity));
+ UpdateDimensions();
+ return Status::OK();
+}
+
+Status RunCompressorBuilder::AppendNulls(int64_t length) {
+ if (ARROW_PREDICT_FALSE(length == 0)) {
+ return Status::OK();
+ }
+ if (ARROW_PREDICT_FALSE(current_run_length_ == 0)) {
+ // Open a new NULL run
+ DCHECK_EQ(current_value_, NULLPTR);
+ current_run_length_ = length;
+ } else if (current_value_ == NULLPTR) {
+ // Extend the currently open NULL run
+ current_run_length_ += length;
+ } else {
+ // Close then non-NULL run
+ ARROW_RETURN_NOT_OK(WillCloseRun(current_value_, current_run_length_));
+ ARROW_RETURN_NOT_OK(inner_builder_->AppendScalar(*current_value_));
+ UpdateDimensions();
+ // Open a new NULL run
+ current_value_.reset();
+ current_run_length_ = length;
+ }
+ return Status::OK();
+}
+
+Status RunCompressorBuilder::AppendEmptyValues(int64_t length) {
+ return Status::NotImplemented("Append empty values to a run-compressed array.");
+}
+
+Status RunCompressorBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) {
+ if (ARROW_PREDICT_FALSE(n_repeats == 0)) {
+ return Status::OK();
+ }
+ if (ARROW_PREDICT_FALSE(current_run_length_ == 0)) {
+ // Open a new run
+ current_value_ = scalar.is_valid ? scalar.shared_from_this() : NULLPTR;
+ current_run_length_ = n_repeats;
+ } else if ((current_value_ == NULLPTR && !scalar.is_valid) ||
+ (current_value_ != NULLPTR && current_value_->Equals(scalar))) {
+ // Extend the currently open run
+ current_run_length_ += n_repeats;
+ } else {
+ // Close the current run
+ ARROW_RETURN_NOT_OK(WillCloseRun(current_value_, current_run_length_));
+ ARROW_RETURN_NOT_OK(current_value_ ? inner_builder_->AppendScalar(*current_value_)
+ : inner_builder_->AppendNull());
+ UpdateDimensions();
+ // Open a new run
+ current_value_ = scalar.is_valid ? scalar.shared_from_this() : NULLPTR;
+ current_run_length_ = n_repeats;
+ }
+ return Status::OK();
+}
+
+Status RunCompressorBuilder::AppendScalars(const ScalarVector& scalars) {
+ if (scalars.empty()) {
+ return Status::OK();
+ }
+ RETURN_NOT_OK(ArrayBuilder::AppendScalars(scalars));
+ UpdateDimensions();
+ return Status::OK();
+}
+
+Status RunCompressorBuilder::AppendRunCompressedArraySlice(
+ const ArraySpan& run_compressed_array, int64_t offset, int64_t length) {
+ DCHECK(!has_open_run());
+ RETURN_NOT_OK(inner_builder_->AppendArraySlice(run_compressed_array, offset, length));
+ UpdateDimensions();
+ return Status::OK();
+}
+
+Status RunCompressorBuilder::FinishCurrentRun() {
+ if (current_run_length_ > 0) {
+ // Close the current run
+ ARROW_RETURN_NOT_OK(WillCloseRun(current_value_, current_run_length_));
+ ARROW_RETURN_NOT_OK(current_value_ ? inner_builder_->AppendScalar(*current_value_)
+ : inner_builder_->AppendNull());
+ UpdateDimensions();
+ // Clear the current run
+ current_value_.reset();
+ current_run_length_ = 0;
+ }
+ return Status::OK();
+}
+
+Status RunCompressorBuilder::FinishInternal(std::shared_ptr<ArrayData>* out) {
+ ARROW_RETURN_NOT_OK(FinishCurrentRun());
+ return inner_builder_->FinishInternal(out);
+}
+
+} // namespace internal
+
+// ----------------------------------------------------------------------
+// RunEndEncodedBuilder
+
+RunEndEncodedBuilder::ValueRunBuilder::ValueRunBuilder(
+ MemoryPool* pool, const std::shared_ptr<ArrayBuilder>& value_builder,
+ const std::shared_ptr<DataType>& value_type, RunEndEncodedBuilder& ree_builder)
+ : RunCompressorBuilder(pool, std::move(value_builder), std::move(value_type)),
+ ree_builder_(ree_builder) {}
+
+RunEndEncodedBuilder::RunEndEncodedBuilder(
+ MemoryPool* pool, const std::shared_ptr<ArrayBuilder>& run_end_builder,
+ const std::shared_ptr<ArrayBuilder>& value_builder, std::shared_ptr<DataType> type)
+ : ArrayBuilder(pool), type_(internal::checked_pointer_cast<RunEndEncodedType>(type)) {
+ auto value_run_builder =
+ std::make_shared<ValueRunBuilder>(pool, value_builder, type_->value_type(), *this);
+ value_run_builder_ = value_run_builder.get();
+ children_ = {run_end_builder, std::move(value_run_builder)};
+ UpdateDimensions(0, 0);
+ null_count_ = 0;
+}
+
+Status RunEndEncodedBuilder::ResizePhysical(int64_t capacity) {
+ RETURN_NOT_OK(value_run_builder_->ResizePhysical(capacity));
+ RETURN_NOT_OK(run_end_builder().Resize(capacity));
+ UpdateDimensions(committed_logical_length_, 0);
+ return Status::OK();
+}
+
+void RunEndEncodedBuilder::Reset() {
+ value_run_builder_->Reset();
+ run_end_builder().Reset();
+ UpdateDimensions(0, 0);
+}
+
+Status RunEndEncodedBuilder::AppendNulls(int64_t length) {
+ RETURN_NOT_OK(value_run_builder_->AppendNulls(length));
+ UpdateDimensions(committed_logical_length_, value_run_builder_->open_run_length());
+ return Status::OK();
+}
+
+Status RunEndEncodedBuilder::AppendEmptyValues(int64_t length) {
+ return Status::NotImplemented("Append empty values to run-end encoded array.");
+}
+
+Status RunEndEncodedBuilder::AppendScalar(const Scalar& scalar, int64_t n_repeats) {
+ if (scalar.type->id() == Type::RUN_END_ENCODED) {
+ return AppendScalar(*internal::checked_cast<const RunEndEncodedScalar&>(scalar).value,
+ n_repeats);
+ }
+ RETURN_NOT_OK(value_run_builder_->AppendScalar(scalar, n_repeats));
+ UpdateDimensions(committed_logical_length_, value_run_builder_->open_run_length());
+ return Status::OK();
+}
+
+Status RunEndEncodedBuilder::AppendScalars(const ScalarVector& scalars) {
+ RETURN_NOT_OK(this->ArrayBuilder::AppendScalars(scalars));
+ UpdateDimensions(committed_logical_length_, value_run_builder_->open_run_length());
+ return Status::OK();
+}
+
+template <typename RunEndsType>
+Status RunEndEncodedBuilder::DoAppendArray(const ArraySpan& to_append) {
+ DCHECK_GT(to_append.length, 0);
+ DCHECK(!value_run_builder_->has_open_run());
+
+ ree_util::RunEndEncodedArraySpan<RunEndsType> ree_span(to_append);
+ const int64_t physical_offset = ree_span.PhysicalIndex(0);
+ const int64_t physical_length =
+ ree_span.PhysicalIndex(ree_span.length() - 1) + 1 - physical_offset;
+
+ RETURN_NOT_OK(ReservePhysical(physical_length));
+
+ // Append all the run ends from to_append
+ const auto end = ree_span.end();
+ for (auto it = ree_span.iterator(0, physical_offset); it != end; ++it) {
+ const int64_t run_end = committed_logical_length_ + it.run_length();
+ RETURN_NOT_OK(DoAppendRunEnd<RunEndsType>(run_end));
+ UpdateDimensions(run_end, 0);
+ }
+
+ // Append all the values directly
+ RETURN_NOT_OK(value_run_builder_->AppendRunCompressedArraySlice(
+ ree_util::ValuesArray(to_append), physical_offset, physical_length));
+
+ return Status::OK();
+}
+
+Status RunEndEncodedBuilder::AppendArraySlice(const ArraySpan& array, int64_t offset,
+ int64_t length) {
+ ARROW_DCHECK(offset + length <= array.length);
+ ARROW_DCHECK(array.type->Equals(type_));
+
+ // Ensure any open run is closed before appending the array slice.
+ RETURN_NOT_OK(value_run_builder_->FinishCurrentRun());
+
+ if (length == 0) {
+ return Status::OK();
+ }
+
+ ArraySpan to_append = array;
+ to_append.SetSlice(array.offset + offset, length);
+
+ switch (type_->run_end_type()->id()) {
+ case Type::INT16:
+ RETURN_NOT_OK(DoAppendArray<int16_t>(to_append));
+ break;
+ case Type::INT32:
+ RETURN_NOT_OK(DoAppendArray<int32_t>(to_append));
+ break;
+ case Type::INT64:
+ RETURN_NOT_OK(DoAppendArray<int64_t>(to_append));
+ break;
+ default:
+ return Status::Invalid("Invalid type for run ends array: ", type_->run_end_type());
+ }
+
+ return Status::OK();
+}
+
+std::shared_ptr<DataType> RunEndEncodedBuilder::type() const { return type_; }
+
+Status RunEndEncodedBuilder::FinishInternal(std::shared_ptr<ArrayData>* out) {
+ // Finish the values array before so we can close the current run and append
+ // the last run-end.
+ std::shared_ptr<ArrayData> values_data;
+ RETURN_NOT_OK(value_run_builder_->FinishInternal(&values_data));
+ auto values_array = MakeArray(values_data);
+
+ ARROW_ASSIGN_OR_RAISE(auto run_ends_array, run_end_builder().Finish());
+
+ ARROW_ASSIGN_OR_RAISE(auto ree_array,
+ RunEndEncodedArray::Make(length_, run_ends_array, values_array));
+ *out = std::move(ree_array->data());
+ return Status::OK();
+}
+
+Status RunEndEncodedBuilder::FinishCurrentRun() {
+ RETURN_NOT_OK(value_run_builder_->FinishCurrentRun());
+ UpdateDimensions(length_, 0);
+ return Status::OK();
+}
+
+template <typename RunEndsType>
+Status RunEndEncodedBuilder::DoAppendRunEnd(int64_t run_end) {
+ constexpr auto max = std::numeric_limits<RunEndsType>::max();
+ if (ARROW_PREDICT_FALSE(run_end > max)) {
+ return Status::Invalid("Run end value must fit on run ends type but ", run_end, " > ",
+ max, ".");
+ }
+ return internal::checked_cast<typename CTypeTraits<RunEndsType>::BuilderType*>(
+ children_[0].get())
+ ->Append(static_cast<RunEndsType>(run_end));
+}
+
+Status RunEndEncodedBuilder::AppendRunEnd(int64_t run_end) {
+ switch (type_->run_end_type()->id()) {
+ case Type::INT16:
+ RETURN_NOT_OK(DoAppendRunEnd<int16_t>(run_end));
+ break;
+ case Type::INT32:
+ RETURN_NOT_OK(DoAppendRunEnd<int32_t>(run_end));
+ break;
+ case Type::INT64:
+ RETURN_NOT_OK(DoAppendRunEnd<int64_t>(run_end));
+ break;
+ default:
+ return Status::Invalid("Invalid type for run ends array: ", type_->run_end_type());
+ }
+ return Status::OK();
+}
+
+Status RunEndEncodedBuilder::CloseRun(const std::shared_ptr<const Scalar>& value,
+ int64_t run_length) {
+ // TODO(felipecrv): gracefully fragment runs bigger than INT32_MAX
+ if (ARROW_PREDICT_FALSE(run_length > std::numeric_limits<int32_t>::max())) {
Review Comment:
Shouldn't this depend on the offset type?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org