From 1a20f3fb4ae8950fb9ba834b18812a2c098f3ffa Mon Sep 17 00:00:00 2001
From: Edd Robinson <me@edd.io>
Date: Fri, 7 May 2021 12:39:16 +0100
Subject: [PATCH] feat: implement raw column size on float columns

---
 read_buffer/src/column/float.rs | 44 ++++++++++++++++++++++++++++++++-
 1 file changed, 43 insertions(+), 1 deletion(-)

diff --git a/read_buffer/src/column/float.rs b/read_buffer/src/column/float.rs
index 09f5f280ce..be4fb43abd 100644
--- a/read_buffer/src/column/float.rs
+++ b/read_buffer/src/column/float.rs
@@ -1,3 +1,5 @@
+use std::mem::size_of;
+
 use arrow::{self, array::Array};
 
 use super::encoding::{scalar::Fixed, scalar::FixedNull};
@@ -10,7 +12,7 @@ pub enum FloatEncoding {
 }
 
 impl FloatEncoding {
-    /// The total size in bytes of the store columnar data.
+    /// The total size in bytes of to store columnar data in memory.
     pub fn size(&self) -> usize {
         match self {
             Self::Fixed64(enc) => enc.size(),
@@ -18,6 +20,19 @@ impl FloatEncoding {
         }
     }
 
+    /// The estimated total size in bytes of the underlying float values in the
+    /// column if they were stored contiguously and uncompressed. `include_nulls`
+    /// will effectively size each NULL value as 8b if `true`.
+    pub fn size_raw(&self, include_nulls: bool) -> usize {
+        match self {
+            // this will be the size of a Vec<f64>
+            Self::Fixed64(enc) => {
+                size_of::<Vec<f64>>() + (enc.num_rows() as usize * size_of::<f64>())
+            }
+            Self::FixedNull64(enc) => enc.size_raw(include_nulls),
+        }
+    }
+
     /// The total number of rows in the column.
     pub fn num_rows(&self) -> u32 {
         match self {
@@ -212,3 +227,30 @@ impl From<arrow::array::Float64Array> for FloatEncoding {
         Self::FixedNull64(FixedNull::<arrow::datatypes::Float64Type>::from(arr))
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn size_raw() {
+        let enc = FloatEncoding::from(&[2.2, 22.2, 12.2, 31.2][..]);
+        // (4 * 8) + 24
+        assert_eq!(enc.size_raw(true), 56);
+        assert_eq!(enc.size_raw(false), 56);
+
+        let enc = FloatEncoding::FixedNull64(FixedNull::<arrow::datatypes::Float64Type>::from(
+            &[2.0, 2.02, 1.02, 3.01][..],
+        ));
+        // (4 * 8) + 24
+        assert_eq!(enc.size_raw(true), 56);
+        assert_eq!(enc.size_raw(false), 56);
+
+        let enc = FloatEncoding::FixedNull64(FixedNull::<arrow::datatypes::Float64Type>::from(
+            &[Some(2.0), Some(2.02), None, Some(1.02), Some(3.01)][..],
+        ));
+        // (5 * 8) + 24
+        assert_eq!(enc.size_raw(true), 64);
+        assert_eq!(enc.size_raw(false), 56);
+    }
+}