hdfs-table-sink.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_EXEC_HDFS_TABLE_SINK_H
#define IMPALA_EXEC_HDFS_TABLE_SINK_H

#include <hdfs.h>
#include <boost/unordered_map.hpp>
#include <boost/scoped_ptr.hpp>

#include "common/object-pool.h"
#include "exec/data-sink.h"
#include "runtime/descriptors.h"
#include "util/runtime-profile.h"

namespace impala {

class Expr;
class TupleDescriptor;
class TupleRow;
class RuntimeState;
class HdfsTableWriter;
class MemTracker;

struct OutputPartition {

  std::string final_hdfs_file_name_prefix;

  std::string current_file_name;

  std::string tmp_hdfs_dir_name;

  std::string tmp_hdfs_file_name_prefix;

  std::string partition_name;

  hdfsFS hdfs_connection;

  hdfsFile tmp_hdfs_file;

  int64_t num_rows;

  int32_t num_files;

  boost::scoped_ptr<HdfsTableWriter> writer;

  const HdfsPartitionDescriptor* partition_descriptor;

  OutputPartition();
};

//
//
//
class HdfsTableSink : public DataSink {
 public:
  HdfsTableSink(const RowDescriptor& row_desc,
      const std::vector<TExpr>& select_list_texprs, const TDataSink& tsink);

  virtual Status Prepare(RuntimeState* state);

  virtual Status Open(RuntimeState* state);

  virtual Status Send(RuntimeState* state, RowBatch* batch, bool eos);

  virtual void Close(RuntimeState* state);

  static Status GetFileBlockSize(OutputPartition* output_partition, int64_t* size);

  virtual RuntimeProfile* profile() { return runtime_profile_; }
  const HdfsTableDescriptor& TableDesc() { return *table_desc_; }
  MemTracker* mem_tracker() { return mem_tracker_.get(); }

  RuntimeProfile::Counter* rows_inserted_counter() { return rows_inserted_counter_; }
  RuntimeProfile::Counter* bytes_written_counter() { return bytes_written_counter_; }
  RuntimeProfile::Counter* encode_timer() { return encode_timer_; }
  RuntimeProfile::Counter* hdfs_write_timer() { return hdfs_write_timer_; }
  RuntimeProfile::Counter* compress_timer() { return compress_timer_; }

  std::string DebugString() const;

 private:
  Status InitOutputPartition(RuntimeState* state,
                             const HdfsPartitionDescriptor& partition_descriptor,
                             OutputPartition* output_partition);

  Status CreateNewTmpFile(RuntimeState* state, OutputPartition* output_partition);

  typedef std::pair<OutputPartition*, std::vector<int32_t> > PartitionPair;
  typedef boost::unordered_map<std::string, PartitionPair> PartitionMap;


  void GetHashTblKey(const std::vector<ExprContext*>& ctxs, std::string* key);

  Status GetOutputPartition(RuntimeState* state, const std::string& key,
                            PartitionPair** partition_pair);

  Status PrepareExprs(RuntimeState* state);

  void BuildHdfsFileNames(const HdfsPartitionDescriptor& partition_descriptor,
      OutputPartition* output);

  Status FinalizePartitionFile(RuntimeState* state, OutputPartition* partition);

  void ClosePartitionFile(RuntimeState* state, OutputPartition* partition);

  const HdfsTableDescriptor* table_desc_;

  const HdfsPartitionDescriptor* default_partition_;

  std::vector<ExprContext*> output_expr_ctxs_;

  TupleRow* current_row_;

  hdfsFS hdfs_connection_;

  const RowDescriptor& row_desc_;

  TableId table_id_;

  const std::vector<TExpr>& select_list_texprs_;

  const std::vector<TExpr>& partition_key_texprs_;

  std::vector<ExprContext*> partition_key_expr_ctxs_;

  bool overwrite_;

  std::string staging_dir_;

  std::string unique_id_str_;

  PartitionMap partition_keys_to_output_partitions_;

  std::vector<ExprContext*> dynamic_partition_key_expr_ctxs_;

  typedef boost::unordered_map<std::string, HdfsPartitionDescriptor*>
      PartitionDescriptorMap;
  PartitionDescriptorMap partition_descriptor_map_;

  boost::scoped_ptr<MemTracker> mem_tracker_;

  RuntimeProfile* runtime_profile_;
  RuntimeProfile::Counter* partitions_created_counter_;
  RuntimeProfile::Counter* files_created_counter_;
  RuntimeProfile::Counter* rows_inserted_counter_;
  RuntimeProfile::Counter* bytes_written_counter_;

  RuntimeProfile::Counter* encode_timer_;
  RuntimeProfile::Counter* hdfs_write_timer_;
  RuntimeProfile::Counter* compress_timer_;

  bool has_empty_input_batch_;
};
}
#endif