refactor(rust): Add streaming IO sink components (#25594)

Author: nameexhaustion

crates/polars-stream/src/nodes/io_sinks2/components/arg_sort.rs

@@ -0,0 +1,52 @@
+use std::sync::Arc;
+
+use futures::{StreamExt as _, TryStreamExt as _};
+use polars_core::frame::DataFrame;
+use polars_core::prelude::sort::arg_sort;
+use polars_core::prelude::{Column, IdxArr, SortMultipleOptions};
+use polars_error::PolarsResult;
+use polars_expr::state::ExecutionState;
+
+use crate::async_executor::TaskPriority;
+use crate::async_primitives::opt_spawned_future::parallelize_first_to_local;
+use crate::expression::StreamExpr;
+
+#[derive(Clone)]
+pub struct ArgSortBy {
+    pub by: Arc<[StreamExpr]>,
+    pub sort_options: SortMultipleOptions,
+    pub in_memory_exec_state: Arc<ExecutionState>,
+}
+
+impl ArgSortBy {
+    pub async fn arg_sort_by_par(self, df: &Arc<DataFrame>) -> PolarsResult<IdxArr> {
+        let ArgSortBy {
+            by,
+            sort_options,
+            in_memory_exec_state,
+        } = self;
+
+        let sort_by_cols: Vec<Column> = futures::stream::iter(parallelize_first_to_local(
+            TaskPriority::Low,
+            (0..by.len()).map(|i| {
+                let df = Arc::clone(df);
+                let by = Arc::clone(&by);
+                let in_memory_exec_state = Arc::clone(&in_memory_exec_state);
+
+                async move {
+                    by[i]
+                        .evaluate(&df, in_memory_exec_state.as_ref())
+                        .await
+                        .map(|c| c.rechunk())
+                }
+            }),
+        ))
+        .then(|x| x)
+        .try_collect()
+        .await?;
+
+        Ok(arg_sort(&sort_by_cols, sort_options)?
+            .downcast_as_array()
+            .clone())
+    }
+}

crates/polars-stream/src/nodes/io_sinks2/components/error_capture.rs

@@ -0,0 +1,64 @@
+use std::any::Any;
+use std::panic::AssertUnwindSafe;
+
+use futures::FutureExt;
+use polars_error::{PolarsError, PolarsResult};
+
+/// Utility to capture errors and propagate them to an associated [`ErrorHandle`].
+#[derive(Clone)]
+pub struct ErrorCapture {
+    tx: tokio::sync::mpsc::Sender<ErrorMessage>,
+}
+
+impl ErrorCapture {
+    pub fn new() -> (Self, ErrorHandle) {
+        let (tx, rx) = tokio::sync::mpsc::channel(1);
+        (Self { tx }, ErrorHandle { rx })
+    }
+
+    /// Wraps a future such that its error result is sent to the associated [`ErrorHandle`].
+    pub async fn wrap_future<F, O>(self, fut: F)
+    where
+        F: Future<Output = PolarsResult<O>>,
+    {
+        let err: Result<(), tokio::sync::mpsc::error::TrySendError<ErrorMessage>> =
+            match AssertUnwindSafe(fut).catch_unwind().await {
+                Ok(Ok(_)) => return,
+                Ok(Err(err)) => self.tx.try_send(ErrorMessage::Error(err)),
+                Err(panic) => self.tx.try_send(ErrorMessage::Panic(panic)),
+            };
+        drop(err);
+    }
+}
+
+enum ErrorMessage {
+    Error(PolarsError),
+    Panic(Box<dyn Any + Send + 'static>),
+}
+
+/// Handle to await the completion of multiple tasks. Propagates error results
+/// and resumes unwinds when joined.
+pub struct ErrorHandle {
+    rx: tokio::sync::mpsc::Receiver<ErrorMessage>,
+}
+
+impl ErrorHandle {
+    pub fn has_errored(&self) -> bool {
+        !self.rx.is_empty()
+    }
+
+    /// Block until either an error is received, or all [`ErrorCapture`]s associated with this
+    /// handle are dropped (i.e. successful completion of all wrapped futures).
+    ///
+    /// # Panics
+    /// If a panic is received, this will resume unwinding.
+    pub async fn join(self) -> PolarsResult<()> {
+        let ErrorHandle { mut rx } = self;
+
+        match rx.recv().await {
+            None => Ok(()),
+            Some(ErrorMessage::Error(e)) => Err(e),
+            Some(ErrorMessage::Panic(panic)) => std::panic::resume_unwind(panic),
+        }
+    }
+}

crates/polars-stream/src/nodes/io_sinks2/components/exclude_keys_projection.rs

@@ -0,0 +1,27 @@
+use std::sync::Arc;
+
+#[derive(Clone)]
+pub enum ExcludeKeysProjection {
+    /// Project these indices
+    Indices(Arc<[usize]>),
+    /// Project this many columns from the left.
+    Width(usize),
+}
+
+impl ExcludeKeysProjection {
+    pub fn iter_indices(&self) -> impl ExactSizeIterator<Item = usize> {
+        let (indices, end) = match self {
+            Self::Indices(indices) => (indices.as_ref(), indices.len()),
+            Self::Width(width) => (&[][..], *width),
+        };
+
+        (0..end).map(|i| if indices.is_empty() { i } else { indices[i] })
+    }
+
+    pub fn len(&self) -> usize {
+        match self {
+            Self::Indices(indices) => indices.len(),
+            Self::Width(len) => *len,
+        }
+    }
+}

crates/polars-stream/src/nodes/io_sinks2/components/file_provider.rs

@@ -0,0 +1,114 @@
+use std::fmt::Write;
+use std::sync::Arc;
+
+use polars_core::prelude::{Column, DataType};
+use polars_error::PolarsResult;
+use polars_io::cloud::CloudOptions;
+use polars_io::pl_async;
+use polars_io::utils::HIVE_VALUE_ENCODE_CHARSET;
+use polars_io::utils::file::Writeable;
+use polars_plan::dsl::sink2::{FileProviderReturn, FileProviderType};
+use polars_plan::prelude::sink2::FileProviderArgs;
+use polars_utils::plpath::PlPath;
+
+pub struct FileProvider {
+    pub base_path: PlPath,
+    pub cloud_options: Option<Arc<CloudOptions>>,
+    pub provider_type: FileProviderType,
+}
+
+impl FileProvider {
+    pub async fn open_file(&self, args: FileProviderArgs) -> PolarsResult<Writeable> {
+        let provided_path: String = match &self.provider_type {
+            FileProviderType::Hive { extension } => {
+                let FileProviderArgs {
+                    index_in_partition,
+                    partition_keys,
+                } = args;
+
+                let mut partition_parts = String::new();
+
+                let partition_keys: &[Column] = partition_keys.get_columns();
+
+                write!(
+                    &mut partition_parts,
+                    "{}",
+                    HivePathFormatter::new(partition_keys)
+                )
+                .unwrap();
+
+                assert!(index_in_partition <= 0xffff_ffff);
+
+                write!(&mut partition_parts, "{index_in_partition:08x}.{extension}").unwrap();
+
+                partition_parts
+            },
+
+            FileProviderType::Function(f) => {
+                let f = f.clone();
+
+                let out = pl_async::get_runtime()
+                    .spawn_blocking(move || f.get_file(args))
+                    .await
+                    .unwrap()?;
+
+                match out {
+                    FileProviderReturn::Path(p) => p,
+                    FileProviderReturn::Writeable(v) => return Ok(v),
+                }
+            },
+
+            FileProviderType::Legacy(_) => unreachable!(),
+        };
+
+        let path = self.base_path.as_ref().join(&provided_path);
+        let path = path.as_ref();
+
+        if let Some(path) = path.as_local_path().and_then(|p| p.parent()) {
+            // Ignore errors from directory creation - the `Writeable::try_new()` below will raise
+            // appropriate errors.
+            let _ = tokio::fs::DirBuilder::new()
+                .recursive(true)
+                .create(path)
+                .await;
+        }
+
+        Writeable::try_new(path, self.cloud_options.as_deref())
+    }
+}
+
+/// # Panics
+/// The `Display` impl of this will panic if a column has non-unit length.
+pub struct HivePathFormatter<'a> {
+    keys: &'a [Column],
+}
+
+impl<'a> HivePathFormatter<'a> {
+    pub fn new(keys: &'a [Column]) -> Self {
+        Self { keys }
+    }
+}
+
+impl std::fmt::Display for HivePathFormatter<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        for column in self.keys {
+            assert_eq!(column.len(), 1);
+            let column = column.cast(&DataType::String).unwrap();
+
+            let key = column.name();
+            let value = percent_encoding::percent_encode(
+                column
+                    .str()
+                    .unwrap()
+                    .get(0)
+                    .unwrap_or("__HIVE_DEFAULT_PARTITION__")
+                    .as_bytes(),
+                HIVE_VALUE_ENCODE_CHARSET,
+            );
+
+            write!(f, "{key}={value}/")?
+        }
+
+        Ok(())
+    }
+}

