refactor(python): Expose transform as a submodule for pyiceberg_core (apache#628)

Author: Xuanwo

bindings/python/Cargo.toml

@@ -32,5 +32,5 @@ crate-type = ["cdylib"]
 
 [dependencies]
 iceberg = { path = "../../crates/iceberg" }
-pyo3 = { version = "0.21.1", features = ["extension-module"] }
-arrow = { version = "52.2.0", features = ["pyarrow"] }
+pyo3 = { version = "0.21", features = ["extension-module"] }
+arrow = { version = "52", features = ["pyarrow"] }

bindings/python/src/error.rs

@@ -0,0 +1,24 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use pyo3::exceptions::PyValueError;
+use pyo3::PyErr;
+
+/// Convert an iceberg error to a python error
+pub fn to_py_err(err: iceberg::Error) -> PyErr {
+ PyValueError::new_err(err.to_string())
+}

bindings/python/src/lib.rs

@@ -15,23 +15,13 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use iceberg::io::FileIOBuilder;
 use pyo3::prelude::*;
-use pyo3::wrap_pyfunction;
 
+mod error;
 mod transform;
 
-#[pyfunction]
-fn hello_world() -> PyResult<String> {
- let _ = FileIOBuilder::new_fs_io().build().unwrap();
- Ok("Hello, world!".to_string())
-}
-
-
 #[pymodule]
-fn pyiceberg_core_rust(m: &Bound<'_, PyModule>) -> PyResult<()> {
- m.add_function(wrap_pyfunction!(hello_world, m)?)?;
-
- m.add_class::<transform::ArrowArrayTransform>()?;
+fn pyiceberg_core_rust(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
+ transform::register_module(py, m)?;
  Ok(())
 }

bindings/python/src/transform.rs

@@ -15,24 +15,55 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use arrow::array::{make_array, Array, ArrayData};
+use arrow::pyarrow::{FromPyArrow, ToPyArrow};
 use iceberg::spec::Transform;
 use iceberg::transform::create_transform_function;
+use pyo3::prelude::*;
 
-use arrow::{
- array::{make_array, Array, ArrayData},
-};
-use arrow::pyarrow::{FromPyArrow, ToPyArrow};
-use pyo3::{exceptions::PyValueError, prelude::*};
+use crate::error::to_py_err;
+
+#[pyfunction]
+pub fn identity(py: Python, array: PyObject) -> PyResult<PyObject> {
+ apply(py, array, Transform::Identity)
+}
+
+#[pyfunction]
+pub fn void(py: Python, array: PyObject) -> PyResult<PyObject> {
+ apply(py, array, Transform::Void)
+}
+
+#[pyfunction]
+pub fn year(py: Python, array: PyObject) -> PyResult<PyObject> {
+ apply(py, array, Transform::Year)
+}
+
+#[pyfunction]
+pub fn month(py: Python, array: PyObject) -> PyResult<PyObject> {
+ apply(py, array, Transform::Month)
+}
 
-fn to_py_err(err: iceberg::Error) -> PyErr {
- PyValueError::new_err(err.to_string())
+#[pyfunction]
+pub fn day(py: Python, array: PyObject) -> PyResult<PyObject> {
+ apply(py, array, Transform::Day)
 }
 
-#[pyclass]
-pub struct ArrowArrayTransform {
+#[pyfunction]
+pub fn hour(py: Python, array: PyObject) -> PyResult<PyObject> {
+ apply(py, array, Transform::Hour)
 }
 
-fn apply(array: PyObject, transform: Transform, py: Python) -> PyResult<PyObject> {
+#[pyfunction]
+pub fn bucket(py: Python, array: PyObject, num_buckets: u32) -> PyResult<PyObject> {
+ apply(py, array, Transform::Bucket(num_buckets))
+}
+
+#[pyfunction]
+pub fn truncate(py: Python, array: PyObject, width: u32) -> PyResult<PyObject> {
+ apply(py, array, Transform::Truncate(width))
+}
+
+fn apply(py: Python, array: PyObject, transform: Transform) -> PyResult<PyObject> {
  // import
  let array = ArrayData::from_pyarrow_bound(array.bind(py))?;
  let array = make_array(array);
@@ -43,45 +74,20 @@ fn apply(array: PyObject, transform: Transform, py: Python) -> PyResult<PyObject
  array.to_pyarrow(py)
 }
 
-#[pymethods]
-impl ArrowArrayTransform {
- #[staticmethod]
- pub fn identity(array: PyObject, py: Python) -> PyResult<PyObject> {
- apply(array, Transform::Identity, py)
- }
-
- #[staticmethod]
- pub fn void(array: PyObject, py: Python) -> PyResult<PyObject> {
- apply(array, Transform::Void, py)
- }
-
- #[staticmethod]
- pub fn year(array: PyObject, py: Python) -> PyResult<PyObject> {
- apply(array, Transform::Year, py)
- }
-
- #[staticmethod]
- pub fn month(array: PyObject, py: Python) -> PyResult<PyObject> {
- apply(array, Transform::Month, py)
- }
-
- #[staticmethod]
- pub fn day(array: PyObject, py: Python) -> PyResult<PyObject> {
- apply(array, Transform::Day, py)
- }
-
- #[staticmethod]
- pub fn hour(array: PyObject, py: Python) -> PyResult<PyObject> {
- apply(array, Transform::Hour, py)
- }
+pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
+ let this = PyModule::new_bound(py, "transform")?;
 
- #[staticmethod]
- pub fn bucket(array: PyObject, num_buckets: u32, py: Python) -> PyResult<PyObject> {
- apply(array, Transform::Bucket(num_buckets), py)
- }
+ this.add_function(wrap_pyfunction!(identity, &this)?)?;
+ this.add_function(wrap_pyfunction!(void, &this)?)?;
+ this.add_function(wrap_pyfunction!(year, &this)?)?;
+ this.add_function(wrap_pyfunction!(month, &this)?)?;
+ this.add_function(wrap_pyfunction!(day, &this)?)?;
+ this.add_function(wrap_pyfunction!(hour, &this)?)?;
+ this.add_function(wrap_pyfunction!(bucket, &this)?)?;
+ this.add_function(wrap_pyfunction!(truncate, &this)?)?;
 
- #[staticmethod]
- pub fn truncate(array: PyObject, width: u32, py: Python) -> PyResult<PyObject> {
- apply(array, Transform::Truncate(width), py)
- }
+ m.add_submodule(&this)?;
+ py.import_bound("sys")?
+ .getattr("modules")?
+  .set_item("pyiceberg_core.transform", this)
 }

bindings/python/tests/test_basic.py

@@ -19,18 +19,18 @@
 
 import pyarrow as pa
 import pytest
-from pyiceberg_core import ArrowArrayTransform
+from pyiceberg_core import transform
 
 
 def test_identity_transform():
  arr = pa.array([1, 2])
- result = ArrowArrayTransform.identity(arr)
+ result = transform.identity(arr)
  assert result == arr
 
 
 def test_bucket_transform():
  arr = pa.array([1, 2])
- result = ArrowArrayTransform.bucket(arr, 10)
+ result = transform.bucket(arr, 10)
  expected = pa.array([6, 2], type=pa.int32())
  assert result == expected
 
@@ -41,14 +41,14 @@ def test_bucket_transform_fails_for_list_type_input():
  ValueError,
  match=r"FeatureUnsupported => Unsupported data type for bucket transform",
  ):
- ArrowArrayTransform.bucket(arr, 10)
+ transform.bucket(arr, 10)
 
 
 def test_bucket_chunked_array():
  chunked = pa.chunked_array([pa.array([1, 2]), pa.array([3, 4])])
  result_chunks = []
  for arr in chunked.iterchunks():
- result_chunks.append(ArrowArrayTransform.bucket(arr, 10))
+ result_chunks.append(transform.bucket(arr, 10))
 
  expected = pa.chunked_array(
  [pa.array([6, 2], type=pa.int32()), pa.array([5, 0], type=pa.int32())]
@@ -58,34 +58,42 @@ def test_bucket_chunked_array():
 
 def test_year_transform():
  arr = pa.array([date(1970, 1, 1), date(2000, 1, 1)])
- result = ArrowArrayTransform.year(arr)
+ result = transform.year(arr)
  expected = pa.array([0, 30], type=pa.int32())
  assert result == expected
 
 
 def test_month_transform():
  arr = pa.array([date(1970, 1, 1), date(2000, 4, 1)])
- result = ArrowArrayTransform.month(arr)
+ result = transform.month(arr)
  expected = pa.array([0, 30 * 12 + 3], type=pa.int32())
  assert result == expected
 
 
 def test_day_transform():
  arr = pa.array([date(1970, 1, 1), date(2000, 4, 1)])
- result = ArrowArrayTransform.day(arr)
+ result = transform.day(arr)
  expected = pa.array([0, 11048], type=pa.int32())
  assert result == expected
 
 
 def test_hour_transform():
  arr = pa.array([datetime(1970, 1, 1, 19, 1, 23), datetime(2000, 3, 1, 12, 1, 23)])
- result = ArrowArrayTransform.hour(arr)
+ result = transform.hour(arr)
  expected = pa.array([19, 264420], type=pa.int32())
  assert result == expected
 
 
 def test_truncate_transform():
  arr = pa.array(["this is a long string", "hi my name is sung"])
- result = ArrowArrayTransform.truncate(arr, 5)
+ result = transform.truncate(arr, 5)
  expected = pa.array(["this ", "hi my"])
  assert result == expected
+
+
+def test_identity_transform_with_direct_import():
+ from pyiceberg_core.transform import identity
+
+ arr = pa.array([1, 2])
+ result = identity(arr)
+ assert result == arr