include/arrow/stl_iterator.h

0001 // Licensed to the Apache Software Foundation (ASF) under one
0002 // or more contributor license agreements.  See the NOTICE file
0003 // distributed with this work for additional information
0004 // regarding copyright ownership.  The ASF licenses this file
0005 // to you under the Apache License, Version 2.0 (the
0006 // "License"); you may not use this file except in compliance
0007 // with the License.  You may obtain a copy of the License at
0008 //
0009 //   http://www.apache.org/licenses/LICENSE-2.0
0010 //
0011 // Unless required by applicable law or agreed to in writing,
0012 // software distributed under the License is distributed on an
0013 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
0014 // KIND, either express or implied.  See the License for the
0015 // specific language governing permissions and limitations
0016 // under the License.
0017
0018 #pragma once
0019
0020 #include <cassert>
0021 #include <cstddef>
0022 #include <iterator>
0023 #include <optional>
0024 #include <utility>
0025
0026 #include "arrow/chunked_array.h"
0027 #include "arrow/type.h"
0028 #include "arrow/type_fwd.h"
0029 #include "arrow/type_traits.h"
0030 #include "arrow/util/macros.h"
0031
0032 namespace arrow {
0033 namespace stl {
0034
0035 namespace detail {
0036
0037 template <typename ArrayType>
0038 struct DefaultValueAccessor {
0039   using ValueType = decltype(std::declval<ArrayType>().GetView(0));
0040
0041   ValueType operator()(const ArrayType& array, int64_t index) {
0042     return array.GetView(index);
0043   }
0044 };
0045
0046 }  // namespace detail
0047
0048 template <typename ArrayType,
0049           typename ValueAccessor = detail::DefaultValueAccessor<ArrayType>>
0050 class ArrayIterator {
0051  public:
0052   using value_type = std::optional<typename ValueAccessor::ValueType>;
0053   using difference_type = int64_t;
0054   using pointer = value_type*;
0055   using reference = value_type&;
0056   using iterator_category = std::random_access_iterator_tag;
0057
0058   // Some algorithms need to default-construct an iterator
0059   ArrayIterator() : array_(NULLPTR), index_(0) {}
0060
0061   explicit ArrayIterator(const ArrayType& array, int64_t index = 0)
0062       : array_(&array), index_(index) {}
0063
0064   // Value access
0065   value_type operator*() const {
0066     assert(array_);
0067     return array_->IsNull(index_) ? value_type{} : array_->GetView(index_);
0068   }
0069
0070   value_type operator[](difference_type n) const {
0071     assert(array_);
0072     return array_->IsNull(index_ + n) ? value_type{} : array_->GetView(index_ + n);
0073   }
0074
0075   int64_t index() const { return index_; }
0076
0077   // Forward / backward
0078   ArrayIterator& operator++() {
0079     ++index_;
0080     return *this;
0081   }
0082   ArrayIterator& operator--() {
0083     --index_;
0084     return *this;
0085   }
0086   ArrayIterator operator++(int) {
0087     ArrayIterator tmp(*this);
0088     ++index_;
0089     return tmp;
0090   }
0091   ArrayIterator operator--(int) {
0092     ArrayIterator tmp(*this);
0093     --index_;
0094     return tmp;
0095   }
0096
0097   // Arithmetic
0098   difference_type operator-(const ArrayIterator& other) const {
0099     return index_ - other.index_;
0100   }
0101   ArrayIterator operator+(difference_type n) const {
0102     return ArrayIterator(*array_, index_ + n);
0103   }
0104   ArrayIterator operator-(difference_type n) const {
0105     return ArrayIterator(*array_, index_ - n);
0106   }
0107   friend inline ArrayIterator operator+(difference_type diff,
0108                                         const ArrayIterator& other) {
0109     return ArrayIterator(*other.array_, diff + other.index_);
0110   }
0111   friend inline ArrayIterator operator-(difference_type diff,
0112                                         const ArrayIterator& other) {
0113     return ArrayIterator(*other.array_, diff - other.index_);
0114   }
0115   ArrayIterator& operator+=(difference_type n) {
0116     index_ += n;
0117     return *this;
0118   }
0119   ArrayIterator& operator-=(difference_type n) {
0120     index_ -= n;
0121     return *this;
0122   }
0123
0124   // Comparisons
0125   bool operator==(const ArrayIterator& other) const { return index_ == other.index_; }
0126   bool operator!=(const ArrayIterator& other) const { return index_ != other.index_; }
0127   bool operator<(const ArrayIterator& other) const { return index_ < other.index_; }
0128   bool operator>(const ArrayIterator& other) const { return index_ > other.index_; }
0129   bool operator<=(const ArrayIterator& other) const { return index_ <= other.index_; }
0130   bool operator>=(const ArrayIterator& other) const { return index_ >= other.index_; }
0131
0132  private:
0133   const ArrayType* array_;
0134   int64_t index_;
0135 };
0136
0137 template <typename ArrayType,
0138           typename ValueAccessor = detail::DefaultValueAccessor<ArrayType>>
0139 class ChunkedArrayIterator {
0140  public:
0141   using value_type = std::optional<typename ValueAccessor::ValueType>;
0142   using difference_type = int64_t;
0143   using pointer = value_type*;
0144   using reference = value_type&;
0145   using iterator_category = std::random_access_iterator_tag;
0146
0147   // Some algorithms need to default-construct an iterator
0148   ChunkedArrayIterator() noexcept : chunked_array_(NULLPTR), index_(0) {}
0149
0150   explicit ChunkedArrayIterator(const ChunkedArray& chunked_array,
0151                                 int64_t index = 0) noexcept
0152       : chunked_array_(&chunked_array), index_(index) {}
0153
0154   // Value access
0155   value_type operator*() const {
0156     auto chunk_location = GetChunkLocation(index_);
0157     ArrayIterator<ArrayType> target_iterator{
0158         arrow::internal::checked_cast<const ArrayType&>(
0159             *chunked_array_->chunk(static_cast<int>(chunk_location.chunk_index)))};
0160     return target_iterator[chunk_location.index_in_chunk];
0161   }
0162
0163   value_type operator[](difference_type n) const { return *(*this + n); }
0164
0165   int64_t index() const { return index_; }
0166
0167   // Forward / backward
0168   ChunkedArrayIterator& operator++() {
0169     (*this) += 1;
0170     return *this;
0171   }
0172   ChunkedArrayIterator& operator--() {
0173     (*this) -= 1;
0174     return *this;
0175   }
0176
0177   ChunkedArrayIterator operator++(int) {
0178     ChunkedArrayIterator tmp(*this);
0179     ++*this;
0180     return tmp;
0181   }
0182   ChunkedArrayIterator operator--(int) {
0183     ChunkedArrayIterator tmp(*this);
0184     --*this;
0185     return tmp;
0186   }
0187
0188   // Arithmetic
0189   difference_type operator-(const ChunkedArrayIterator& other) const {
0190     return index_ - other.index_;
0191   }
0192   ChunkedArrayIterator operator+(difference_type n) const {
0193     assert(chunked_array_);
0194     return ChunkedArrayIterator(*chunked_array_, index_ + n);
0195   }
0196   ChunkedArrayIterator operator-(difference_type n) const {
0197     assert(chunked_array_);
0198     return ChunkedArrayIterator(*chunked_array_, index_ - n);
0199   }
0200   friend inline ChunkedArrayIterator operator+(difference_type diff,
0201                                                const ChunkedArrayIterator& other) {
0202     assert(other.chunked_array_);
0203     return ChunkedArrayIterator(*other.chunked_array_, diff + other.index_);
0204   }
0205   friend inline ChunkedArrayIterator operator-(difference_type diff,
0206                                                const ChunkedArrayIterator& other) {
0207     assert(other.chunked_array_);
0208     return ChunkedArrayIterator(*other.chunked_array_, diff - other.index_);
0209   }
0210   ChunkedArrayIterator& operator+=(difference_type n) {
0211     index_ += n;
0212     return *this;
0213   }
0214   ChunkedArrayIterator& operator-=(difference_type n) {
0215     (*this) += -n;
0216     return *this;
0217   }
0218
0219   // Comparisons
0220   bool operator==(const ChunkedArrayIterator& other) const {
0221     return index_ == other.index_;
0222   }
0223   bool operator!=(const ChunkedArrayIterator& other) const {
0224     return index_ != other.index_;
0225   }
0226   bool operator<(const ChunkedArrayIterator& other) const {
0227     return index_ < other.index_;
0228   }
0229   bool operator>(const ChunkedArrayIterator& other) const {
0230     return index_ > other.index_;
0231   }
0232   bool operator<=(const ChunkedArrayIterator& other) const {
0233     return index_ <= other.index_;
0234   }
0235   bool operator>=(const ChunkedArrayIterator& other) const {
0236     return index_ >= other.index_;
0237   }
0238
0239  private:
0240   arrow::ChunkLocation GetChunkLocation(int64_t index) const {
0241     assert(chunked_array_);
0242     return chunked_array_->chunk_resolver_.Resolve(index);
0243   }
0244
0245   const ChunkedArray* chunked_array_;
0246   int64_t index_;
0247 };
0248
0249 /// Return an iterator to the beginning of the chunked array
0250 template <typename Type, typename ArrayType = typename TypeTraits<Type>::ArrayType>
0251 ChunkedArrayIterator<ArrayType> Begin(const ChunkedArray& chunked_array) {
0252   return ChunkedArrayIterator<ArrayType>(chunked_array);
0253 }
0254
0255 /// Return an iterator to the end of the chunked array
0256 template <typename Type, typename ArrayType = typename TypeTraits<Type>::ArrayType>
0257 ChunkedArrayIterator<ArrayType> End(const ChunkedArray& chunked_array) {
0258   return ChunkedArrayIterator<ArrayType>(chunked_array, chunked_array.length());
0259 }
0260
0261 template <typename ArrayType>
0262 struct ChunkedArrayRange {
0263   const ChunkedArray* chunked_array;
0264
0265   ChunkedArrayIterator<ArrayType> begin() {
0266     return stl::ChunkedArrayIterator<ArrayType>(*chunked_array);
0267   }
0268   ChunkedArrayIterator<ArrayType> end() {
0269     return stl::ChunkedArrayIterator<ArrayType>(*chunked_array, chunked_array->length());
0270   }
0271 };
0272
0273 /// Return an iterable range over the chunked array
0274 template <typename Type, typename ArrayType = typename TypeTraits<Type>::ArrayType>
0275 ChunkedArrayRange<ArrayType> Iterate(const ChunkedArray& chunked_array) {
0276   return stl::ChunkedArrayRange<ArrayType>{&chunked_array};
0277 }
0278
0279 }  // namespace stl
0280 }  // namespace arrow
0281
0282 namespace std {
0283
0284 template <typename ArrayType>
0285 struct iterator_traits<::arrow::stl::ArrayIterator<ArrayType>> {
0286   using IteratorType = ::arrow::stl::ArrayIterator<ArrayType>;
0287   using difference_type = typename IteratorType::difference_type;
0288   using value_type = typename IteratorType::value_type;
0289   using pointer = typename IteratorType::pointer;
0290   using reference = typename IteratorType::reference;
0291   using iterator_category = typename IteratorType::iterator_category;
0292 };
0293
0294 template <typename ArrayType>
0295 struct iterator_traits<::arrow::stl::ChunkedArrayIterator<ArrayType>> {
0296   using IteratorType = ::arrow::stl::ChunkedArrayIterator<ArrayType>;
0297   using difference_type = typename IteratorType::difference_type;
0298   using value_type = typename IteratorType::value_type;
0299   using pointer = typename IteratorType::pointer;
0300   using reference = typename IteratorType::reference;
0301   using iterator_category = typename IteratorType::iterator_category;
0302 };
0303
0304 }  // namespace std