Better decoder management.

Author: Nong Li

example/compute_stats.cc

@@ -50,9 +50,9 @@ int main(int argc, char** argv) {
       InMemoryInputStream input(&column_buffer[0], column_buffer.size());
       ColumnReader reader(&metadata.schema[c + 1], &input);
       while (reader.HasNext()) {
-        int32_t val;
-        bool is_null = reader.GetInt32(&val);
-        if (!is_null) cout << val << endl;
+        int def_level, rep_level;
+        int32_t val = reader.GetInt32(&def_level, &rep_level);;
+        if (def_level >= rep_level) cout << val << endl;
       }
     }
   }

src/encodings.h

@@ -13,26 +13,46 @@ class Decoder {
  public:
   virtual ~Decoder() {}
 
+  // Sets the data for a new page. This will be called multiple times on the same
+  // decoder and should reset all internal state.
   virtual void SetData(int num_values, const uint8_t* data, int len) = 0;
 
   // Subclasses should override the ones they support
-  virtual bool GetBool() { return false; }
-  virtual int32_t GetInt32() { return 0; }
-  virtual int64_t GetInt64() { return 0; }
-  virtual float GetFloat() { return 0; }
-  virtual String GetString() { return String(); }
+  virtual bool GetBool() {
+    throw ParquetException("Decoder does not implement this type.");
+  }
+  virtual int32_t GetInt32() {
+    throw ParquetException("Decoder does not implement this type.");
+  }
+  virtual int64_t GetInt64() {
+    throw ParquetException("Decoder does not implement this type.");
+  }
+  virtual float GetFloat() {
+    throw ParquetException("Decoder does not implement this type.");
+  }
+  virtual String GetString() {
+    throw ParquetException("Decoder does not implement this type.");
+  }
 
-  int value_left() const { return num_values_; }
+  // Returns the number of values left (for the last call to SetData()). This is
+  // the number of values left in this page.
+  int values_left() const { return num_values_; }
+
+  const parquet::Encoding::type encoding() const { return encoding_; }
 
  protected:
-  Decoder(const parquet::SchemaElement* schema) : schema_(schema), num_values_(0) {}
+  Decoder(const parquet::SchemaElement* schema, const parquet::Encoding::type& encoding)
+    : schema_(schema), encoding_(encoding), num_values_(0) {}
+
   const parquet::SchemaElement* schema_;
+  const parquet::Encoding::type encoding_;
   int num_values_;
 };
 
 class BoolDecoder : public Decoder {
  public:
-  BoolDecoder(const parquet::SchemaElement* schema) : Decoder(schema) { }
+  BoolDecoder(const parquet::SchemaElement* schema)
+    : Decoder(schema, parquet::Encoding::PLAIN) { }
 
   virtual void SetData(int num_values, const uint8_t* data, int len) {
     num_values_ = num_values;
@@ -41,7 +61,7 @@ class BoolDecoder : public Decoder {
 
   virtual bool GetBool() {
     bool result;
-    if (!decoder_.Get(&result)) throw "EOF";
+    if (!decoder_.Get(&result)) ParquetException::EofException();
     --num_values_;
     return result;
   }
@@ -53,7 +73,7 @@ class BoolDecoder : public Decoder {
 class PlainDecoder : public Decoder {
  public:
   PlainDecoder(const parquet::SchemaElement* schema)
-    : Decoder(schema), data_(NULL), len_(0) {
+    : Decoder(schema, parquet::Encoding::PLAIN), data_(NULL), len_(0) {
   }
 
   virtual void SetData(int num_values, const uint8_t* data, int len) {
@@ -63,7 +83,7 @@ class PlainDecoder : public Decoder {
   }
 
   virtual int32_t GetInt32() {
-    if (len_ < sizeof(int32_t)) throw "EOF";
+    if (len_ < sizeof(int32_t)) ParquetException::EofException();
     int32_t val = *reinterpret_cast<const int32_t*>(data_);
     data_ += sizeof(int32_t);
     len_ -= sizeof(int32_t);
@@ -72,7 +92,7 @@ class PlainDecoder : public Decoder {
   }
 
   virtual int64_t GetInt64() {
-    if (len_ < sizeof(int64_t)) throw "EOF";
+    if (len_ < sizeof(int64_t)) ParquetException::EofException();
     int64_t val = *reinterpret_cast<const int64_t*>(data_);
     data_ += sizeof(int64_t);
     len_ -= sizeof(int64_t);
@@ -81,7 +101,7 @@ class PlainDecoder : public Decoder {
   }
 
   virtual float GetFloat() {
-    if (len_ < sizeof(float)) throw "EOF";
+    if (len_ < sizeof(float)) ParquetException::EofException();
     float val = *reinterpret_cast<const float*>(data_);
     data_ += sizeof(float);
     len_ -= sizeof(float);
@@ -91,11 +111,11 @@ class PlainDecoder : public Decoder {
 
   virtual String GetString() {
     String result;
-    if (len_ < sizeof(uint32_t)) throw "EOF";
+    if (len_ < sizeof(uint32_t)) ParquetException::EofException();
     result.len = *reinterpret_cast<const uint32_t*>(data_);
     data_ += sizeof(uint32_t);
     len_ -= sizeof(uint32_t);
-    if (len_ < result.len) throw "EOF";
+    if (len_ < result.len) ParquetException::EofException();
     result.ptr = data_;
     data_ += result.len;
     len_ -= result.len;
@@ -111,10 +131,12 @@ class PlainDecoder : public Decoder {
 class DictionaryDecoder : public Decoder {
  public:
   DictionaryDecoder(const parquet::SchemaElement* schema, Decoder* dictionary)
-    : Decoder(schema) {
-    int num_dictionary_values = dictionary->value_left();
+    : Decoder(schema, parquet::Encoding::RLE_DICTIONARY) {
+    int num_dictionary_values = dictionary->values_left();
     switch (schema->type) {
-      case parquet::Type::BOOLEAN: throw "Boolean cols should not be dictionary encoded.";
+      case parquet::Type::BOOLEAN:
+        throw ParquetException("Boolean cols should not be dictionary encoded.");
+
       case parquet::Type::INT32:
         int32_dictionary_.resize(num_dictionary_values);
         for (int i = 0; i < num_dictionary_values; ++i) {
@@ -140,7 +162,7 @@ class DictionaryDecoder : public Decoder {
         }
         break;
       default:
-        throw "NYI";
+        ParquetException::NYI();
     }
   }
 
@@ -161,7 +183,7 @@ class DictionaryDecoder : public Decoder {
  private:
   int index() {
     int idx;
-    if (!idx_decoder_.Get(&idx)) throw "EOF";
+    if (!idx_decoder_.Get(&idx)) ParquetException::EofException();
     --num_values_;
     return idx;
   }

src/parquet.cc

@@ -8,6 +8,7 @@
 
 const int DATA_PAGE_SIZE = 64 * 1024;
 
+using namespace boost;
 using namespace parquet;
 using namespace std;
 
@@ -31,16 +32,22 @@ const uint8_t* InMemoryInputStream::Read(int num_to_read, int* num_bytes) {
 ColumnReader::ColumnReader(const SchemaElement* schema, InputStream* stream)
   : schema_(schema),
     stream_(stream),
-    num_buffered_values_(0) {
+    num_buffered_values_(0),
+    current_decoder_(NULL) {
 }
 
-bool ColumnReader::GetInt32(int32_t* result) {
+int32_t ColumnReader::GetInt32(int* definition_level, int* repetition_level) {
   --num_buffered_values_;
-  int def_level;
-  if (!definition_level_decoder_.Get(&def_level)) throw "EOF";
-  if (def_level == 0) return true;
-  *result = decoder_->GetInt32();
-  return false;
+  *repetition_level = 1;
+  if (!definition_level_decoder_.Get(definition_level)) ParquetException::EofException();
+  if (*definition_level == 0) return true;
+  return current_decoder_->GetInt32();
+}
+
+static bool IsDictionaryIndexEncoding(const Encoding::type& e) {
+  // PLAIN_DICTIONARY is deprecated but used to be used as a dictionary index
+  // encoding.
+  return e == Encoding::RLE_DICTIONARY || e == Encoding::PLAIN_DICTIONARY;
 }
 
 bool ColumnReader::ReadNewPage() {
@@ -54,18 +61,25 @@ bool ColumnReader::ReadNewPage() {
       return false;
     }
     stream_->Read(header_size, &bytes_read);
-    cout << apache::thrift::ThriftDebugString(current_page_header_) << endl;
 
     // TODO: handle decompression.
     int uncompressed_len = current_page_header_.uncompressed_page_size;
     buffer = stream_->Read(uncompressed_len, &bytes_read);
-    if (bytes_read != uncompressed_len) throw "EOF";
+    if (bytes_read != uncompressed_len) ParquetException::EofException();
 
     if (current_page_header_.type == PageType::DICTIONARY_PAGE) {
+      unordered_map<Encoding::type, shared_ptr<Decoder> >::iterator it =
+          decoders_.find(Encoding::RLE_DICTIONARY);
+      if (it != decoders_.end()) {
+        throw ParquetException("Column cannot have more than one dictionary.");
+      }
+
       PlainDecoder dictionary(schema_);
       dictionary.SetData(current_page_header_.dictionary_page_header.num_values,
           buffer, uncompressed_len);
-      decoder_.reset(new DictionaryDecoder(schema_, &dictionary));
+      shared_ptr<Decoder> decoder(new DictionaryDecoder(schema_, &dictionary));
+      decoders_[Encoding::RLE_DICTIONARY] = decoder;
+      current_decoder_ = decoders_[Encoding::RLE_DICTIONARY].get();
       continue;
     } else if (current_page_header_.type == PageType::DATA_PAGE) {
       // Read a data page.
@@ -79,8 +93,36 @@ bool ColumnReader::ReadNewPage() {
 
       // TODO: repetition levels
 
-      // Now buffer is at the start of the data.
-      decoder_->SetData(num_buffered_values_, buffer,
+      // Get a decoder object for this page or create a new decoder if this is the
+      // first page with this encoding.
+      Encoding::type encoding = current_page_header_.data_page_header.encoding;
+      if (IsDictionaryIndexEncoding(encoding)) encoding = Encoding::RLE_DICTIONARY;
+
+      unordered_map<Encoding::type, shared_ptr<Decoder> >::iterator it =
+          decoders_.find(encoding);
+      if (it != decoders_.end()) {
+        current_decoder_ = it->second.get();
+      } else {
+        switch (encoding) {
+          case Encoding::PLAIN: {
+            shared_ptr<Decoder> decoder(new PlainDecoder(schema_));
+            decoders_[encoding] = decoder;
+            current_decoder_ = decoder.get();
+            break;
+          }
+          case Encoding::RLE_DICTIONARY:
+            throw ParquetException("Dictionary page must be before data page.");
+
+          case Encoding::DELTA_BINARY_PACKED:
+          case Encoding::DELTA_LENGTH_BYTE_ARRAY:
+          case Encoding::DELTA_BYTE_ARRAY:
+            ParquetException::NYI();
+
+          default:
+            throw ParquetException("Unknown encoding type.");
+        }
+      }
+      current_decoder_->SetData(num_buffered_values_, buffer,
           uncompressed_len - sizeof(uint32_t) - num_definition_bytes);
     } else {
       // We don't know what this page type is. just skip it.

src/parquet/parquet.h

@@ -1,7 +1,9 @@
 #ifndef PARQUET_PARQUET_H_
 #define PARQUET_PARQUET_H_
 
+#include <exception>
 #include <boost/cstdint.hpp>
+#include <boost/unordered_map.hpp>
 #include "gen-cpp/parquet_constants.h"
 #include "gen-cpp/parquet_types.h"
 
@@ -26,16 +28,45 @@ struct String {
   const uint8_t* ptr;
 };
 
+class ParquetException : public std::exception {
+ public:
+  static void EofException() { throw ParquetException("Expected end of stream."); }
+  static void NYI() { throw ParquetException("Not yet implemented."); }
+
+  explicit ParquetException(const char* msg) : msg_(msg) {}
+  explicit ParquetException(const std::string& msg) : msg_(msg) {}
+
+  virtual ~ParquetException() throw() {}
+  virtual const char* what() const throw() { return msg_.c_str(); }
+
+ private:
+  std::string msg_;
+};
+
+// Interface for the column reader to get the bytes. The interface is a stream
+// interface, meaning the bytes in order and once a byte is read, it does not
+// need to be read again.
 class InputStream {
  public:
-  virtual ~InputStream() {}
+  // Returns the next 'num_to_peek' without advancing the current position.
+  // *num_bytes will contain the number of bytes returned which can only be
+  // less than num_to_peek at end of stream cases.
+  // Since the position is not advanced, calls to this function are idempotent.
+  // The buffer returned to the caller is still owned by the input stream and must
+  // stay valid until the next call to Peek() or Read().
   virtual const uint8_t* Peek(int num_to_peek, int* num_bytes) = 0;
+
+  // Identical to Peek(), except the current position in the stream is advanced by
+  // *num_bytes.
   virtual const uint8_t* Read(int num_to_read, int* num_bytes) = 0;
 
+  virtual ~InputStream() {}
+
  protected:
   InputStream() {}
 };
 
+// Implementation of an InputStream when all the bytes are in memory.
 class InMemoryInputStream : public InputStream {
  public:
   InMemoryInputStream(const uint8_t* buffer, int64_t len);
@@ -48,12 +79,13 @@ class InMemoryInputStream : public InputStream {
   int64_t offset_;
 };
 
+// API to read values from a single column.
 class ColumnReader {
  public:
   ColumnReader(const parquet::SchemaElement* schema, InputStream* stream);
-  bool HasNext();
 
-  bool GetInt32(int32_t* result);
+  bool HasNext();
+  int32_t GetInt32(int* definition_level, int* repetition_level);
 
  private:
   bool ReadNewPage();
@@ -64,7 +96,8 @@ class ColumnReader {
   parquet::PageHeader current_page_header_;
 
   impala::RleDecoder definition_level_decoder_;
-  boost::shared_ptr<Decoder> decoder_;
+  boost::unordered_map<parquet::Encoding::type, boost::shared_ptr<Decoder> > decoders_;
+  Decoder* current_decoder_;
 };
 
 // Deserialize a thrift message from buf/len.  buf/len must at least contain