crates/polars-stream/src/nodes/io_sinks2/components/file_sink.rs

@@ -0,0 +1,21 @@
+use polars_error::PolarsResult;
+
+use crate::async_executor;
+use crate::async_primitives::connector;
+use crate::nodes::io_sinks2::components::sink_morsel::SinkMorsel;
+use crate::nodes::io_sinks2::components::size::RowCountAndSize;
+
+pub type FileSinkPermit = tokio::sync::OwnedSemaphorePermit;
+
+pub struct FileSinkTaskData {
+    pub morsel_tx: connector::Sender<SinkMorsel>,
+    pub start_position: RowCountAndSize,
+    pub task_handle: async_executor::JoinHandle<PolarsResult<FileSinkPermit>>,
+}
+
+impl FileSinkTaskData {
+    /// Signals to the writer to close, and returns its task handle.
+    pub fn close(self) -> async_executor::JoinHandle<PolarsResult<FileSinkPermit>> {
+        self.task_handle
+    }
+}

crates/polars-stream/src/nodes/io_sinks2/components/hstack_columns.rs

@@ -0,0 +1,67 @@
+use std::sync::Arc;
+
+use polars_core::prelude::Column;
+use polars_core::schema::Schema;
+use polars_utils::marked_usize::MarkedUsize;
+
+/// Applies a `with_columns()` operation with pre-computed indices.
+#[derive(Clone)]
+pub struct HStackColumns {
+    gather_indices: Arc<[MarkedUsize]>,
+}
+
+impl HStackColumns {
+    /// Note:
+    /// * Dtypes of the schemas are unused.
+    pub fn new(output_schema: &Schema, schema_left: &Schema, schema_right: &Schema) -> Self {
+        assert!(schema_left.len() <= MarkedUsize::UNMARKED_MAX);
+        assert!(schema_right.len() <= MarkedUsize::UNMARKED_MAX);
+
+        let gather_indices: Arc<[MarkedUsize]> = output_schema
+            .iter_names()
+            .map(|name| {
+                if let Some((idx, ..)) = schema_right.get_full(name) {
+                    MarkedUsize::new(idx, true)
+                } else {
+                    MarkedUsize::new(schema_left.get_full(name).unwrap().0, false)
+                }
+            })
+            .collect();
+
+        Self { gather_indices }
+    }
+
+    #[expect(unused)]
+    pub fn output_width(&self) -> usize {
+        self.gather_indices.len()
+    }
+
+    /// Broadcasts unit-length columns from the RHS.
+    pub fn hstack_columns_broadcast(
+        &self,
+        height: usize,
+        cols_left: &[Column],
+        cols_right: &[Column],
+    ) -> Vec<Column> {
+        self.gather_indices
+            .iter()
+            .copied()
+            .map(|mi| {
+                let i = mi.idx();
+
+                if mi.marked() {
+                    let c = &cols_right[i];
+
+                    if c.len() != height {
+                        assert_eq!(c.len(), 1);
+                        c.new_from_index(0, height)
+                    } else {
+                        c.clone()
+                    }
+                } else {
+                    cols_left[i].clone()
+                }
+            })
+            .collect()
+    }
+}

