Read a dictionary encoded int column.

Author: Nong Li

example/compute_stats.cc

@@ -49,7 +49,11 @@ int main(int argc, char** argv) {
       }
       InMemoryInputStream input(&column_buffer[0], column_buffer.size());
       ColumnReader reader(&metadata.schema[c + 1], &input);
-      printf("%d\n", reader.HasNext());
+      while (reader.HasNext()) {
+        int32_t val;
+        bool is_null = reader.GetInt32(&val);
+        if (!is_null) cout << val << endl;
+      }
     }
   }
   fclose(file);

src/encodings.h

@@ -13,7 +13,7 @@ class Decoder {
  public:
   virtual ~Decoder() {}
 
-  virtual void SetData(const uint8_t* data, int len) = 0;
+  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; }
@@ -25,17 +25,17 @@ class Decoder {
   int value_left() const { return num_values_; }
 
  protected:
-  Decoder(parquet::SchemaElement* schema, int num_values)
-    : schema_(schema), num_values_(num_values) {}
-  parquet::SchemaElement* schema_;
+  Decoder(const parquet::SchemaElement* schema) : schema_(schema), num_values_(0) {}
+  const parquet::SchemaElement* schema_;
   int num_values_;
 };
 
 class BoolDecoder : public Decoder {
  public:
-  BoolDecoder(parquet::SchemaElement* schema, int num_values) : Decoder(schema, num_values) { }
+  BoolDecoder(const parquet::SchemaElement* schema) : Decoder(schema) { }
 
-  virtual void SetData(const uint8_t* data, int len) {
+  virtual void SetData(int num_values, const uint8_t* data, int len) {
+    num_values_ = num_values;
     decoder_ = impala::RleDecoder(data, len, 1);
   }
 
@@ -52,11 +52,12 @@ class BoolDecoder : public Decoder {
 
 class PlainDecoder : public Decoder {
  public:
-  PlainDecoder(parquet::SchemaElement* schema, int num_values)
-    : Decoder(schema, num_values), data_(NULL), len_(0) {
+  PlainDecoder(const parquet::SchemaElement* schema)
+    : Decoder(schema), data_(NULL), len_(0) {
   }
 
-  virtual void SetData(const uint8_t* data, int len) {
+  virtual void SetData(int num_values, const uint8_t* data, int len) {
+    num_values_ = num_values;
     data_ = data;
     len_ = len;
   }
@@ -109,8 +110,8 @@ class PlainDecoder : public Decoder {
 
 class DictionaryDecoder : public Decoder {
  public:
-  DictionaryDecoder(parquet::SchemaElement* schema, int num_values, Decoder* dictionary)
-    : Decoder(schema, num_values) {
+  DictionaryDecoder(const parquet::SchemaElement* schema, Decoder* dictionary)
+    : Decoder(schema) {
     int num_dictionary_values = dictionary->value_left();
     switch (schema->type) {
       case parquet::Type::BOOLEAN: throw "Boolean cols should not be dictionary encoded.";
@@ -143,7 +144,8 @@ class DictionaryDecoder : public Decoder {
     }
   }
 
-  virtual void SetData(const uint8_t* data, int len) {
+  virtual void SetData(int num_values, const uint8_t* data, int len) {
+    num_values_ = num_values;
     if (len == 0) return;
     uint8_t bit_width = *data;
     ++data;

src/parquet.cc

@@ -34,37 +34,62 @@ ColumnReader::ColumnReader(const SchemaElement* schema, InputStream* stream)
   buffered_bytes_offset_ = 0;
 }
 
+bool ColumnReader::GetInt32(int32_t* result) {
+  --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;
+}
+
 bool ColumnReader::ReadNewPage() {
   uint32_t bytes_read = stream_->Read(&buffered_bytes_[0], buffered_bytes_.size());
+  if (bytes_read == 0) return false;
   uint32_t header_size = bytes_read;
   if (!DeserializeThriftMsg(&buffered_bytes_[0], &header_size, &current_page_header_)) {
     return false;
   }
   buffered_bytes_offset_ += header_size;
   // TODO: handle decompression.
+  cout << apache::thrift::ThriftDebugString(current_page_header_) << endl;
 
   if (current_page_header_.type == PageType::DICTIONARY_PAGE) {
     InitDictionary();
-//    return ReadNewPage();
+    header_size = bytes_read - buffered_bytes_offset_;
+    if (!DeserializeThriftMsg(&buffered_bytes_[buffered_bytes_offset_],
+          &header_size, &current_page_header_)) {
+      return false;
+    }
+    cout << apache::thrift::ThriftDebugString(current_page_header_) << endl;
+    buffered_bytes_offset_ += header_size;
   }
-
+  int num_definition_bytes = *reinterpret_cast<uint32_t*>(&buffered_bytes_[buffered_bytes_offset_]);
+  buffered_bytes_offset_ += 4;
+  definition_level_decoder_ =
+      impala::RleDecoder(&buffered_bytes_[buffered_bytes_offset_], num_definition_bytes, 1);
+  buffered_bytes_offset_ += num_definition_bytes;
+  num_buffered_values_ = current_page_header_.data_page_header.num_values;
+  decoder_->SetData(num_buffered_values_,
+      &buffered_bytes_[buffered_bytes_offset_], bytes_read - buffered_bytes_offset_);
   return true;
 }
 
 void ColumnReader::InitDictionary() {
-  for (int i = 0; i < current_page_header_.dictionary_page_header.num_values; ++i) {
-    uint8_t* data = &buffered_bytes_[buffered_bytes_offset_];
-    cout << *reinterpret_cast<int32_t*>(data) << endl;
-    buffered_bytes_offset_ += sizeof(int32_t);
-  }
+  uint8_t* data = &buffered_bytes_[buffered_bytes_offset_];
+  int len = current_page_header_.uncompressed_page_size;
+  buffered_bytes_offset_ += len;
+  PlainDecoder dictionary(schema_);
+  dictionary.SetData(current_page_header_.dictionary_page_header.num_values, data, len);
+  decoder_.reset(new DictionaryDecoder(schema_, &dictionary));
 }
 
 bool ColumnReader::HasNext() {
   if (num_buffered_values_ == 0) {
     ReadNewPage();
     if (num_buffered_values_ == 0) return false;
   }
-  return false;
+  return true;
 }
 
 }

src/parquet/parquet.h

@@ -5,6 +5,8 @@
 #include "gen-cpp/parquet_constants.h"
 #include "gen-cpp/parquet_types.h"
 
+#include "impala/rle-encoding.h"
+
 // TCompactProtocol requires some #defines to work right.
 #define SIGNED_RIGHT_SHIFT_IS 1
 #define ARITHMETIC_RIGHT_SHIFT 1
@@ -49,6 +51,8 @@ class ColumnReader {
   ColumnReader(const parquet::SchemaElement* schema, InputStream* stream);
   bool HasNext();
 
+  bool GetInt32(int32_t* result);
+
  private:
   bool ReadNewPage();
   void InitDictionary();
@@ -61,6 +65,7 @@ class ColumnReader {
   size_t num_buffered_bytes_;
   size_t buffered_bytes_offset_;
 
+  impala::RleDecoder definition_level_decoder_;
   boost::shared_ptr<Decoder> decoder_;
 };