row-batch.h

Go to the documentation of this file.

// Copyright 2012 Cloudera Inc.
//
// Licensed 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.


#ifndef IMPALA_RUNTIME_ROW_BATCH_H
#define IMPALA_RUNTIME_ROW_BATCH_H

#include <vector>
#include <cstring>
#include <boost/scoped_ptr.hpp>

#include "common/logging.h"
#include "runtime/descriptors.h"
#include "runtime/disk-io-mgr.h"
#include "runtime/mem-pool.h"
#include "runtime/mem-tracker.h"

namespace impala {

class BufferedTupleStream;
class MemTracker;
class TRowBatch;
class Tuple;
class TupleRow;
class TupleDescriptor;

//
//
class RowBatch {
 public:
  RowBatch(const RowDescriptor& row_desc, int capacity, MemTracker* tracker);

  RowBatch(const RowDescriptor& row_desc, const TRowBatch& input_batch,
      MemTracker* tracker);

  ~RowBatch();

  static const int INVALID_ROW_INDEX = -1;

  int AddRows(int n) {
    if (num_rows_ + n > capacity_) return INVALID_ROW_INDEX;
    has_in_flight_row_ = true;
    return num_rows_;
  }

  int AddRow() { return AddRows(1); }

  void CommitRows(int n) {
    DCHECK_GE(n, 0);
    DCHECK_LE(num_rows_ + n, capacity_);
    num_rows_ += n;
    has_in_flight_row_ = false;
  }

  void CommitLastRow() { CommitRows(1); }

  void set_num_rows(int num_rows) {
    DCHECK_LE(num_rows, num_rows_);
    num_rows_ = num_rows;
  }

  bool AtCapacity() {
    return num_rows_ == capacity_ || auxiliary_mem_usage_ >= AT_CAPACITY_MEM_USAGE ||
      num_tuple_streams() > 0 || need_to_return_;
  }

  bool AtCapacity(MemPool* tuple_pool) {
    DCHECK(tuple_pool != NULL);
    return AtCapacity() ||
        (tuple_pool->total_allocated_bytes() > AT_CAPACITY_MEM_USAGE && num_rows_ > 0);
  }

  int TotalByteSize();

  TupleRow* GetRow(int row_idx) {
    DCHECK(tuple_ptrs_ != NULL);
    DCHECK_GE(row_idx, 0);
    DCHECK_LT(row_idx, num_rows_ + (has_in_flight_row_ ? 1 : 0));
    return reinterpret_cast<TupleRow*>(tuple_ptrs_ + row_idx * num_tuples_per_row_);
  }

  int row_byte_size() { return num_tuples_per_row_ * sizeof(Tuple*); }
  MemPool* tuple_data_pool() { return tuple_data_pool_.get(); }
  int num_io_buffers() const { return io_buffers_.size(); }
  int num_tuple_streams() const { return tuple_streams_.size(); }

  void Reset();

  void AddIoBuffer(DiskIoMgr::BufferDescriptor* buffer);

  void AddTupleStream(BufferedTupleStream* stream);

  void MarkNeedToReturn() { need_to_return_ = true; }

  void TransferResourceOwnership(RowBatch* dest);

  void CopyRow(TupleRow* src, TupleRow* dest) {
    memcpy(dest, src, num_tuples_per_row_ * sizeof(Tuple*));
  }

  void CopyRows(int dest, int src, int num_rows) {
    DCHECK_LE(dest, src);
    DCHECK_LE(src + num_rows, capacity_);
    memmove(tuple_ptrs_ + num_tuples_per_row_ * dest,
        tuple_ptrs_ + num_tuples_per_row_ * src,
        num_rows * num_tuples_per_row_ * sizeof(Tuple*));
  }

  void ClearRow(TupleRow* row) {
    memset(row, 0, num_tuples_per_row_ * sizeof(Tuple*));
  }

  void AcquireState(RowBatch* src);

  int Serialize(TRowBatch* output_batch);

  static int GetBatchSize(const TRowBatch& batch);

  int num_rows() const { return num_rows_; }
  int capacity() const { return capacity_; }

  const RowDescriptor& row_desc() const { return row_desc_; }

  static const int AT_CAPACITY_MEM_USAGE;

  int MaxTupleBufferSize();

 private:
  MemTracker* mem_tracker_;  // not owned


  bool has_in_flight_row_;  // if true, last row hasn't been committed yet
  int num_rows_;  // # of committed rows
  int capacity_;  // maximum # of rows

  int num_tuples_per_row_;
  RowDescriptor row_desc_;

  Tuple** tuple_ptrs_;
  int tuple_ptrs_size_;

  int64_t auxiliary_mem_usage_;

  bool need_to_return_;

  boost::scoped_ptr<MemPool> tuple_data_pool_;

  std::vector<DiskIoMgr::BufferDescriptor*> io_buffers_;

  std::vector<BufferedTupleStream*> tuple_streams_;

  std::string compression_scratch_;
};

}

#endif