feat: Add unstable LazyFrame.sink_iceberg (#26799)

Author: nameexhaustion

Cargo.lock

@@ -40,8 +40,8 @@ pub use datatypes::SchemaExtPl;
 pub use hashing::IdBuildHasher;
 use rayon::{ThreadPool, ThreadPoolBuilder};
 
-// A secret ID used to limit deserialization of raw pointers to those
-// generated by this instance of Polars.
+/// A secret ID used to limit deserialization of raw pointers to those
+/// generated by this instance of Polars.
 pub static PROCESS_ID: LazyLock<u128> = LazyLock::new(|| {
     let mut bytes = [0u8; 16];
     getrandom::fill(&mut bytes).unwrap();

crates/polars-core/src/lib.rs

@@ -28,7 +28,7 @@ polars-utils = { workspace = true }
 
 arrow = { workspace = true }
 bitflags = { workspace = true }
-blake3 = { workspace = true, optional = true }
+blake3 = { workspace = true }
 bytemuck = { workspace = true }
 bytes = { workspace = true, features = ["serde"] }
 chrono = { workspace = true, optional = true }
@@ -196,7 +196,7 @@ meta = []
 pivot = ["polars-core/rows", "polars-ops/pivot"]
 top_k = ["polars-ops/top_k"]
 semi_anti_join = ["polars-ops/semi_anti_join"]
-cse = ["blake3"]
+cse = []
 propagate_nans = ["polars-ops/propagate_nans"]
 coalesce = []
 fused = ["polars-ops/fused"]

crates/polars-plan/Cargo.toml

@@ -60,7 +60,7 @@
   "ExtraColumnsPolicy": "eb81efadce58eb148e658db4f2b5c1f38155d617431b81121043e9f9c21acd30",
   "Field": "dd95c2b6d7aa44004b900ef31fcf18e70f862d97488ef46c67b7c64c226b50d8",
   "Field2": "5a81d8772b4c18be0a0de8fc79d433d4b1d54b4008e211f2ca9217c15cf5611c",
-  "FileProviderType": "330497d01ca3e4ab79e838481eeab485603958a74534e214a9962e396c5b1770",
+  "FileProviderType": "1ddca3724d728cceed106b44479a92d317d8d612ef28e3e77797bef02d466090",
   "FileScanDsl": "ed77d7dc8af8845915011232bab3c31a90a1755c8e3897cd1473c435eb75accd",
   "FileSinkOptions": "edebcf5e3965add5e4fd1be14ca6bdddc55fa22e6e829dca04beb321de0c992c",
   "FileWriteFormat": "1a685aba7dd5d6c0aefc99a9060d1b57f166ea44ef57ad0d0d0c565dbabda811",
@@ -77,6 +77,7 @@
   "IcebergColumn": "171ff56c222358389754a7ff774eec6fc958478df2317720c63b4addc8f9a4c5",
   "IcebergColumnType": "e612983b0dfce78d172af2e4bb4726e3303ede09ea3c1de8ec40e12ee7922dac",
   "IcebergIdentityTransformedPartitionFields": "a9ea26367a6a3a97560aa9010f711a211cabfbffb6a318cb834ceccc672d3ae1",
+  "IcebergPathProvider": "20732e7c4d3e6386d7e2a9675973ea4527a4e474a194c9a837a73839acbef715",
   "IcebergSchema": "2341b76e5aca7780e28fcee6bd7a2650ce7a9df61e043b839dd3e74bd95efb3b",
   "IntDataTypeExpr": "cd66dcd9c44cdddd8864c0fe642e5fcef5263f6f142cce906011a0180e0fd161",
   "IntegerType": "2e73fb811a2830b8b114dfe914512bfa6031325da9ea5513875a6e49b6ab1a58",

crates/polars-plan/dsl-schema-hashes.json

@@ -2,7 +2,7 @@ use std::hash::Hash;
 use std::sync::Arc;
 
 use polars_core::frame::DataFrame;
-use polars_core::prelude::Column;
+use polars_core::prelude::{Column, DataType};
 use polars_error::PolarsResult;
 use polars_io::hive::HivePathFormatter;
 use polars_io::utils::file::Writeable;
@@ -29,6 +29,7 @@ pub type FileProviderFunction = PlanCallback<FileProviderArgs, FileProviderRetur
 #[derive(Clone, Debug, Hash, PartialEq)]
 pub enum FileProviderType {
     Hive(HivePathProvider),
+    Iceberg(IcebergPathProvider),
     Function(FileProviderFunction),
 }
 
@@ -40,10 +41,28 @@ pub struct HivePathProvider {
 }
 
 impl FileProviderType {
+    /// Get a mutable reference to the file part prefix for this file provider.
+    ///
+    /// File part prefixes are inserted after the partition prefix, before the file part number.
+    ///
+    /// # Returns
+    /// Returns `None` if this file provider does not support attaching file part prefixes.
+    pub fn file_part_prefix_mut(&mut self) -> Option<&mut String> {
+        use FileProviderType::*;
+
+        match self {
+            Iceberg(p) => Some(p.file_part_prefix_mut()),
+            Hive(_) | Function(_) => None,
+        }
+    }
+
     pub fn get_path_or_file(&self, args: FileProviderArgs) -> PolarsResult<FileProviderReturn> {
+        use FileProviderType::*;
+
         match self {
-            Self::Hive(v) => v.get_path(args).map(FileProviderReturn::Path),
-            Self::Function(v) => v.get_path_or_file(args),
+            Hive(p) => p.get_path(args).map(FileProviderReturn::Path),
+            Iceberg(p) => p.get_path(args).map(FileProviderReturn::Path),
+            Function(p) => p.get_path_or_file(args),
         }
     }
 }
@@ -59,22 +78,82 @@ impl HivePathProvider {
             partition_keys,
         } = args;
 
-        let mut partition_parts = String::new();
+        let mut path = String::new();
 
         let partition_keys: &[Column] = partition_keys.columns();
 
-        write!(
-            &mut partition_parts,
-            "{}",
-            HivePathFormatter::new(partition_keys)
-        )
-        .unwrap();
+        write!(&mut path, "{}", HivePathFormatter::new(partition_keys)).unwrap();
 
         assert!(index_in_partition <= 0xffff_ffff);
 
-        write!(&mut partition_parts, "{index_in_partition:08x}.{extension}").unwrap();
+        write!(&mut path, "{index_in_partition:08x}.{extension}").unwrap();
+
+        Ok(path)
+    }
+}
+
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
+#[derive(Clone, Debug, Hash, PartialEq)]
+pub struct IcebergPathProvider {
+    pub extension: PlSmallStr,
+    pub file_part_prefix: String,
+}
+
+impl IcebergPathProvider {
+    pub fn file_part_prefix_mut(&mut self) -> &mut String {
+        &mut self.file_part_prefix
+    }
+
+    /// # Panics
+    /// Panics if `self.file_part_prefix` is `None`.
+    pub fn get_path(&self, args: FileProviderArgs) -> PolarsResult<String> {
+        use std::fmt::Write;
+
+        let IcebergPathProvider {
+            extension,
+            file_part_prefix,
+        } = self;
+
+        assert!(!file_part_prefix.is_empty());
+
+        let FileProviderArgs {
+            index_in_partition,
+            partition_keys,
+        } = args;
+
+        let mut partition_keys_hash = None;
+
+        if partition_keys.width() != 0 {
+            let mut hasher = blake3::Hasher::new();
+
+            for column in partition_keys.columns() {
+                let column = column.cast(&DataType::String).unwrap();
+
+                let value = column.str().unwrap().get(0);
+
+                hasher.update(&[value.is_some() as u8]);
+                hasher.update(value.unwrap_or_default().as_bytes());
+            }
+
+            partition_keys_hash = Some(hasher.finalize().to_hex());
+        }
+
+        let partition_key_prefix: &str = partition_keys_hash.as_ref().map_or("", |x| &x[..32]);
+
+        let mut path = String::with_capacity(
+            partition_key_prefix.len() + file_part_prefix.len() + 8 + 1 + extension.len(),
+        );
+
+        assert!(index_in_partition <= 0xffff_ffff);
+
+        write!(
+            &mut path,
+            "{partition_key_prefix}{file_part_prefix}{index_in_partition:08x}.{extension}"
+        )
+        .unwrap();
 
-        Ok(partition_parts)
+        Ok(path)
     }
 }
 

crates/polars-plan/src/dsl/options/file_provider.rs

@@ -1,9 +1,11 @@
-use polars::prelude::file_provider::{FileProviderFunction, FileProviderType};
-use polars::prelude::{PartitionStrategy, PlRefPath, SinkDestination, SpecialEq};
+use polars::prelude::file_provider::{FileProviderFunction, FileProviderType, IcebergPathProvider};
+use polars::prelude::{PartitionStrategy, PlRefPath, PlSmallStr, SinkDestination, SpecialEq};
 use polars_utils::IdxSize;
 use polars_utils::python_function::PythonObject;
+use pyo3::exceptions::PyValueError;
 use pyo3::intern;
 use pyo3::prelude::*;
+use pyo3::pybacked::PyBackedStr;
 
 use crate::PyExpr;
 use crate::prelude::Wrap;
@@ -68,13 +70,40 @@ impl PyFileSinkDestination<'_> {
             PartitionStrategy::FileSize
         };
 
