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, PyObject, 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        let val = value.call_method0("as_py")?;
34
35        // construct pyarrow array from the python value and pyarrow type
36        let factory = py.import("pyarrow")?.getattr("array")?;
37        let args = PyList::new(py, [val])?;
38        let array = factory.call1((args, typ))?;
39
40        // convert the pyarrow array to rust array using C data interface
41        let array = arrow::array::make_array(ArrayData::from_pyarrow_bound(&array)?);
42        let scalar = ScalarValue::try_from_array(&array, 0).map_err(PyDataFusionError::from)?;
43
44        Ok(PyScalarValue(scalar))
45    }
46}
47
48impl<'source> FromPyObject<'source> for PyScalarValue {
49    fn extract_bound(value: &Bound<'source, PyAny>) -> PyResult<Self> {
50        Self::from_pyarrow_bound(value)
51    }
52}
53
54pub fn scalar_to_pyarrow(scalar: &ScalarValue, py: Python) -> PyResult<PyObject> {
55    let array = scalar.to_array().map_err(PyDataFusionError::from)?;
56    // convert to pyarrow array using C data interface
57    let pyarray = array.to_data().to_pyarrow(py)?;
58    let pyscalar = pyarray.call_method1(py, "__getitem__", (0,))?;
59
60    Ok(pyscalar)
61}
datafusion_python/pyarrow_util.rs

datafusion_python/
pyarrow_util.rs
