refactor(rust): replace exact instances of Option/Result combinators (#…

…6088)
pola-rs · Jan 7, 2023 · 8150e71 · 8150e71
1 parent 822e8c9
commit 8150e71
Show file tree

Hide file tree

Showing 13 changed files with 91 additions and 126 deletions.
diff --git a/polars/polars-core/src/chunked_array/iterator/par/list.rs b/polars/polars-core/src/chunked_array/iterator/par/list.rs
@@ -14,19 +14,16 @@ unsafe fn idx_to_array(idx: usize, arr: &ListArray<i64>, dtype: &DataType) -> Op
 impl ListChunked {
     // Get a parallel iterator over the [`Series`] in this [`ListChunked`].
     pub fn par_iter(&self) -> impl ParallelIterator<Item = Option<Series>> + '_ {
-        self.chunks
-            .par_iter()
-            .map(move |arr| {
-                let dtype = self.inner_dtype();
-                // Safety:
-                // guarded by the type system
-                let arr = &**arr;
-                let arr = unsafe { &*(arr as *const dyn Array as *const ListArray<i64>) };
-                (0..arr.len())
-                    .into_par_iter()
-                    .map(move |idx| unsafe { idx_to_array(idx, arr, &dtype) })
-            })
-            .flatten()
+        self.chunks.par_iter().flat_map(move |arr| {
+            let dtype = self.inner_dtype();
+            // Safety:
+            // guarded by the type system
+            let arr = &**arr;
+            let arr = unsafe { &*(arr as *const dyn Array as *const ListArray<i64>) };
+            (0..arr.len())
+                .into_par_iter()
+                .map(move |idx| unsafe { idx_to_array(idx, arr, &dtype) })
+        })
     }
 
     // Get an indexed parallel iterator over the [`Series`] in this [`ListChunked`].

diff --git a/polars/polars-core/src/chunked_array/iterator/par/utf8.rs b/polars/polars-core/src/chunked_array/iterator/par/utf8.rs
@@ -24,17 +24,14 @@ impl Utf8Chunked {
     }
 
     pub fn par_iter(&self) -> impl ParallelIterator<Item = Option<&str>> + '_ {
-        self.chunks
-            .par_iter()
-            .map(move |arr| {
-                // Safety:
-                // guarded by the type system
-                let arr = &**arr;
-                let arr = unsafe { &*(arr as *const dyn Array as *const Utf8Array<i64>) };
-                (0..arr.len())
-                    .into_par_iter()
-                    .map(move |idx| unsafe { idx_to_str(idx, arr) })
-            })
-            .flatten()
+        self.chunks.par_iter().flat_map(move |arr| {
+            // Safety:
+            // guarded by the type system
+            let arr = &**arr;
+            let arr = unsafe { &*(arr as *const dyn Array as *const Utf8Array<i64>) };
+            (0..arr.len())
+                .into_par_iter()
+                .map(move |idx| unsafe { idx_to_str(idx, arr) })
+        })
     }
 }
diff --git a/polars/polars-core/src/frame/asof_join/groups.rs b/polars/polars-core/src/frame/asof_join/groups.rs
@@ -280,7 +280,7 @@ where
             .zip(offsets)
             // probes_hashes: Vec<u64> processed by this thread
             // offset: offset index
-            .map(|(vals_left, offset)| {
+            .flat_map(|(vals_left, offset)| {
                 // local reference
                 let hash_tbls = &hash_tbls;
 
@@ -321,7 +321,6 @@ where
                 });
                 results
             })
-            .flatten()
             .collect()
     }))
 }
@@ -402,7 +401,7 @@ where
             .zip(offsets)
             // probes_hashes: Vec<u64> processed by this thread
             // offset: offset index
-            .map(|(vals_left, offset)| {
+            .flat_map(|(vals_left, offset)| {
                 // local reference
                 let hash_tbls = &hash_tbls;
 
@@ -443,7 +442,6 @@ where
                 });
                 results
             })
-            .flatten()
             .collect()
     })
 }
@@ -514,7 +512,7 @@ where
         probe_hashes
             .into_par_iter()
             .zip(offsets)
-            .map(|(probe_hashes, offset)| {
+            .flat_map(|(probe_hashes, offset)| {
                 // local reference
                 let hash_tbls = &hash_tbls;
 
@@ -566,7 +564,6 @@ where
 
                 results
             })
-            .flatten()
             .collect()
     })
 }

