polars_python/dataframe/
export.rs

1use arrow::datatypes::IntegerType;
2use arrow::record_batch::RecordBatch;
3use polars::prelude::*;
4use polars_compute::cast::CastOptionsImpl;
5use pyo3::IntoPyObjectExt;
6use pyo3::prelude::*;
7use pyo3::types::{PyCapsule, PyList, PyTuple};
8
9use super::PyDataFrame;
10use crate::conversion::{ObjectValue, Wrap};
11use crate::error::PyPolarsErr;
12use crate::interop;
13use crate::interop::arrow::to_py::dataframe_to_stream;
14use crate::prelude::PyCompatLevel;
15use crate::utils::EnterPolarsExt;
16
17#[pymethods]
18impl PyDataFrame {
19    #[cfg(feature = "object")]
20    pub fn row_tuple<'py>(&self, idx: i64, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
21        let idx = if idx < 0 {
22            (self.df.height() as i64 + idx) as usize
23        } else {
24            idx as usize
25        };
26        if idx >= self.df.height() {
27            return Err(PyPolarsErr::from(polars_err!(oob = idx, self.df.height())).into());
28        }
29        PyTuple::new(
30            py,
31            self.df.get_columns().iter().map(|s| match s.dtype() {
32                DataType::Object(_) => {
33                    let obj: Option<&ObjectValue> = s.get_object(idx).map(|any| any.into());
34                    obj.into_py_any(py).unwrap()
35                },
36                _ => Wrap(s.get(idx).unwrap()).into_py_any(py).unwrap(),
37            }),
38        )
39    }
40
41    #[cfg(feature = "object")]
42    pub fn row_tuples<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyList>> {
43        let mut rechunked;
44        // Rechunk if random access would become rather expensive.
45        // TODO: iterate over the chunks directly instead of using random access.
46        let df = if self.df.max_n_chunks() > 16 {
47            rechunked = self.df.clone();
48            rechunked.as_single_chunk_par();
49            &rechunked
50        } else {
51            &self.df
52        };
53        PyList::new(
54            py,
55            (0..df.height()).map(|idx| {
56                PyTuple::new(
57                    py,
58                    df.get_columns().iter().map(|c| match c.dtype() {
59                        DataType::Null => py.None(),
60                        DataType::Object(_) => {
61                            let obj: Option<&ObjectValue> = c.get_object(idx).map(|any| any.into());
62                            obj.into_py_any(py).unwrap()
63                        },
64                        _ => {
65                            // SAFETY: we are in bounds.
66                            let av = unsafe { c.get_unchecked(idx) };
67                            Wrap(av).into_py_any(py).unwrap()
68                        },
69                    }),
70                )
71                .unwrap()
72            }),
73        )
74    }
75
76    #[allow(clippy::wrong_self_convention)]
77    pub fn to_arrow(&mut self, py: Python, compat_level: PyCompatLevel) -> PyResult<Vec<PyObject>> {
78        py.enter_polars_ok(|| self.df.align_chunks_par())?;
79        let pyarrow = py.import("pyarrow")?;
80
81        let rbs = self
82            .df
83            .iter_chunks(compat_level.0, true)
84            .map(|rb| interop::arrow::to_py::to_py_rb(&rb, py, &pyarrow))
85            .collect::<PyResult<_>>()?;
86        Ok(rbs)
87    }
88
89    /// Create a `Vec` of PyArrow RecordBatch instances.
90    ///
91    /// Note this will give bad results for columns with dtype `pl.Object`,
92    /// since those can't be converted correctly via PyArrow. The calling Python
93    /// code should make sure these are not included.
94    #[allow(clippy::wrong_self_convention)]
95    pub fn to_pandas(&mut self, py: Python) -> PyResult<Vec<PyObject>> {
96        py.enter_polars_ok(|| self.df.as_single_chunk_par())?;
97        Python::with_gil(|py| {
98            let pyarrow = py.import("pyarrow")?;
99            let cat_columns = self
100                .df
101                .get_columns()
102                .iter()
103                .enumerate()
104                .filter(|(_i, s)| {
105                    matches!(
106                        s.dtype(),
107                        DataType::Categorical(_, _) | DataType::Enum(_, _)
108                    )
109                })
110                .map(|(i, _)| i)
111                .collect::<Vec<_>>();
112
113            let enum_and_categorical_dtype = ArrowDataType::Dictionary(
114                IntegerType::Int64,
115                Box::new(ArrowDataType::LargeUtf8),
116                false,
117            );
118
119            let mut replaced_schema = None;
120            let rbs = self
121                .df
122                .iter_chunks(CompatLevel::oldest(), true)
123                .map(|rb| {
124                    let length = rb.len();
125                    let (schema, mut arrays) = rb.into_schema_and_arrays();
126
127                    // Pandas does not allow unsigned dictionary indices so we replace them.
128                    replaced_schema =
129                        (replaced_schema.is_none() && !cat_columns.is_empty()).then(|| {
130                            let mut schema = schema.as_ref().clone();
131                            for i in &cat_columns {
132                                let (_, field) = schema.get_at_index_mut(*i).unwrap();
133                                field.dtype = enum_and_categorical_dtype.clone();
134                            }
135                            Arc::new(schema)
136                        });
137
138                    for i in &cat_columns {
139                        let arr = arrays.get_mut(*i).unwrap();
140                        let out = polars_compute::cast::cast(
141                            &**arr,
142                            &enum_and_categorical_dtype,
143                            CastOptionsImpl::default(),
144                        )
145                        .unwrap();
146                        *arr = out;
147                    }
148                    let schema = replaced_schema
149                        .as_ref()
150                        .map_or(schema, |replaced| replaced.clone());
151                    let rb = RecordBatch::new(length, schema, arrays);
152
153                    interop::arrow::to_py::to_py_rb(&rb, py, &pyarrow)
154                })
155                .collect::<PyResult<_>>()?;
156            Ok(rbs)
157        })
158    }
159
160    #[allow(unused_variables)]
161    #[pyo3(signature = (requested_schema=None))]
162    fn __arrow_c_stream__<'py>(
163        &'py mut self,
164        py: Python<'py>,
165        requested_schema: Option<PyObject>,
166    ) -> PyResult<Bound<'py, PyCapsule>> {
167        py.enter_polars_ok(|| self.df.align_chunks_par())?;
168        dataframe_to_stream(&self.df, py)
169    }
170}
polars_python/dataframe/export.rs

polars_python/dataframe/
export.rs
