refactor(rust): make Parquet Statistics into enum instead of `tra…

…it` (#16485)
pola-rs · May 25, 2024 · d68d499 · d68d499
1 parent f2bbc39
commit d68d499
Show file tree

Hide file tree

Showing 30 changed files with 513 additions and 585 deletions.
diff --git a/crates/polars-parquet/src/arrow/read/statistics/binary.rs b/crates/polars-parquet/src/arrow/read/statistics/binary.rs
@@ -2,10 +2,10 @@ use arrow::array::{MutableArray, MutableBinaryArray};
 use arrow::offset::Offset;
 use polars_error::PolarsResult;
 
-use crate::parquet::statistics::{BinaryStatistics, Statistics as ParquetStatistics};
+use crate::parquet::statistics::BinaryStatistics;
 
 pub(super) fn push<O: Offset>(
-    from: Option<&dyn ParquetStatistics>,
+    from: Option<&BinaryStatistics>,
     min: &mut dyn MutableArray,
     max: &mut dyn MutableArray,
 ) -> PolarsResult<()> {
@@ -17,8 +17,9 @@ pub(super) fn push<O: Offset>(
         .as_mut_any()
         .downcast_mut::<MutableBinaryArray<O>>()
         .unwrap();
-    let from = from.map(|s| s.as_any().downcast_ref::<BinaryStatistics>().unwrap());
+
     min.push(from.and_then(|s| s.min_value.as_ref()));
     max.push(from.and_then(|s| s.max_value.as_ref()));
+
     Ok(())
 }
diff --git a/crates/polars-parquet/src/arrow/read/statistics/binview.rs b/crates/polars-parquet/src/arrow/read/statistics/binview.rs
@@ -1,10 +1,10 @@
 use arrow::array::{MutableArray, MutableBinaryViewArray, ViewType};
 use polars_error::PolarsResult;
 
-use crate::parquet::statistics::{BinaryStatistics, Statistics as ParquetStatistics};
+use crate::parquet::statistics::BinaryStatistics;
 
 pub(super) fn push<T: ViewType + ?Sized>(
-    from: Option<&dyn ParquetStatistics>,
+    from: Option<&BinaryStatistics>,
     min: &mut dyn MutableArray,
     max: &mut dyn MutableArray,
 ) -> PolarsResult<()> {
@@ -16,7 +16,7 @@ pub(super) fn push<T: ViewType + ?Sized>(
         .as_mut_any()
         .downcast_mut::<MutableBinaryViewArray<T>>()
         .unwrap();
-    let from = from.map(|s| s.as_any().downcast_ref::<BinaryStatistics>().unwrap());
+
     min.push(from.and_then(|s| {
         let opt_b = s.min_value.as_deref();
         unsafe { opt_b.map(|b| T::from_bytes_unchecked(b)) }
@@ -25,5 +25,6 @@ pub(super) fn push<T: ViewType + ?Sized>(
         let opt_b = s.max_value.as_deref();
         unsafe { opt_b.map(|b| T::from_bytes_unchecked(b)) }
     }));
+
     Ok(())
 }
diff --git a/crates/polars-parquet/src/arrow/read/statistics/boolean.rs b/crates/polars-parquet/src/arrow/read/statistics/boolean.rs
@@ -1,10 +1,10 @@
 use arrow::array::{MutableArray, MutableBooleanArray};
 use polars_error::PolarsResult;
 
-use crate::parquet::statistics::{BooleanStatistics, Statistics as ParquetStatistics};
+use crate::parquet::statistics::BooleanStatistics;
 
 pub(super) fn push(
-    from: Option<&dyn ParquetStatistics>,
+    from: Option<&BooleanStatistics>,
     min: &mut dyn MutableArray,
     max: &mut dyn MutableArray,
 ) -> PolarsResult<()> {
@@ -16,8 +16,9 @@ pub(super) fn push(
         .as_mut_any()
         .downcast_mut::<MutableBooleanArray>()
         .unwrap();
-    let from = from.map(|s| s.as_any().downcast_ref::<BooleanStatistics>().unwrap());
+
     min.push(from.and_then(|s| s.min_value));
     max.push(from.and_then(|s| s.max_value));
+
     Ok(())
 }
diff --git a/crates/polars-parquet/src/arrow/read/statistics/fixlen.rs b/crates/polars-parquet/src/arrow/read/statistics/fixlen.rs
@@ -5,10 +5,10 @@ use polars_error::PolarsResult;
 
 use super::super::{convert_days_ms, convert_i128};
 use crate::arrow::read::convert_i256;
-use crate::parquet::statistics::{FixedLenStatistics, Statistics as ParquetStatistics};
+use crate::parquet::statistics::FixedLenStatistics;
 
 pub(super) fn push_i128(
-    from: Option<&dyn ParquetStatistics>,
+    from: Option<&FixedLenStatistics>,
     n: usize,
     min: &mut dyn MutableArray,
     max: &mut dyn MutableArray,
@@ -21,7 +21,6 @@ pub(super) fn push_i128(
         .as_mut_any()
         .downcast_mut::<MutablePrimitiveArray<i128>>()
         .unwrap();
-    let from = from.map(|s| s.as_any().downcast_ref::<FixedLenStatistics>().unwrap());
 
     min.push(from.and_then(|s| s.min_value.as_deref().map(|x| convert_i128(x, n))));
     max.push(from.and_then(|s| s.max_value.as_deref().map(|x| convert_i128(x, n))));
@@ -30,7 +29,7 @@ pub(super) fn push_i128(
 }
 
 pub(super) fn push_i256_with_i128(
-    from: Option<&dyn ParquetStatistics>,
+    from: Option<&FixedLenStatistics>,
     n: usize,
     min: &mut dyn MutableArray,
     max: &mut dyn MutableArray,
@@ -43,7 +42,6 @@ pub(super) fn push_i256_with_i128(
         .as_mut_any()
         .downcast_mut::<MutablePrimitiveArray<i256>>()
         .unwrap();
-    let from = from.map(|s| s.as_any().downcast_ref::<FixedLenStatistics>().unwrap());
 
     min.push(from.and_then(|s| {
         s.min_value
@@ -60,7 +58,7 @@ pub(super) fn push_i256_with_i128(
 }
 
 pub(super) fn push_i256(
-    from: Option<&dyn ParquetStatistics>,
+    from: Option<&FixedLenStatistics>,
     min: &mut dyn MutableArray,
     max: &mut dyn MutableArray,
 ) -> PolarsResult<()> {
@@ -72,7 +70,6 @@ pub(super) fn push_i256(
         .as_mut_any()
         .downcast_mut::<MutablePrimitiveArray<i256>>()
         .unwrap();
-    let from = from.map(|s| s.as_any().downcast_ref::<FixedLenStatistics>().unwrap());
 
     min.push(from.and_then(|s| s.min_value.as_deref().map(convert_i256)));
     max.push(from.and_then(|s| s.max_value.as_deref().map(convert_i256)));
@@ -81,7 +78,7 @@ pub(super) fn push_i256(
 }
 
 pub(super) fn push(
-    from: Option<&dyn ParquetStatistics>,
+    from: Option<&FixedLenStatistics>,
     min: &mut dyn MutableArray,
     max: &mut dyn MutableArray,
 ) -> PolarsResult<()> {
@@ -93,9 +90,10 @@ pub(super) fn push(
         .as_mut_any()
         .downcast_mut::<MutableFixedSizeBinaryArray>()
         .unwrap();
-    let from = from.map(|s| s.as_any().downcast_ref::<FixedLenStatistics>().unwrap());
+
     min.push(from.and_then(|s| s.min_value.as_ref()));
     max.push(from.and_then(|s| s.max_value.as_ref()));
+
     Ok(())
 }
 
@@ -104,7 +102,7 @@ fn convert_year_month(value: &[u8]) -> i32 {
 }
 
 pub(super) fn push_year_month(
-    from: Option<&dyn ParquetStatistics>,
+    from: Option<&FixedLenStatistics>,
     min: &mut dyn MutableArray,
     max: &mut dyn MutableArray,
 ) -> PolarsResult<()> {
@@ -116,7 +114,6 @@ pub(super) fn push_year_month(
         .as_mut_any()
         .downcast_mut::<MutablePrimitiveArray<i32>>()
         .unwrap();
-    let from = from.map(|s| s.as_any().downcast_ref::<FixedLenStatistics>().unwrap());
 
     min.push(from.and_then(|s| s.min_value.as_deref().map(convert_year_month)));
     max.push(from.and_then(|s| s.max_value.as_deref().map(convert_year_month)));
@@ -125,7 +122,7 @@ pub(super) fn push_year_month(
 }
 
 pub(super) fn push_days_ms(
-    from: Option<&dyn ParquetStatistics>,
+    from: Option<&FixedLenStatistics>,
     min: &mut dyn MutableArray,
     max: &mut dyn MutableArray,
 ) -> PolarsResult<()> {
@@ -137,7 +134,6 @@ pub(super) fn push_days_ms(
         .as_mut_any()
         .downcast_mut::<MutablePrimitiveArray<days_ms>>()
         .unwrap();
-    let from = from.map(|s| s.as_any().downcast_ref::<FixedLenStatistics>().unwrap());
 
     min.push(from.and_then(|s| s.min_value.as_deref().map(convert_days_ms)));
     max.push(from.and_then(|s| s.max_value.as_deref().map(convert_days_ms)));