include/arrow/table_builder.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 <cstdint>
0021 #include <memory>
0022 #include <vector>
0023
0024 #include "arrow/array/builder_base.h"
0025 #include "arrow/status.h"
0026 #include "arrow/type.h"
0027 #include "arrow/util/checked_cast.h"
0028 #include "arrow/util/macros.h"
0029 #include "arrow/util/visibility.h"
0030
0031 namespace arrow {
0032
0033 class MemoryPool;
0034 class RecordBatch;
0035
0036 /// \class RecordBatchBuilder
0037 /// \brief Helper class for creating record batches iteratively given a known
0038 /// schema
0039 class ARROW_EXPORT RecordBatchBuilder {
0040  public:
0041   /// \brief Create and initialize a RecordBatchBuilder
0042   /// \param[in] schema The schema for the record batch
0043   /// \param[in] pool A MemoryPool to use for allocations
0044   /// \return the created builder instance
0045   static Result<std::unique_ptr<RecordBatchBuilder>> Make(
0046       const std::shared_ptr<Schema>& schema, MemoryPool* pool);
0047
0048   /// \brief Create and initialize a RecordBatchBuilder
0049   /// \param[in] schema The schema for the record batch
0050   /// \param[in] pool A MemoryPool to use for allocations
0051   /// \param[in] initial_capacity The initial capacity for the builders
0052   /// \return the created builder instance
0053   static Result<std::unique_ptr<RecordBatchBuilder>> Make(
0054       const std::shared_ptr<Schema>& schema, MemoryPool* pool, int64_t initial_capacity);
0055
0056   /// \brief Get base pointer to field builder
0057   /// \param i the field index
0058   /// \return pointer to ArrayBuilder
0059   ArrayBuilder* GetField(int i) { return raw_field_builders_[i]; }
0060
0061   /// \brief Return field builder casted to indicated specific builder type
0062   /// \param i the field index
0063   /// \return pointer to template type
0064   template <typename T>
0065   T* GetFieldAs(int i) {
0066     return internal::checked_cast<T*>(raw_field_builders_[i]);
0067   }
0068
0069   /// \brief Finish current batch and optionally reset
0070   /// \param[in] reset_builders the resulting RecordBatch
0071   /// \return the resulting RecordBatch
0072   Result<std::shared_ptr<RecordBatch>> Flush(bool reset_builders);
0073
0074   /// \brief Finish current batch and reset
0075   /// \return the resulting RecordBatch
0076   Result<std::shared_ptr<RecordBatch>> Flush();
0077
0078   /// \brief Set the initial capacity for new builders
0079   void SetInitialCapacity(int64_t capacity);
0080
0081   /// \brief The initial capacity for builders
0082   int64_t initial_capacity() const { return initial_capacity_; }
0083
0084   /// \brief The number of fields in the schema
0085   int num_fields() const { return schema_->num_fields(); }
0086
0087   /// \brief The number of fields in the schema
0088   std::shared_ptr<Schema> schema() const { return schema_; }
0089
0090  private:
0091   ARROW_DISALLOW_COPY_AND_ASSIGN(RecordBatchBuilder);
0092
0093   RecordBatchBuilder(const std::shared_ptr<Schema>& schema, MemoryPool* pool,
0094                      int64_t initial_capacity);
0095
0096   Status CreateBuilders();
0097   Status InitBuilders();
0098
0099   std::shared_ptr<Schema> schema_;
0100   int64_t initial_capacity_;
0101   MemoryPool* pool_;
0102
0103   std::vector<std::unique_ptr<ArrayBuilder>> field_builders_;
0104   std::vector<ArrayBuilder*> raw_field_builders_;
0105 };
0106
0107 }  // namespace arrow