crates/polars-stream/src/nodes/io_sinks2/components/mod.rs

@@ -0,0 +1,8 @@
+pub mod arg_sort;
+pub mod error_capture;
+pub mod exclude_keys_projection;
+pub mod file_provider;
+pub mod file_sink;
+pub mod hstack_columns;
+pub mod sink_morsel;
+pub mod size;

crates/polars-stream/src/nodes/io_sinks2/components/sink_morsel.rs

@@ -0,0 +1,29 @@
+use polars_core::frame::DataFrame;
+
+pub type SinkMorselPermit = tokio::sync::OwnedSemaphorePermit;
+
+/// In-flight morsel in the IO sink. Holds a permit against a semaphore that restricts
+/// the total number of sink morsels in memory.
+pub struct SinkMorsel {
+    df: DataFrame,
+    /// Should only be dropped once the data associated with this morsel has been dropped from memory.
+    permit: SinkMorselPermit,
+}
+
+impl SinkMorsel {
+    pub fn new(df: DataFrame, permit: SinkMorselPermit) -> Self {
+        Self { df, permit }
+    }
+
+    pub fn into_inner(self) -> (DataFrame, SinkMorselPermit) {
+        (self.df, self.permit)
+    }
+
+    pub fn df(&self) -> &DataFrame {
+        &self.df
+    }
+
+    pub fn df_mut(&mut self) -> &mut DataFrame {
+        &mut self.df
+    }
+}