Merge branch 'main' into refactor/upsample-recursion

Author: hallmason17

crates/polars-core/src/chunked_array/ops/arity.rs

@@ -556,7 +556,18 @@ where
     ca
 }
 
-#[inline]
+pub fn try_unary_to_series<T, F>(ca: &ChunkedArray<T>, op: F) -> PolarsResult<Series>
+where
+    T: PolarsDataType,
+    F: FnMut(&T::Array) -> PolarsResult<Box<dyn Array>>,
+{
+    let chunks = ca
+        .downcast_iter()
+        .map(op)
+        .collect::<PolarsResult<Vec<_>>>()?;
+    Series::try_from((ca.name().clone(), chunks))
+}
+
 pub fn binary_to_series<T, U, F>(
     lhs: &ChunkedArray<T>,
     rhs: &ChunkedArray<U>,
@@ -576,6 +587,25 @@ where
     Series::try_from((lhs.name().clone(), chunks))
 }
 
+pub fn try_binary_to_series<T, U, F>(
+    lhs: &ChunkedArray<T>,
+    rhs: &ChunkedArray<U>,
+    mut op: F,
+) -> PolarsResult<Series>
+where
+    T: PolarsDataType,
+    U: PolarsDataType,
+    F: FnMut(&T::Array, &U::Array) -> PolarsResult<Box<dyn Array>>,
+{
+    let (lhs, rhs) = align_chunks_binary(lhs, rhs);
+    let chunks = lhs
+        .downcast_iter()
+        .zip(rhs.downcast_iter())
+        .map(|(lhs_arr, rhs_arr)| op(lhs_arr, rhs_arr))
+        .collect::<PolarsResult<Vec<_>>>()?;
+    Series::try_from((lhs.name().clone(), chunks))
+}
+
 /// Applies a kernel that produces `ArrayRef` of the same type.
 ///
 /// # Safety

crates/polars-expr/src/expressions/gather.rs

@@ -2,7 +2,7 @@ use polars_core::chunked_array::cast::CastOptions;
 use polars_core::prelude::arity::unary_elementwise_values;
 use polars_core::prelude::*;
 use polars_ops::prelude::lst_get;
-use polars_ops::series::convert_to_unsigned_index;
+use polars_ops::series::convert_and_bound_index;
 use polars_utils::index::ToIdx;
 
 use super::*;
