datafusion_python/
pyarrow_util.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Conversions between PyArrow and DataFusion types
19
20use arrow::array::{Array, ArrayData};
21use arrow::pyarrow::{FromPyArrow, ToPyArrow};
22use datafusion::scalar::ScalarValue;
23use pyo3::types::{PyAnyMethods, PyList};
24use pyo3::{Bound, FromPyObject, PyAny, PyResult, Python};
25
26use crate::common::data_type::PyScalarValue;
27use crate::errors::PyDataFusionError;
28
29impl FromPyArrow for PyScalarValue {
30    fn from_pyarrow_bound(value: &Bound<'_, PyAny>) -> PyResult<Self> {
31        let py = value.py();
32        let typ = value.getattr("type")?;
33
34        // construct pyarrow array from the python value and pyarrow type
35        let factory = py.import("pyarrow")?.getattr("array")?;
36        let args = PyList::new(py, [value])?;
37        let array = factory.call1((args, typ))?;
38
39        // convert the pyarrow array to rust array using C data interface
40        let array = arrow::array::make_array(ArrayData::from_pyarrow_bound(&array)?);
41        let scalar = ScalarValue::try_from_array(&array, 0).map_err(PyDataFusionError::from)?;
42
43        Ok(PyScalarValue(scalar))
44    }
45}
46
47impl<'source> FromPyObject<'source> for PyScalarValue {
48    fn extract_bound(value: &Bound<'source, PyAny>) -> PyResult<Self> {
49        Self::from_pyarrow_bound(value)
50    }
51}
52
53pub fn scalar_to_pyarrow<'py>(
54    scalar: &ScalarValue,
55    py: Python<'py>,
56) -> PyResult<Bound<'py, PyAny>> {
57    let array = scalar.to_array().map_err(PyDataFusionError::from)?;
58    // convert to pyarrow array using C data interface
59    let pyarray = array.to_data().to_pyarrow(py)?;
60    let pyscalar = pyarray.call_method1("__getitem__", (0,))?;
61
62    Ok(pyscalar)
63}
datafusion_python/pyarrow_util.rs

datafusion_python/
pyarrow_util.rs
