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

polars_plan/dsl/python_dsl/
source.rs
