polars-plan 0.54.2

Lazy query engine for the Polars DataFrame library
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

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 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,
        existing_resolved_version_key: Option<&str>,
        limit: Option<usize>,
        projection: Option<&[PlSmallStr]>,
        filter_columns: Option<&[PlSmallStr]>,
        pyarrow_predicate: Option<&str>,
    ) -> PolarsResult<Option<(DslPlan, PlSmallStr)>>,
}

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))]
#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
pub struct PythonDatasetProvider {
    dataset_object: PythonObject,
}

impl PythonDatasetProvider {
    pub fn new(dataset_object: PythonObject) -> Self {
        Self { dataset_object }
    }

    pub fn name(&self) -> PlSmallStr {
        (dataset_provider_vtable().unwrap().name)(&self.dataset_object)
    }

    pub fn schema(&self) -> PolarsResult<SchemaRef> {
        (dataset_provider_vtable().unwrap().schema)(&self.dataset_object)
    }

    pub fn to_dataset_scan(
        &self,
        existing_resolved_version_key: Option<&str>,
        limit: Option<usize>,
        projection: Option<&[PlSmallStr]>,
        filter_columns: Option<&[PlSmallStr]>,
        pyarrow_predicate: Option<&str>,
    ) -> PolarsResult<Option<(DslPlan, PlSmallStr)>> {
        (dataset_provider_vtable().unwrap().to_dataset_scan)(
            &self.dataset_object,
            existing_resolved_version_key,
            limit,
            projection,
            filter_columns,
            pyarrow_predicate,
        )
    }
}