diff --git a/polars/polars-core/src/frame/hash_join/multiple_keys.rs b/polars/polars-core/src/frame/hash_join/multiple_keys.rs
@@ -200,7 +200,7 @@ pub fn _inner_join_multiple_keys(
         probe_hashes
             .into_par_iter()
             .zip(offsets)
-            .map(|(probe_hashes, offset)| {
+            .flat_map(|(probe_hashes, offset)| {
                 // local reference
                 let hash_tbls = &hash_tbls;
                 let mut results =
@@ -233,7 +233,6 @@ pub fn _inner_join_multiple_keys(
 
                 results
             })
-            .flatten()
             .unzip()
     })
 }
@@ -414,7 +413,7 @@ pub(crate) fn semi_anti_join_multiple_keys_impl<'a>(
         probe_hashes
             .into_par_iter()
             .zip(offsets)
-            .map(move |(probe_hashes, offset)| {
+            .flat_map(move |(probe_hashes, offset)| {
                 // local reference
                 let hash_tbls = &hash_tbls;
                 let mut results =
@@ -448,7 +447,6 @@ pub(crate) fn semi_anti_join_multiple_keys_impl<'a>(
 
                 results
             })
-            .flatten()
     })
 }
 

diff --git a/polars/polars-core/src/frame/hash_join/single_keys_inner.rs b/polars/polars-core/src/frame/hash_join/single_keys_inner.rs
@@ -52,7 +52,7 @@ where
         probe
             .into_par_iter()
             .zip(offsets)
-            .map(|(probe, offset)| {
+            .flat_map(|(probe, offset)| {
                 let probe = probe.as_ref();
                 // local reference
                 let hash_tbls = &hash_tbls;
@@ -82,7 +82,6 @@ where
 
                 results
             })
-            .flatten()
             .unzip()
     })
 }
diff --git a/polars/polars-core/src/frame/hash_join/single_keys_semi_anti.rs b/polars/polars-core/src/frame/hash_join/single_keys_semi_anti.rs
@@ -56,7 +56,7 @@ where
             .zip(offsets)
             // probes_hashes: Vec<u64> processed by this thread
             // offset: offset index
-            .map(move |(probe, offset)| {
+            .flat_map(move |(probe, offset)| {
                 // local reference
                 let hash_sets = &hash_sets;
                 let probe = probe.as_ref();
@@ -83,7 +83,6 @@ where
                 });
                 results
             })
-            .flatten()
     })
 }
 

diff --git a/polars/polars-core/src/frame/mod.rs b/polars/polars-core/src/frame/mod.rs
@@ -1537,10 +1537,7 @@ impl DataFrame {
     fn select_mut(&mut self, name: &str) -> Option<&mut Series> {
         let opt_idx = self.find_idx_by_name(name);
 
-        match opt_idx {
-            Some(idx) => self.select_at_idx_mut(idx),
-            None => None,
-        }
+        opt_idx.and_then(|idx| self.select_at_idx_mut(idx))
     }
 
     /// Does a filter but splits thread chunks vertically instead of horizontally

diff --git a/polars/polars-core/src/utils/supertype.rs b/polars/polars-core/src/utils/supertype.rs
@@ -3,12 +3,11 @@ use super::*;
 /// Given two datatypes, determine the supertype that both types can safely be cast to
 #[cfg(feature = "private")]
 pub fn try_get_supertype(l: &DataType, r: &DataType) -> PolarsResult<DataType> {
-    match get_supertype(l, r) {
-        Some(dt) => Ok(dt),
-        None => Err(PolarsError::ComputeError(
+    get_supertype(l, r).ok_or_else(|| {
+        PolarsError::ComputeError(
             format!("Failed to determine supertype of {l:?} and {r:?}").into(),
-        )),
-    }
+        )
+    })
 }
 
 /// Given two datatypes, determine the supertype that both types can safely be cast to
@@ -270,10 +269,7 @@ pub fn get_supertype(l: &DataType, r: &DataType) -> Option<DataType> {
         }
     }
 
-    match inner(l, r) {
-        Some(dt) => Some(dt),
-        None => inner(r, l),
-    }
+    inner(l, r).or_else(|| inner(r, l))
 }
 
 #[cfg(feature = "dtype-struct")]

diff --git a/polars/polars-io/src/ndjson_core/buffer.rs b/polars/polars-io/src/ndjson_core/buffer.rs
@@ -151,13 +151,12 @@ where
         Value::String(s) => s,
         _ => return None,
     };
-    match infer_pattern_single(val) {
-        None => None,
-        Some(pattern) => match DatetimeInfer::<T::Native>::try_from(pattern) {
+    infer_pattern_single(val).and_then(|pattern| {
+        match DatetimeInfer::<T::Native>::try_from(pattern) {
             Ok(mut infer) => infer.parse(val),
             Err(_) => None,
-        },
-    }
+        }
+    })
 }
 
 #[cfg(feature = "dtype-struct")]

