Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
adopt review comment
  • Loading branch information
sungwy committed Aug 21, 2024
commit 6e1b411eb8f891a08cc65520dbbbda64457eb4f7
42 changes: 4 additions & 38 deletions bindings/python/src/transform.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,60 +15,26 @@
// specific language governing permissions and limitations
// under the License.

use std::{error, fmt};
use iceberg::spec::Transform;
use iceberg::transform::create_transform_function;
use iceberg::Error;

use arrow::{
array::{make_array, Array, ArrayData},
};
use arrow::pyarrow::{FromPyArrow, ToPyArrow};
use pyo3::{exceptions::PyValueError, prelude::*};

#[derive(Debug)]
enum PyO3IcebergError {
Error(Error),
}

impl fmt::Display for PyO3IcebergError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match *self {
PyO3IcebergError::Error(ref e) => e.fmt(f),
}
}
}

impl error::Error for PyO3IcebergError {
fn source(&self) -> Option<&(dyn error::Error + 'static)> {
match *self {
// The cause is the underlying implementation error type. Is implicitly
// cast to the trait object `&error::Error`. This works because the
// underlying type already implements the `Error` trait.
PyO3IcebergError::Error(ref e) => Some(e),
}
}
}

impl From<Error> for PyO3IcebergError {
fn from(err: Error) -> PyO3IcebergError {
PyO3IcebergError::Error(err)
}
}

impl From<PyO3IcebergError> for PyErr {
fn from(err: PyO3IcebergError) -> PyErr {
PyValueError::new_err(err.to_string())
}
fn to_py_err(err: iceberg::Error) -> PyErr {
PyValueError::new_err(err.to_string())
}

#[pyfunction]
pub fn bucket_transform(array: PyObject, num_buckets: u32, py: Python) -> PyResult<PyObject> {
// import
let array = ArrayData::from_pyarrow_bound(array.bind(py))?;
let array = make_array(array);
let bucket = create_transform_function(&Transform::Bucket(num_buckets)).map_err(PyO3IcebergError::from)?;
let array = bucket.transform(array).map_err(PyO3IcebergError::from)?;
let bucket = create_transform_function(&Transform::Bucket(num_buckets)).map_err(to_py_err)?;
let array = bucket.transform(array).map_err(to_py_err)?;
// export
let array = array.into_data();
array.to_pyarrow(py)
Expand Down