Upgrade to arrow 54 and datafusion 45 (#2271)

including the pyo3 version bump (#2270 and #2042)

Author: AdamGS

Cargo.lock

@@ -52,15 +52,15 @@ categories = ["database-implementations", "data-structures", "compression"]
 anyhow = "1.0.95"
 arbitrary = "1.3.2"
 arrayref = "0.3.7"
-arrow = "53.3"
-arrow-arith = "53.3"
-arrow-array = "53.3"
-arrow-buffer = "53.3"
-arrow-cast = "53.3"
-arrow-ord = "53.3"
-arrow-schema = "53.3"
-arrow-select = "53.3"
-arrow-string = "53.3"
+arrow = "54.1"
+arrow-arith = "54.1"
+arrow-array = "54.1"
+arrow-buffer = "54.1"
+arrow-cast = "54.1"
+arrow-ord = "54.1"
+arrow-schema = "54.1"
+arrow-select = "54.1"
+arrow-string = "54.1"
 async-once-cell = "0.5.4"
 async-trait = "0.1.86"
 pkg-config = "0.3.31"
@@ -76,12 +76,12 @@ criterion = { package = "codspeed-criterion-compat", features = [
 ], version = "2.7.2" }
 croaring = "2.1.0"
 crossterm = "0.28"
-datafusion = { version = "44.0.0", default-features = false }
-datafusion-common = "44"
-datafusion-execution = "44"
-datafusion-expr = "44"
-datafusion-physical-expr = "44"
-datafusion-physical-plan = "44"
+datafusion = { version = "45", default-features = false }
+datafusion-common = "45"
+datafusion-execution = "45"
+datafusion-expr = "45"
+datafusion-physical-expr = "45"
+datafusion-physical-plan = "45"
 divan = "0.1.14"
 dyn-hash = "0.2.0"
 enum-iterator = "2.0.0"
@@ -111,15 +111,15 @@ num-traits = "0.2.19"
 num_enum = "0.7.2"
 object_store = "0.11.0"
 oneshot = "0.1.10"
-parquet = "53.3"
+parquet = "54.1"
 paste = "1.0.15"
 pin-project = "1.1.5"
 pin-project-lite = "0.2.15"
 prost = "0.13.4"
 prost-build = "0.13.4"
 prost-types = "0.13.4"
-pyo3 = { version = ">= 0.22", features = ["extension-module", "abi3-py310"] }
-pyo3-log = ">= 0.11"
+pyo3 = { version = "0.23.4", features = ["extension-module", "abi3-py310"] }
+pyo3-log = "0.12.1"
 rancor = "0.1.0"
 rand = "0.8.5"
 rand_distr = "0.4"
@@ -154,7 +154,7 @@ tikv-jemallocator = "0.6"
 thiserror = "2.0.3"
 tokio = "1.36"
 tracing = "0.1.37"
-url = "2.4"
+url = "2.5.4"
 uuid = "1.8.0"
 wasm-bindgen-futures = "0.4.39"
 

Cargo.toml

@@ -15,7 +15,7 @@ use crate::arrays::PyArray;
 
 /// Convert an Arrow object to a Vortex array.
 pub(super) fn from_arrow<'py>(obj: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyArray>> {
-    let pa = obj.py().import_bound("pyarrow")?;
+    let pa = obj.py().import("pyarrow")?;
     let pa_array = pa.getattr("Array")?;
     let chunked_array = pa.getattr("ChunkedArray")?;
     let table = pa.getattr("Table")?;

pyvortex/src/arrays/from_arrow.rs

@@ -7,7 +7,7 @@ use arrow::array::{Array as ArrowArray, ArrayRef};
 use arrow::pyarrow::ToPyArrow;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
-use pyo3::types::{IntoPyDict, PyList};
+use pyo3::types::{PyDict, PyList};
 use pyo3::PyClass;
 use vortex::array::ChunkedArray;
 use vortex::arrow::{infer_data_type, IntoArrowArray};
@@ -29,7 +29,7 @@ use crate::python_repr::PythonRepr;
 use crate::scalar::PyScalar;
 
 pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {
-    let m = PyModule::new_bound(py, "arrays")?;
+    let m = PyModule::new(py, "arrays")?;
     parent.add_submodule(&m)?;
     install_module("vortex._lib.arrays", &m)?;
 
@@ -250,24 +250,28 @@ impl PyArray {
             // the preferred type of each chunk may be different.
             let arrow_dtype = infer_data_type(chunked_array.dtype())?;
 
-            let chunks: Vec<ArrayRef> = chunked_array
+            let chunks = chunked_array
                 .chunks()
-                .map(|chunk| -> PyResult<ArrayRef> { Ok(chunk.into_arrow(&arrow_dtype)?) })
+                .map(|chunk| PyResult::Ok(chunk.into_arrow(&arrow_dtype)?))
                 .collect::<PyResult<Vec<ArrayRef>>>()?;
             if chunks.is_empty() {
                 return Err(PyValueError::new_err("No chunks in array"));
             }
+
             let pa_data_type = chunks[0].data_type().clone().to_pyarrow(py)?;
-            let chunks: PyResult<Vec<PyObject>> = chunks
+            let chunks = chunks
                 .iter()
                 .map(|arrow_array| arrow_array.into_data().to_pyarrow(py))
-                .collect();
+                .collect::<Result<Vec<_>, _>>()?;
+
+            let kwargs =
+                PyDict::from_sequence(&PyList::new(py, vec![("type", pa_data_type)])?.into_any())?;
 
             // Combine into a chunked array
-            PyModule::import_bound(py, "pyarrow")?.call_method(
+            PyModule::import(py, "pyarrow")?.call_method(
                 "chunked_array",
-                (PyList::new_bound(py, chunks?),),
-                Some(&[("type", pa_data_type)].into_py_dict_bound(py)),
+                (PyList::new(py, chunks)?,),
+                Some(&kwargs),
             )
         } else {
             Ok(vortex

pyvortex/src/arrays/mod.rs

@@ -5,7 +5,7 @@ use crate::arrays::PyArray;
 use crate::install_module;
 
 pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {
-    let m = PyModule::new_bound(py, "compress")?;
+    let m = PyModule::new(py, "compress")?;
     parent.add_submodule(&m)?;
     install_module("vortex._lib.compress", &m)?;
 

pyvortex/src/compress.rs

@@ -24,7 +24,7 @@ use crate::record_batch_reader::VortexRecordBatchReader;
 use crate::{install_module, TOKIO_RUNTIME};
 
 pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {
-    let m = PyModule::new_bound(py, "dataset")?;
+    let m = PyModule::new(py, "dataset")?;
     parent.add_submodule(&m)?;
     install_module("vortex._lib.dataset", &m)?;
 

pyvortex/src/dataset.rs

@@ -34,7 +34,7 @@ use crate::python_repr::PythonRepr;
 
 /// Register DType functions and classes.
 pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {
-    let m = PyModule::new_bound(py, "dtype")?;
+    let m = PyModule::new(py, "dtype")?;
     parent.add_submodule(&m)?;
     install_module("vortex._lib.dtype", &m)?;
 

pyvortex/src/dtype/mod.rs

@@ -18,7 +18,7 @@ use fastlanes::{PyFastLanesBitPackedEncoding, PyFastLanesDeltaEncoding, PyFastLa
 use crate::install_module;
 
 pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {
-    let m = PyModule::new_bound(py, "encoding")?;
+    let m = PyModule::new(py, "encoding")?;
     parent.add_submodule(&m)?;
     install_module("vortex._lib.encoding", &m)?;
 

pyvortex/src/encoding/mod.rs

@@ -9,7 +9,7 @@ use crate::install_module;
 use crate::scalar::factory::scalar_helper;
 
 pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {
-    let m = PyModule::new_bound(py, "expr")?;
+    let m = PyModule::new(py, "expr")?;
     parent.add_submodule(&m)?;
     install_module("vortex._lib.expr", &m)?;
 
@@ -159,7 +159,7 @@ fn coerce_expr<'py>(value: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyExpr>> {
         Ok(value.clone())
     } else if let Ok(value) = value.downcast::<PyNone>() {
         scalar(DType::Null, value)
-    } else if let Ok(value) = value.downcast::<PyLong>() {
+    } else if let Ok(value) = value.downcast::<PyInt>() {
         scalar(DType::Primitive(PType::I64, nonnull), value)
     } else if let Ok(value) = value.downcast::<PyFloat>() {
         scalar(DType::Primitive(PType::F64, nonnull), value)

pyvortex/src/expr.rs

@@ -14,7 +14,7 @@ use crate::expr::PyExpr;
 use crate::{install_module, TOKIO_RUNTIME};
 
 pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {
-    let m = PyModule::new_bound(py, "io")?;
+    let m = PyModule::new(py, "io")?;
     parent.add_submodule(&m)?;
     install_module("vortex._lib.io", &m)?;