From 247f0b17042401551af77205bcf28e4e739887c8 Mon Sep 17 00:00:00 2001
From: Gijs Burghoorn <me@gburghoorn.com>
Date: Wed, 8 Jan 2025 15:09:05 +0100
Subject: [PATCH] perf: Make Parquet `verify_dict_indices` SIMD (#20623)

---
 crates/polars-parquet/Cargo.toml              |  1 +
 .../deserialize/dictionary_encoded/mod.rs     | 21 +++++++++++++++++++
 crates/polars-parquet/src/lib.rs              |  1 +
 3 files changed, 23 insertions(+)

diff --git a/crates/polars-parquet/Cargo.toml b/crates/polars-parquet/Cargo.toml
index 27dda4f55e14..b61a01f33b64 100644
--- a/crates/polars-parquet/Cargo.toml
+++ b/crates/polars-parquet/Cargo.toml
@@ -63,3 +63,4 @@ lz4_flex = ["dep:lz4_flex"]
 async = ["async-stream", "futures", "polars-parquet-format/async"]
 bloom_filter = ["xxhash-rust"]
 serde_types = ["serde"]
+simd = ["polars-compute/simd"]
diff --git a/crates/polars-parquet/src/arrow/read/deserialize/dictionary_encoded/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/dictionary_encoded/mod.rs
index e025fa2b1ddd..27952215897b 100644
--- a/crates/polars-parquet/src/arrow/read/deserialize/dictionary_encoded/mod.rs
+++ b/crates/polars-parquet/src/arrow/read/deserialize/dictionary_encoded/mod.rs
@@ -179,6 +179,27 @@ fn no_more_bitpacked_values() -> ParquetError {
 }
 
 #[inline(always)]
+#[cfg(feature = "simd")]
+fn verify_dict_indices(indices: &[u32; 32], dict_size: usize) -> ParquetResult<()> {
+    // You would think that the compiler can do this itself, but it does not always do this
+    // properly. So we help it a bit.
+
+    use std::simd::cmp::SimdPartialOrd;
+    use std::simd::u32x32;
+
+    let dict_size = u32x32::splat(dict_size as u32);
+    let indices = u32x32::from_slice(indices);
+
+    let is_invalid = indices.simd_ge(dict_size);
+    if is_invalid.any() {
+        Err(oob_dict_idx())
+    } else {
+        Ok(())
+    }
+}
+
+#[inline(always)]
+#[cfg(not(feature = "simd"))]
 fn verify_dict_indices(indices: &[u32; 32], dict_size: usize) -> ParquetResult<()> {
     let mut is_valid = true;
     for &idx in indices {
diff --git a/crates/polars-parquet/src/lib.rs b/crates/polars-parquet/src/lib.rs
index 04fc2f6211b7..c429e83ad328 100644
--- a/crates/polars-parquet/src/lib.rs
+++ b/crates/polars-parquet/src/lib.rs
@@ -1,3 +1,4 @@
+#![cfg_attr(feature = "simd", feature(portable_simd))]
 #![allow(clippy::len_without_is_empty)]
 pub mod arrow;
 pub use crate::arrow::{read, write};