lifetimes ¯\_(ツ)_/¯

spiraldb · May 8, 2024 · 4667b2d · 4667b2d
1 parent 2bfcb46
commit 4667b2d
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 55 deletions.
diff --git a/vortex-array/src/stats/flatbuffers.rs b/vortex-array/src/stats/flatbuffers.rs
@@ -0,0 +1,49 @@
+use flatbuffers::{FlatBufferBuilder, WIPOffset};
+use itertools::Itertools;
+use vortex_flatbuffers::WriteFlatBuffer;
+
+use crate::stats::{Stat, Statistics};
+
+impl WriteFlatBuffer for &dyn Statistics {
+    type Target<'t> = crate::flatbuffers::ArrayStats<'t>;
+
+    fn write_flatbuffer<'fb>(
+        &self,
+        fbb: &mut FlatBufferBuilder<'fb>,
+    ) -> WIPOffset<Self::Target<'fb>> {
+        let trailing_zero_freq = self
+            .get_as::<Vec<u64>>(Stat::TrailingZeroFreq)
+            .ok()
+            .map(|v| v.iter().copied().collect_vec())
+            .map(|v| fbb.create_vector(v.as_slice()));
+
+        let bit_width_freq = self
+            .get_as::<Vec<u64>>(Stat::BitWidthFreq)
+            .ok()
+            .map(|v| v.iter().copied().collect_vec())
+            .map(|v| fbb.create_vector(v.as_slice()));
+
+        let min = self
+            .get(Stat::Min)
+            .map(|min| min.value().write_flatbuffer(fbb));
+
+        let max = self
+            .get(Stat::Max)
+            .map(|max| max.value().write_flatbuffer(fbb));
+
+        let stat_args = &crate::flatbuffers::ArrayStatsArgs {
+            min,
+            max,
+            is_sorted: self.get_as::<bool>(Stat::IsSorted).ok(),
+            is_strict_sorted: self.get_as::<bool>(Stat::IsStrictSorted).ok(),
+            is_constant: self.get_as::<bool>(Stat::IsConstant).ok(),
+            run_count: self.get_as::<u64>(Stat::RunCount).ok(),
+            true_count: self.get_as::<u64>(Stat::TrueCount).ok(),
+            null_count: self.get_as::<u64>(Stat::NullCount).ok(),
+            bit_width_freq,
+            trailing_zero_freq,
+        };
+
+        crate::flatbuffers::ArrayStats::create(fbb, stat_args)
+    }
+}
diff --git a/vortex-array/src/stats/mod.rs b/vortex-array/src/stats/mod.rs
@@ -8,6 +8,7 @@ use vortex_dtype::{DType, NativePType};
 use vortex_error::{vortex_err, VortexError, VortexResult};
 use vortex_scalar::Scalar;
 
+pub mod flatbuffers;
 mod statsset;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Sequence)]

diff --git a/vortex-array/src/view.rs b/vortex-array/src/view.rs
@@ -212,9 +212,9 @@ impl Statistics for ArrayView<'_> {
 
         self.to_array()
             .with_dyn(|a| a.compute_statistics(stat))
-            .ok()?;
-
-        self.get(stat)
+            .ok()?
+            .get(stat)
+            .cloned()
     }
 }
 

diff --git a/vortex-ipc/src/messages.rs b/vortex-ipc/src/messages.rs
@@ -1,7 +1,6 @@
 use flatbuffers::{FlatBufferBuilder, WIPOffset};
 use itertools::Itertools;
-use vortex::flatbuffers as vfb;
-use vortex::stats::Stat;
+use vortex::flatbuffers as fba;
 use vortex::{ArrayData, Context, ViewContext};
 use vortex_dtype::DType;
 use vortex_error::{vortex_err, VortexError};
@@ -159,7 +158,7 @@ impl<'a> WriteFlatBuffer for IPCChunk<'a> {
 }
 
 impl<'a> WriteFlatBuffer for IPCArray<'a> {
-    type Target<'t> = vfb::Array<'t>;
+    type Target<'t> = fba::Array<'t>;
 
     fn write_flatbuffer<'fb>(
         &self,
@@ -191,63 +190,18 @@ impl<'a> WriteFlatBuffer for IPCArray<'a> {
             .collect_vec();
         let children = Some(fbb.create_vector(&children));
 
-        let stats = collect_array_stats(fbb, self.1);
+        let stats = Some(self.1.statistics().write_flatbuffer(fbb));
 
-        vfb::Array::create(
+        fba::Array::create(
             fbb,
-            &vfb::ArrayArgs {
+            &fba::ArrayArgs {
                 version: Default::default(),
                 has_buffer: column_data.buffer().is_some(),
                 encoding,
                 metadata,
-                stats: Some(stats),
+                stats,
                 children,
             },
         )
     }
 }
-
-/// Computes all stats and uses the results to create an ArrayStats table for the flatbuffer message
-fn collect_array_stats<'a>(
-    fbb: &'_ mut FlatBufferBuilder<'a>,
-    array: &'_ ArrayData,
-) -> WIPOffset<vfb::ArrayStats<'a>> {
-    let trailing_zero_freq = array
-        .statistics()
-        .get_as::<Vec<u64>>(Stat::TrailingZeroFreq)
-        .ok()
-        .map(|v| v.iter().copied().collect_vec())
-        .map(|v| fbb.create_vector(v.as_slice()));
-
-    let bit_width_freq = array
-        .statistics()
-        .get_as::<Vec<u64>>(Stat::BitWidthFreq)
-        .ok()
-        .map(|v| v.iter().copied().collect_vec())
-        .map(|v| fbb.create_vector(v.as_slice()));
-
-    let min = array
-        .statistics()
-        .get(Stat::Min)
-        .map(|min| min.value().write_flatbuffer(fbb));
-
-    let max = array
-        .statistics()
-        .get(Stat::Max)
-        .map(|max| max.value().write_flatbuffer(fbb));
-
-    let stat_args = &vfb::ArrayStatsArgs {
-        min,
-        max,
-        is_sorted: array.statistics().get_as::<bool>(Stat::IsSorted).ok(),
-        is_strict_sorted: array.statistics().get_as::<bool>(Stat::IsStrictSorted).ok(),
-        is_constant: array.statistics().get_as::<bool>(Stat::IsConstant).ok(),
-        run_count: array.statistics().get_as::<u64>(Stat::RunCount).ok(),
-        true_count: array.statistics().get_as::<u64>(Stat::TrueCount).ok(),
-        null_count: array.statistics().get_as::<u64>(Stat::NullCount).ok(),
-        bit_width_freq,
-        trailing_zero_freq,
-    };
-
-    vfb::ArrayStats::create(fbb, stat_args)
-}