diff --git a/polars/polars-io/src/parquet/write.rs b/polars/polars-io/src/parquet/write.rs
@@ -261,7 +261,7 @@ fn create_serializer<'a>(
         .par_iter()
         .zip(fields)
         .zip(encodings)
-        .map(move |((array, type_), encoding)| {
+        .flat_map(move |((array, type_), encoding)| {
             let encoded_columns =
                 array_to_columns(array, type_.clone(), options, encoding).unwrap();
 
@@ -288,7 +288,6 @@ fn create_serializer<'a>(
                 })
                 .collect::<Vec<_>>()
         })
-        .flatten()
         .collect::<Vec<_>>();
 
     let row_group = DynIter::new(columns.into_iter());

diff --git a/polars/polars-lazy/polars-plan/src/global.rs b/polars/polars-lazy/polars-plan/src/global.rs
@@ -5,10 +5,7 @@ thread_local! {pub static FETCH_ROWS: Cell<Option<usize>> = Cell::new(None)}
 
 pub fn _set_n_rows_for_scan(n_rows: Option<usize>) -> Option<usize> {
     let fetch_rows = FETCH_ROWS.with(|fetch_rows| fetch_rows.get());
-    match fetch_rows {
-        None => n_rows,
-        Some(n) => Some(n),
-    }
+    fetch_rows.or(n_rows)
 }
 
 pub fn _is_fetch_query() -> bool {

diff --git a/polars/polars-time/src/chunkedarray/utf8/infer.rs b/polars/polars-time/src/chunkedarray/utf8/infer.rs
@@ -160,24 +160,22 @@ fn transform_datetime_ns(val: &str, fmt: &str) -> Option<i64> {
     let out = NaiveDateTime::parse_from_str(val, fmt)
         .ok()
         .map(datetime_to_timestamp_ns);
-    match out {
-        Some(out) => Some(out),
-        None => NaiveDate::parse_from_str(val, fmt)
+    out.or_else(|| {
+        NaiveDate::parse_from_str(val, fmt)
             .ok()
-            .map(|nd| datetime_to_timestamp_ns(nd.and_hms_opt(0, 0, 0).unwrap())),
-    }
+            .map(|nd| datetime_to_timestamp_ns(nd.and_hms_opt(0, 0, 0).unwrap()))
+    })
 }
 
 fn transform_datetime_us(val: &str, fmt: &str) -> Option<i64> {
     let out = NaiveDateTime::parse_from_str(val, fmt)
         .ok()
         .map(datetime_to_timestamp_us);
-    match out {
-        Some(out) => Some(out),
-        None => NaiveDate::parse_from_str(val, fmt)
+    out.or_else(|| {
+        NaiveDate::parse_from_str(val, fmt)
             .ok()
-            .map(|nd| datetime_to_timestamp_us(nd.and_hms_opt(0, 0, 0).unwrap())),
-    }
+            .map(|nd| datetime_to_timestamp_us(nd.and_hms_opt(0, 0, 0).unwrap()))
+    })
 }
 
 fn transform_datetime_us_bytes(val: &[u8], fmt: &[u8], fmt_len: u16) -> Option<i64> {
@@ -188,20 +186,16 @@ fn transform_datetime_ms(val: &str, fmt: &str) -> Option<i64> {
     let out = NaiveDateTime::parse_from_str(val, fmt)
         .ok()
         .map(datetime_to_timestamp_ms);
-    match out {
-        Some(out) => Some(out),
-        None => NaiveDate::parse_from_str(val, fmt)
+    out.or_else(|| {
+        NaiveDate::parse_from_str(val, fmt)
             .ok()
-            .map(|nd| datetime_to_timestamp_ms(nd.and_hms_opt(0, 0, 0).unwrap())),
-    }
+            .map(|nd| datetime_to_timestamp_ms(nd.and_hms_opt(0, 0, 0).unwrap()))
+    })
 }
 
 pub fn infer_pattern_single(val: &str) -> Option<Pattern> {
     // Dates come first, because we see datetimes as superset of dates
-    match infer_pattern_date_single(val) {
-        Some(pat) => Some(pat),
-        None => infer_pattern_datetime_single(val),
-    }
+    infer_pattern_date_single(val).or_else(|| infer_pattern_datetime_single(val))
 }
 
 fn infer_pattern_datetime_single(val: &str) -> Option<Pattern> {