+        let file_path_provider = if let Some(file_path_provider) = file_path_provider {
+            let py = self.0.py();
+
+            Some(
+                match file_path_provider.getattr(py, intern!(py, "pl_path_provider_id")) {
+                    Ok(v) => match &*v.extract::<PyBackedStr>(py)? {
+                        "iceberg" => {
+                            let extension: PyBackedStr = file_path_provider
+                                .getattr(py, intern!(py, "extension"))?
+                                .extract(py)?;
+
+                            FileProviderType::Iceberg(IcebergPathProvider {
+                                extension: PlSmallStr::from_str(&extension),
+                                file_part_prefix: String::new(),
+                            })
+                        },
+                        id => {
+                            return Err(PyValueError::new_err(format!(
+                                "unknown pl_path_provider_id: '{id}'"
+                            )));
+                        },
+                    },
+                    Err(_) => FileProviderType::Function(FileProviderFunction::Python(
+                        SpecialEq::new(PythonObject(file_path_provider).into()),
+                    )),
+                },
+            )
+        } else {
+            None
+        };
+
         Ok(SinkDestination::Partitioned {
             base_path: base_path.0,
-            file_path_provider: file_path_provider.map(|x| {
-                FileProviderType::Function(FileProviderFunction::Python(SpecialEq::new(
-                    PythonObject(x).into(),
-                )))
-            }),
+            file_path_provider,
             partition_strategy,
             max_rows_per_file: max_rows_per_file.unwrap_or(IdxSize::MAX),
             approximate_bytes_per_file,

crates/polars-python/src/io/sink_output.rs

@@ -34,6 +34,7 @@ serde = { workspace = true, optional = true }
 serde_json = { workspace = true, optional = true }
 slotmap = { workspace = true }
 tokio = { workspace = true, features = ["sync"] }
+uuid = { workspace = true }
 
 arrow = { workspace = true }
 memchr = { workspace = true }

crates/polars-stream/Cargo.toml

@@ -21,7 +21,8 @@ pub struct FileProvider {
 impl FileProvider {
     pub async fn open_file(&self, args: FileProviderArgs) -> PolarsResult<Writeable> {
         let provided_path: String = match &self.provider_type {
-            FileProviderType::Hive(v) => v.get_path(args)?,
+            FileProviderType::Hive(p) => p.get_path(args)?,
+            FileProviderType::Iceberg(p) => p.get_path(args)?,
             FileProviderType::Function(f) => {
                 let f = f.clone();
 

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

@@ -55,7 +55,7 @@ pub fn start_partition_sink_pipeline(
 
     let PartitionedTarget {
         base_path,
-        file_path_provider,
+        mut file_path_provider,
         partitioner,
         hstack_keys,
         include_keys_in_file,
@@ -68,6 +68,13 @@ pub fn start_partition_sink_pipeline(
     let in_memory_exec_state = Arc::new(execution_state.in_memory_exec_state.clone());
     let io_metrics_is_some = io_metrics.is_some();
 
+    if let Some(file_part_prefix) = file_path_provider.file_part_prefix_mut() {
+        use std::fmt::Write as _;
+        let uuid = uuid::Uuid::new_v4();
+        let uuid = uuid.as_simple();
+        write!(file_part_prefix, "{uuid}").unwrap();
+    }
+
     let file_provider = Arc::new(FileProvider {
         base_path,
         cloud_options,

crates/polars-stream/src/nodes/io_sinks/pipeline_initialization/partition_by.rs

@@ -1,3 +1,7 @@
+from polars.io.iceberg._dataset import IcebergCatalogConfig
 from polars.io.iceberg.functions import scan_iceberg
 
-__all__ = ["scan_iceberg"]
+__all__ = [
+    "IcebergCatalogConfig",
+    "scan_iceberg",
+]