Skip to main content

polars_lazy/frame/
python.rs

1use std::sync::Arc;
2
3use either::Either;
4use polars_core::schema::SchemaRef;
5use polars_utils::pl_str::PlSmallStr;
6use pyo3::{Py, PyAny};
7
8use self::python_dsl::{PythonOptionsDsl, PythonScanSource};
9use crate::prelude::*;
10
11impl LazyFrame {
12    pub fn scan_from_python_function(
13        schema: Either<Py<PyAny>, SchemaRef>,
14        scan_fn: Py<PyAny>,
15        pyarrow: bool,
16        // Validate that the source gives the proper schema
17        validate_schema: bool,
18        is_pure: bool,
19        explain_name: Option<PlSmallStr>,
20        explain_detail: Option<PlSmallStr>,
21    ) -> Self {
22        DslPlan::PythonScan {
23            options: PythonOptionsDsl {
24                // Should be a python function that returns a generator
25                scan_fn: Some(scan_fn.into()),
26                schema_fn: Some(SpecialEq::new(Arc::new(schema.map_left(|obj| obj.into())))),
27                python_source: if pyarrow {
28                    PythonScanSource::Pyarrow
29                } else {
30                    PythonScanSource::IOPlugin
31                },
32                validate_schema,
33                is_pure,
34                explain_name,
35                explain_detail,
36            },
37        }
38        .into()
39    }
40}