polars_plan/dsl/python_dsl/
source.rs

1use std::sync::Arc;
2
3use either::Either;
4use polars_core::error::{PolarsResult, polars_err};
5use polars_core::schema::SchemaRef;
6use polars_utils::python_function::PythonFunction;
7use pyo3::prelude::*;
8#[cfg(feature = "serde")]
9use serde::{Deserialize, Serialize};
10
11use crate::dsl::SpecialEq;
12
13#[derive(Clone, PartialEq, Eq, Debug, Default)]
14#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
15#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
16pub struct PythonOptionsDsl {
17    /// A function that returns a Python Generator.
18    /// The generator should produce Polars DataFrame's.
19    pub scan_fn: Option<PythonFunction>,
20    /// Either the schema fn or schema is set.
21    pub schema_fn: Option<SpecialEq<Arc<Either<PythonFunction, SchemaRef>>>>,
22    pub python_source: PythonScanSource,
23    pub validate_schema: bool,
24}
25
26impl PythonOptionsDsl {
27    pub fn get_schema(&self) -> PolarsResult<SchemaRef> {
28        match self.schema_fn.as_ref().expect("should be set").as_ref() {
29            Either::Left(func) => Python::with_gil(|py| {
30                let schema = func
31                    .0
32                    .call0(py)
33                    .map_err(|e| polars_err!(ComputeError: "schema callable failed: {}", e))?;
34                crate::plans::python::python_schema_to_rust(py, schema.into_bound(py))
35            }),
36            Either::Right(schema) => Ok(schema.clone()),
37        }
38    }
39}
40
41#[derive(Clone, PartialEq, Eq, Debug, Default)]
42#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
43#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
44pub enum PythonScanSource {
45    Pyarrow,
46    Cuda,
47    #[default]
48    IOPlugin,
49}
polars_plan/dsl/python_dsl/source.rs

polars_plan/dsl/python_dsl/
source.rs