@@ -24,11 +24,8 @@ impl PhysicalExpr for GatherExpr {
     fn evaluate(&self, df: &DataFrame, state: &ExecutionState) -> PolarsResult<Column> {
         let series = self.phys_expr.evaluate(df, state)?;
         let idx = self.idx.evaluate(df, state)?;
-        let idx = convert_to_unsigned_index(
-            idx.as_materialized_series(),
-            series.len(),
-            self.null_on_oob,
-        )?;
+        let idx =
+            convert_and_bound_index(idx.as_materialized_series(), series.len(), self.null_on_oob)?;
         series.take(&idx)
     }
 

crates/polars-io/src/cloud/object_store_setup.rs

@@ -11,7 +11,7 @@ use polars_utils::{format_pl_smallstr, pl_serialize};
 use tokio::sync::RwLock;
 
 use super::{CloudLocation, CloudOptions, CloudType, PolarsObjectStore};
-use crate::cloud::CloudConfig;
+use crate::cloud::{CloudConfig, CloudRetryConfig};
 
 /// Object stores must be cached. Every object-store will do DNS lookups and
 /// get rate limited when querying the DNS (can take up to 5s).
@@ -38,19 +38,20 @@ fn path_and_creds_to_key(path: &PlPath, options: Option<&CloudOptions>) -> Vec<u
     // We include credentials as they can expire, so users will send new credentials for the same url.
     let cloud_options = options.map(
         |CloudOptions {
+             max_retries: _,
              // Destructure to ensure this breaks if anything changes.
-             max_retries,
              #[cfg(feature = "file_cache")]
              file_cache_ttl,
              config,
+             retry_config,
              #[cfg(feature = "cloud")]
              credential_provider,
          }| {
-            CloudOptions2 {
-                max_retries: *max_retries,
+            CloudOptionsKey {
                 #[cfg(feature = "file_cache")]
                 file_cache_ttl: *file_cache_ttl,
                 config: config.clone(),
+                retry_config: *retry_config,
                 #[cfg(feature = "cloud")]
                 credential_provider: credential_provider.as_ref().map_or(0, |x| x.func_addr()),
             }
@@ -75,18 +76,18 @@ fn path_and_creds_to_key(path: &PlPath, options: Option<&CloudOptions>) -> Vec<u
     #[cfg_attr(feature = "serde", derive(serde::Serialize))]
     struct CacheKey {
         url_base: PlSmallStr,
-        cloud_options: Option<CloudOptions2>,
+        cloud_options: Option<CloudOptionsKey>,
     }
 
     /// Variant of CloudOptions for serializing to a cache key. The credential
     /// provider is replaced by the function address.
     #[derive(Clone, Debug, PartialEq, Hash, Eq)]
     #[cfg_attr(feature = "serde", derive(serde::Serialize))]
-    struct CloudOptions2 {
-        max_retries: usize,
+    struct CloudOptionsKey {
         #[cfg(feature = "file_cache")]
         file_cache_ttl: u64,
         config: Option<CloudConfig>,
+        retry_config: CloudRetryConfig,
         #[cfg(feature = "cloud")]
         credential_provider: usize,
     }

crates/polars-io/src/cloud/options.rs

@@ -19,12 +19,11 @@ use object_store::azure::MicrosoftAzureBuilder;
 use object_store::gcp::GoogleCloudStorageBuilder;
 #[cfg(feature = "gcp")]
 pub use object_store::gcp::GoogleConfigKey;
-#[cfg(any(feature = "aws", feature = "gcp", feature = "azure"))]
-use object_store::{BackoffConfig, RetryConfig};
 use polars_error::*;
 #[cfg(feature = "aws")]
 use polars_utils::cache::LruCache;
 use polars_utils::pl_path::{CloudScheme, PlRefPath};
+use polars_utils::total_ord::TotalOrdWrap;
 #[cfg(feature = "http")]
 use reqwest::header::HeaderMap;
 #[cfg(feature = "serde")]
@@ -79,10 +78,12 @@ pub(crate) enum CloudConfig {
 #[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
 /// Options to connect to various cloud providers.
 pub struct CloudOptions {
-    pub max_retries: usize,
+    pub max_retries: usize, // TODO: Remove in a breaking DSL change
     #[cfg(feature = "file_cache")]
     pub file_cache_ttl: u64,
     pub(crate) config: Option<CloudConfig>,
+    #[cfg_attr(feature = "serde", serde(default))]
+    pub retry_config: CloudRetryConfig,
     #[cfg(feature = "cloud")]
     /// Note: In most cases you will want to access this via [`CloudOptions::initialized_credential_provider`]
     /// rather than directly.
@@ -102,6 +103,7 @@ impl CloudOptions {
             #[cfg(feature = "file_cache")]
             file_cache_ttl: get_env_file_cache_ttl(),
             config: None,
+            retry_config: CloudRetryConfig::default(),
             #[cfg(feature = "cloud")]
             credential_provider: None,
         });
@@ -110,6 +112,66 @@ impl CloudOptions {
     }
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Hash, Eq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
+pub struct CloudRetryConfig {
+    pub max_retries: usize,
+    pub retry_timeout: std::time::Duration,
+    pub retry_init_backoff: std::time::Duration,
+    pub retry_max_backoff: std::time::Duration,
+    pub retry_base_multiplier: TotalOrdWrap<f64>,
+}
+
+impl Default for CloudRetryConfig {
+    fn default() -> Self {
+        use std::time::Duration;
+
+        return Self {
+            max_retries: parse_env_var(2, "POLARS_CLOUD_MAX_RETRIES"),
+            retry_timeout: Duration::from_millis(parse_env_var(
+                10 * 1000,
+                "POLARS_CLOUD_RETRY_TIMEOUT_MS",
+            )),
+            retry_init_backoff: Duration::from_millis(parse_env_var(
+                100,
+                "POLARS_CLOUD_RETRY_INIT_BACKOFF_MS",
+            )),
+            retry_max_backoff: Duration::from_millis(parse_env_var(
+                15 * 1000,
+                "POLARS_CLOUD_RETRY_MAX_BACKOFF_MS",
+            )),
+            retry_base_multiplier: TotalOrdWrap(parse_env_var(
+                2.,
+                "POLARS_CLOUD_RETRY_BASE_MULTIPLIER",
+            )),
+        };
+
+        fn parse_env_var<T: FromStr>(default: T, name: &'static str) -> T {
+            std::env::var(name).map_or(default, |x| {
+                x.parse::<T>()
+                    .ok()
+                    .unwrap_or_else(|| panic!("invalid value for {name}: {x}"))
+            })
+        }
+    }
+}
+
+#[cfg(any(feature = "aws", feature = "gcp", feature = "azure"))]
+impl From<CloudRetryConfig> for object_store::RetryConfig {
+    fn from(value: CloudRetryConfig) -> Self {
+        object_store::RetryConfig {
+            backoff: object_store::BackoffConfig {
+                init_backoff: value.retry_init_backoff,
+                max_backoff: value.retry_max_backoff,
+                base: value.retry_base_multiplier.0,
+            },
+            max_retries: value.max_retries,
+            retry_timeout: value.retry_timeout,
+        }
+    }
+}
+
 #[cfg(feature = "http")]
 pub(crate) fn try_build_http_header_map_from_items_slice<S: AsRef<str>>(
     headers: &[(S, S)],
@@ -182,15 +244,6 @@ impl CloudType {
     }
 }
 
-#[cfg(any(feature = "aws", feature = "gcp", feature = "azure"))]
-fn get_retry_config(max_retries: usize) -> RetryConfig {
-    RetryConfig {
-        backoff: BackoffConfig::default(),
-        max_retries,
-        retry_timeout: std::time::Duration::from_secs(10),
-    }
-}
-
 pub static USER_AGENT: &str = concat!("polars", "/", env!("CARGO_PKG_VERSION"),);
 
 #[cfg(any(feature = "aws", feature = "gcp", feature = "azure", feature = "http"))]
@@ -253,9 +306,9 @@ fn read_config(
 }
 
 impl CloudOptions {
-    /// Set the maximum number of retries.
-    pub fn with_max_retries(mut self, max_retries: usize) -> Self {
-        self.max_retries = max_retries;
+    pub fn with_retry_config(mut self, retry_config: CloudRetryConfig) -> Self {
+        self.max_retries = retry_config.max_retries;
+        self.retry_config = retry_config;
         self
     }
 
@@ -399,7 +452,7 @@ impl CloudOptions {
             };
         };
 
-        let builder = builder.with_retry(get_retry_config(self.max_retries));
+        let builder = builder.with_retry(self.retry_config.into());
 
         let opt_credential_provider = match opt_credential_provider {
             #[cfg(feature = "python")]
@@ -473,7 +526,7 @@ impl CloudOptions {
 
         let builder = builder
             .with_url(url.to_string())
-            .with_retry(get_retry_config(self.max_retries));
+            .with_retry(self.retry_config.into());
 
         let builder =
             if let Some(v) = self.initialized_credential_provider(clear_cached_credentials)? {
@@ -535,7 +588,7 @@ impl CloudOptions {
 
         let builder = builder
             .with_url(url.to_string())
-            .with_retry(get_retry_config(self.max_retries));
+            .with_retry(self.retry_config.into());
 
         let builder = if let Some(v) = credential_provider {
             builder.with_credentials(v.into_gcp_provider())

crates/polars-io/src/parquet/read/mmap.rs

@@ -1,11 +1,13 @@
+use std::io::Cursor;
+
 use arrow::array::Array;
 use arrow::bitmap::Bitmap;
 use arrow::datatypes::Field;
 use polars_error::PolarsResult;
 use polars_parquet::read::{
     BasicDecompressor, ColumnChunkMetadata, Filter, PageReader, column_iter_to_arrays,
 };
-use polars_utils::mmap::{MemReader, MemSlice};
+use polars_utils::mmap::MemSlice;
 
 /// Store columns data in two scenarios:
 /// 1. a local memory mapped file
@@ -60,7 +62,7 @@ pub fn to_deserializer(
             // Advise fetching the data for the column chunk
             chunk.prefetch();
 
-            let pages = PageReader::new(MemReader::new(chunk), column_meta, vec![], usize::MAX);
+            let pages = PageReader::new(Cursor::new(chunk), column_meta, vec![], usize::MAX);
             (
                 BasicDecompressor::new(pages, vec![]),
                 &column_meta.descriptor().descriptor.primitive_type,

crates/polars-io/src/utils/compression.rs

@@ -1,9 +1,9 @@
 use std::cmp;
-use std::io::{Read, Write};
+use std::io::{Cursor, Read, Write};
 
 use polars_core::prelude::*;
 use polars_error::{feature_gated, to_compute_err};
-use polars_utils::mmap::{MemReader, MemSlice};
+use polars_utils::mmap::MemSlice;
 
 use crate::utils::file::{Writeable, WriteableTrait};
 use crate::utils::sync_on_close::SyncOnCloseType;
@@ -80,11 +80,11 @@ pub enum CompressedReader {
         offset: usize,
     },
     #[cfg(feature = "decompress")]
-    Gzip(flate2::bufread::MultiGzDecoder<MemReader>),
+    Gzip(flate2::bufread::MultiGzDecoder<Cursor<MemSlice>>),
     #[cfg(feature = "decompress")]
-    Zlib(flate2::bufread::ZlibDecoder<MemReader>),
+    Zlib(flate2::bufread::ZlibDecoder<Cursor<MemSlice>>),
     #[cfg(feature = "decompress")]
-    Zstd(zstd::Decoder<'static, MemReader>),
+    Zstd(zstd::Decoder<'static, Cursor<MemSlice>>),
 }
 
 impl CompressedReader {
@@ -95,15 +95,15 @@ impl CompressedReader {
             None => CompressedReader::Uncompressed { slice, offset: 0 },
             #[cfg(feature = "decompress")]
             Some(SupportedCompression::GZIP) => {
-                CompressedReader::Gzip(flate2::bufread::MultiGzDecoder::new(MemReader::new(slice)))
+                CompressedReader::Gzip(flate2::bufread::MultiGzDecoder::new(Cursor::new(slice)))
             },
             #[cfg(feature = "decompress")]
             Some(SupportedCompression::ZLIB) => {
-                CompressedReader::Zlib(flate2::bufread::ZlibDecoder::new(MemReader::new(slice)))
+                CompressedReader::Zlib(flate2::bufread::ZlibDecoder::new(Cursor::new(slice)))
             },
             #[cfg(feature = "decompress")]
             Some(SupportedCompression::ZSTD) => {
-                CompressedReader::Zstd(zstd::Decoder::with_buffer(MemReader::new(slice))?)
+                CompressedReader::Zstd(zstd::Decoder::with_buffer(Cursor::new(slice))?)
             },
             #[cfg(not(feature = "decompress"))]
             _ => panic!("activate 'decompress' feature"),
@@ -141,14 +141,16 @@ impl CompressedReader {
             CompressedReader::Uncompressed { slice, .. } => slice.len(),
             #[cfg(feature = "decompress")]
             CompressedReader::Gzip(reader) => {
-                reader.get_ref().total_len() * ESTIMATED_DEFLATE_RATIO
+                reader.get_ref().get_ref().len() * ESTIMATED_DEFLATE_RATIO
             },
             #[cfg(feature = "decompress")]
             CompressedReader::Zlib(reader) => {
-                reader.get_ref().total_len() * ESTIMATED_DEFLATE_RATIO
+                reader.get_ref().get_ref().len() * ESTIMATED_DEFLATE_RATIO
             },
             #[cfg(feature = "decompress")]
-            CompressedReader::Zstd(reader) => reader.get_ref().total_len() * ESTIMATED_ZSTD_RATIO,
+            CompressedReader::Zstd(reader) => {
+                reader.get_ref().get_ref().len() * ESTIMATED_ZSTD_RATIO
+            },
         }
     }