refactor(rust): start further use of polars-compute in polars-parquet (…

…#16788)
pola-rs · Jun 7, 2024 · 2e633dc · 2e633dc
1 parent 1df442f
commit 2e633dc
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 18 deletions.
diff --git a/crates/polars-parquet/src/arrow/write/binview/basic.rs b/crates/polars-parquet/src/arrow/write/binview/basic.rs
@@ -1,11 +1,12 @@
 use arrow::array::{Array, BinaryViewArray};
+use polars_compute::min_max::MinMaxKernel;
 use polars_error::PolarsResult;
 
 use crate::parquet::encoding::delta_bitpacked;
 use crate::parquet::schema::types::PrimitiveType;
 use crate::parquet::statistics::{BinaryStatistics, ParquetStatistics};
 use crate::read::schema::is_nullable;
-use crate::write::binary::{encode_non_null_values, ord_binary};
+use crate::write::binary::encode_non_null_values;
 use crate::write::utils::invalid_encoding;
 use crate::write::{utils, Encoding, Page, StatisticsOptions, WriteOptions};
 
@@ -77,7 +78,6 @@ pub fn array_to_page(
     .map(Page::Data)
 }
 
-// TODO! speed this up
 pub(crate) fn build_statistics(
     array: &BinaryViewArray,
     primitive_type: PrimitiveType,
@@ -89,23 +89,11 @@ pub(crate) fn build_statistics(
         distinct_count: None,
         max_value: options
             .max_value
-            .then(|| {
-                array
-                    .iter()
-                    .flatten()
-                    .max_by(|x, y| ord_binary(x, y))
-                    .map(|x| x.to_vec())
-            })
+            .then(|| array.max_propagate_nan_kernel().map(<[u8]>::to_vec))
             .flatten(),
         min_value: options
             .min_value
-            .then(|| {
-                array
-                    .iter()
-                    .flatten()
-                    .min_by(|x, y| ord_binary(x, y))
-                    .map(|x| x.to_vec())
-            })
+            .then(|| array.min_propagate_nan_kernel().map(<[u8]>::to_vec))
             .flatten(),
     }
     .serialize()

diff --git a/crates/polars-parquet/src/arrow/write/primitive/nested.rs b/crates/polars-parquet/src/arrow/write/primitive/nested.rs
@@ -11,13 +11,14 @@ use crate::parquet::page::DataPage;
 use crate::parquet::schema::types::PrimitiveType;
 use crate::parquet::types::NativeType;
 
-pub fn array_to_page<T, R>(
-    array: &PrimitiveArray<T>,
+pub fn array_to_page<'a, T, R>(
+    array: &'a PrimitiveArray<T>,
     options: WriteOptions,
     type_: PrimitiveType,
     nested: &[Nested],
 ) -> PolarsResult<DataPage>
 where
+    PrimitiveArray<T>: polars_compute::min_max::MinMaxKernel<Scalar<'a> = T>,
     T: ArrowNativeType,
     R: NativeType,
     T: num_traits::AsPrimitive<R>,