use BufferedStats for buffered data

Author: wecharyu

cpp/src/parquet/column_writer.cc

@@ -1480,8 +1480,19 @@ class TypedColumnWriterImpl : public ColumnWriterImpl,
     return current_encoder_->EstimatedDataEncodedSize();
   }
 
-  int64_t EstimatedBufferedLevelsBytes() const override {
-    return definition_levels_sink_.length() + repetition_levels_sink_.length();
+  int64_t estimated_buffered_def_level_bytes() const override {
+    return definition_levels_sink_.length();
+  }
+
+  int64_t estimated_buffered_rep_level_bytes() const override {
+    return repetition_levels_sink_.length();
+  }
+
+  int64_t estimated_buffered_dict_bytes() const override {
+    if (current_dict_encoder_) {
+      return current_dict_encoder_->dict_encoded_size();
+    }
+    return 0;
   }
 
  protected:

cpp/src/parquet/column_writer.h

@@ -164,8 +164,14 @@ class PARQUET_EXPORT ColumnWriter {
   /// \brief Estimated size of the values that are not written to a page yet.
   virtual int64_t estimated_buffered_value_bytes() const = 0;
 
-  /// \brief Estimated size of the levels that are not written to a page yet.
-  virtual int64_t EstimatedBufferedLevelsBytes() const = 0;
+  /// \brief Estimated size of the definition levels that are not written to a page yet.
+  virtual int64_t estimated_buffered_def_level_bytes() const = 0;
+
+  /// \brief Estimated size of the repetition levels that are not written to a page yet.
+  virtual int64_t estimated_buffered_rep_level_bytes() const = 0;
+
+  /// \brief Estimated size of the dictionary that are not written to a page yet.
+  virtual int64_t estimated_buffered_dict_bytes() const = 0;
 
   /// \brief The file-level writer properties
   virtual const WriterProperties* properties() = 0;

cpp/src/parquet/file_writer.cc

@@ -68,8 +68,8 @@ int64_t RowGroupWriter::total_compressed_bytes_written() const {
   return contents_->total_compressed_bytes_written();
 }
 
-int64_t RowGroupWriter::total_buffered_bytes() const {
-  return contents_->EstimatedBufferedBytes();
+BufferedStats RowGroupWriter::estimated_buffered_stats() const {
+  return contents_->EstimatedBufferedStats();
 }
 
 bool RowGroupWriter::buffered() const { return contents_->buffered(); }
@@ -199,19 +199,20 @@ class RowGroupSerializer : public RowGroupWriter::Contents {
     return total_compressed_bytes_written;
   }
 
-  int64_t EstimatedBufferedBytes() const override {
+  BufferedStats EstimatedBufferedStats() const override {
+    BufferedStats stats;
     if (closed_) {
-      return 0;
+      return stats;
     }
-    int64_t estimated_buffered_value_bytes = 0;
     for (size_t i = 0; i < column_writers_.size(); i++) {
       if (column_writers_[i]) {
-        estimated_buffered_value_bytes +=
-            column_writers_[i]->estimated_buffered_value_bytes() +
-            column_writers_[i]->EstimatedBufferedLevelsBytes();
+        stats.def_level_bytes += column_writers_[i]->estimated_buffered_def_level_bytes();
+        stats.rep_level_bytes += column_writers_[i]->estimated_buffered_rep_level_bytes();
+        stats.value_bytes += column_writers_[i]->estimated_buffered_value_bytes();
+        stats.dict_bytes += column_writers_[i]->estimated_buffered_dict_bytes();
       }
     }
-    return estimated_buffered_value_bytes;
+    return stats;
   }
 
   bool buffered() const override { return buffered_row_group_; }

cpp/src/parquet/file_writer.h

@@ -34,6 +34,13 @@ class ColumnWriter;
 static constexpr uint8_t kParquetMagic[4] = {'P', 'A', 'R', '1'};
 static constexpr uint8_t kParquetEMagic[4] = {'P', 'A', 'R', 'E'};
 
+struct PARQUET_EXPORT BufferedStats {
+  int64_t def_level_bytes = 0;
+  int64_t rep_level_bytes = 0;
+  int64_t value_bytes = 0;
+  int64_t dict_bytes = 0;
+};
+
 class PARQUET_EXPORT RowGroupWriter {
  public:
   // Forward declare a virtual class 'Contents' to aid dependency injection and more
@@ -58,9 +65,9 @@ class PARQUET_EXPORT RowGroupWriter {
     virtual int64_t total_compressed_bytes() const = 0;
     /// \brief total compressed bytes written by the page writer
     virtual int64_t total_compressed_bytes_written() const = 0;
-    /// \brief Estimated bytes of values and levels that are buffered by the page writer
-    /// but not written to a page yet
-    virtual int64_t EstimatedBufferedBytes() const = 0;
+    /// \brief Estimated sizes of buffered data (levels, values, dict) not yet
+    /// written to a page.
+    virtual BufferedStats EstimatedBufferedStats() const = 0;
 
     virtual bool buffered() const = 0;
   };
@@ -102,8 +109,9 @@ class PARQUET_EXPORT RowGroupWriter {
   int64_t total_compressed_bytes() const;
   /// \brief total compressed bytes written by the page writer
   int64_t total_compressed_bytes_written() const;
-  /// \brief total bytes of values and levels that are buffered by the page writer
-  int64_t total_buffered_bytes() const;
+  /// \brief Estimated sizes of buffered data (levels, values, dict) not yet
+  /// written to a page.
+  BufferedStats estimated_buffered_stats() const;
 
   /// Returns whether the current RowGroupWriter is in the buffered mode and is created
   /// by calling ParquetFileWriter::AppendBufferedRowGroup.