rebase

Author: nameexhaustion

crates/polars-plan/src/dsl/builder_dsl.rs

@@ -34,7 +34,7 @@ impl DslBuilder {
         })?;
 
         Ok(DslPlan::Scan {
-            sources: ScanSources::Buffers(Arc::default()),
+            sources: ScanSources::default(),
             file_info: Some(FileInfo {
                 schema: schema.clone(),
                 reader_schema: Some(either::Either::Right(schema)),
@@ -107,6 +107,25 @@ impl DslBuilder {
         .into())
     }
 
+    #[cfg(feature = "python")]
+    pub fn scan_python_dataset(
+        dataset_object: polars_utils::python_function::PythonObject,
+    ) -> DslBuilder {
+        use super::python_dataset::PythonDatasetProvider;
+
+        DslPlan::Scan {
+            sources: ScanSources::default(),
+            file_info: None,
+            unified_scan_args: Default::default(),
+            scan_type: Box::new(FileScan::PythonDataset {
+                dataset_object: Arc::new(PythonDatasetProvider::new(dataset_object)),
+                cached_ir: Default::default(),
+            }),
+            cached_ir: Default::default(),
+        }
+        .into()
+    }
+
     pub fn cache(self) -> Self {
         let input = Arc::new(self.0);
         let id = input.as_ref() as *const DslPlan as usize;

crates/polars-plan/src/dsl/file_scan.rs renamed to crates/polars-plan/src/dsl/file_scan/mod.rs

@@ -1,4 +1,5 @@
 use std::hash::Hash;
+use std::sync::Mutex;
 
 use polars_io::cloud::CloudOptions;
 #[cfg(feature = "csv")]
@@ -17,6 +18,11 @@ use strum_macros::IntoStaticStr;
 
 use super::*;
 
+#[cfg(feature = "python")]
+pub mod python_dataset;
+#[cfg(feature = "python")]
+pub use python_dataset::{DATASET_PROVIDER_VTABLE, PythonDatasetProviderVTable};
+
 bitflags::bitflags! {
     #[derive(Debug, Clone, Copy, PartialEq, Eq)]
     pub struct ScanFlags : u32 {
@@ -48,6 +54,14 @@ pub enum FileScan {
         metadata: Option<Arc<arrow::io::ipc::read::FileMetadata>>,
     },
 
+    #[cfg(feature = "python")]
+    PythonDataset {
+        dataset_object: Arc<python_dataset::PythonDatasetProvider>,
+
+        #[cfg_attr(feature = "serde", serde(skip, default))]
+        cached_ir: Arc<Mutex<Option<ExpandedDataset>>>,
+    },
+
     #[cfg_attr(feature = "serde", serde(skip))]
     Anonymous {
         options: Arc<AnonymousScanOptions>,
@@ -199,10 +213,20 @@ mod _file_scan_eq_hash {
             metadata: Option<usize>,
         },
 
+        #[cfg(feature = "python")]
+        PythonDataset {
+            dataset_object: usize,
+            cached_ir: usize,
+        },
+
         Anonymous {
             options: &'a crate::dsl::AnonymousScanOptions,
             function: usize,
         },
+
+        /// Variant to ensure the lifetime is used regardless of feature gate combination.
+        #[expect(unused)]
+        Phantom(&'a ()),
     }
 
     impl<'a> From<&'a FileScan> for FileScanEqHashWrap<'a> {
@@ -226,6 +250,15 @@ mod _file_scan_eq_hash {
                     metadata: metadata.as_ref().map(arc_as_ptr),
                 },
 
+                #[cfg(feature = "python")]
+                FileScan::PythonDataset {
+                    dataset_object,
+                    cached_ir,
+                } => FileScanEqHashWrap::PythonDataset {
+                    dataset_object: arc_as_ptr(dataset_object),
+                    cached_ir: arc_as_ptr(cached_ir),
+                },
+
                 FileScan::Anonymous { options, function } => FileScanEqHashWrap::Anonymous {
                     options,
                     function: arc_as_ptr(function),

crates/polars-plan/src/dsl/file_scan/python_dataset.rs

@@ -0,0 +1,65 @@
+use std::fmt::Debug;
+use std::sync::OnceLock;
+
+use polars_core::error::PolarsResult;
+use polars_core::schema::SchemaRef;
+use polars_utils::pl_str::PlSmallStr;
+use polars_utils::python_function::PythonObject;
+
+use crate::dsl::DslPlan;
+
+/// This is for `polars-python` to inject so that the implementation can be done there:
+/// * The impls for converting from Python objects are there.
+pub static DATASET_PROVIDER_VTABLE: OnceLock<PythonDatasetProviderVTable> = OnceLock::new();
+
+pub struct PythonDatasetProviderVTable {
+    pub reader_name: fn(dataset_object: &PythonObject) -> PlSmallStr,
+
+    pub schema: fn(dataset_object: &PythonObject) -> PolarsResult<SchemaRef>,
+
+    #[expect(clippy::type_complexity)]
+    pub to_dataset_scan: fn(
+        dataset_object: &PythonObject,
+        limit: Option<usize>,
+        projection: Option<&[PlSmallStr]>,
+    ) -> PolarsResult<DslPlan>,
+}
+
+pub fn dataset_provider_vtable() -> Result<&'static PythonDatasetProviderVTable, &'static str> {
+    DATASET_PROVIDER_VTABLE
+        .get()
+        .ok_or("DATASET_PROVIDER_VTABLE not initialized")
+}
+
+/// Currently intended only for Iceberg support
+#[derive(Debug)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+pub struct PythonDatasetProvider {
+    dataset_object: PythonObject,
+}
+
+impl PythonDatasetProvider {
+    pub fn new(dataset_object: PythonObject) -> Self {
+        Self { dataset_object }
+    }
+
+    pub fn reader_name(&self) -> PlSmallStr {
+        (dataset_provider_vtable().unwrap().reader_name)(&self.dataset_object)
+    }
+
+    pub fn schema(&self) -> PolarsResult<SchemaRef> {
+        (dataset_provider_vtable().unwrap().schema)(&self.dataset_object)
+    }
+
+    pub fn to_dataset_scan(
+        &self,
+        limit: Option<usize>,
+        projection: Option<&[PlSmallStr]>,
+    ) -> PolarsResult<DslPlan> {
+        (dataset_provider_vtable().unwrap().to_dataset_scan)(
+            &self.dataset_object,
+            limit,
+            projection,
+        )
+    }
+}

crates/polars-plan/src/dsl/python_dsl/source.rs

@@ -23,7 +23,7 @@ pub struct PythonOptionsDsl {
 }
 
 impl PythonOptionsDsl {
-    pub(crate) fn get_schema(&self) -> PolarsResult<SchemaRef> {
+    pub fn get_schema(&self) -> PolarsResult<SchemaRef> {
         match self.schema_fn.as_ref().expect("should be set").as_ref() {
             Either::Left(func) => Python::with_gil(|py| {
                 let schema = func

crates/polars-plan/src/plans/conversion/dsl_to_ir.rs

@@ -180,6 +180,12 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
                         FileScan::NDJson { .. } => {
                             sources.expand_paths(unified_scan_args, cloud_options)?
                         },
+                        #[cfg(feature = "python")]
+                        FileScan::PythonDataset { .. } => {
+                            // There are a lot of places that short-circuit if the paths is empty,
+                            // so we just give a dummy path here.
+                            ScanSources::Paths(Arc::from(["dummy".into()]))
+                        },
                         FileScan::Anonymous { .. } => sources,
                     };
 
@@ -244,6 +250,20 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
                         cloud_options,
                     )
                     .map_err(|e| e.context(failed_here!(ndjson scan)))?,
+                    #[cfg(feature = "python")]
+                    FileScan::PythonDataset { dataset_object, .. } => {
+                        if crate::dsl::DATASET_PROVIDER_VTABLE.get().is_none() {
+                            polars_bail!(ComputeError: "DATASET_PROVIDER_VTABLE (python) not initialized")
+                        }
+
+                        let schema = dataset_object.schema()?;
+
+                        FileInfo {
+                            schema: schema.clone(),
+                            reader_schema: Some(either::Either::Right(schema)),
+                            row_estimation: (None, usize::MAX),
+                        }
+                    },
                     FileScan::Anonymous { .. } => {
                         file_info.expect("FileInfo should be set for AnonymousScan")
                     },
@@ -292,22 +312,6 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
                     file_info.update_schema_with_hive_schema(hive_schema);
                 }
 
-                unified_scan_args.include_file_paths = unified_scan_args
-                    .include_file_paths
-                    .as_ref()
-                    .filter(|_| match &*scan_type {
-                        #[cfg(feature = "parquet")]
-                        FileScan::Parquet { .. } => true,
-                        #[cfg(feature = "ipc")]
-                        FileScan::Ipc { .. } => true,
-                        #[cfg(feature = "csv")]
-                        FileScan::Csv { .. } => true,
-                        #[cfg(feature = "json")]
-                        FileScan::NDJson { .. } => true,
-                        FileScan::Anonymous { .. } => false,
-                    })
-                    .cloned();
-
                 if let Some(ref file_path_col) = unified_scan_args.include_file_paths {
                     let schema = Arc::make_mut(&mut file_info.schema);
 

crates/polars-plan/src/plans/functions/count.rs

@@ -56,6 +56,8 @@ pub fn count_rows(
             ),
             #[cfg(feature = "json")]
             FileScan::NDJson { options } => count_rows_ndjson(sources, cloud_options),
+            #[cfg(feature = "python")]
+            FileScan::PythonDataset { .. } => unreachable!(),
             FileScan::Anonymous { .. } => {
                 unreachable!()
             },