pickling support

Signed-off-by: Onur Satici <[email protected]>

Author: onursatici

Cargo.lock

@@ -156,6 +156,7 @@ prost = "0.14"
 prost-build = "0.14"
 prost-types = "0.14"
 pyo3 = { version = "0.26.0" }
+pyo3-bytes = "0.4"
 pyo3-log = "0.13.0"
 rand = "0.9.0"
 rand_distr = "0.5"

Cargo.toml

@@ -25,13 +25,16 @@ crate-type = ["rlib", "cdylib"]
 arrow-array = { workspace = true }
 arrow-data = { workspace = true }
 arrow-schema = { workspace = true }
+bytes = { workspace = true }
 itertools = { workspace = true }
 log = { workspace = true }
 mimalloc = { workspace = true }
 object_store = { workspace = true, features = ["aws", "gcp", "azure", "http"] }
 parking_lot = { workspace = true }
 pyo3 = { workspace = true, features = ["abi3", "abi3-py311"] }
+pyo3-bytes = { workspace = true }
 pyo3-log = { workspace = true }
 tokio = { workspace = true, features = ["fs", "rt-multi-thread"] }
 url = { workspace = true }
 vortex = { workspace = true, features = ["object_store", "python", "tokio"] }
+vortex-ipc = { workspace = true }

vortex-python/Cargo.toml

@@ -4,6 +4,7 @@
 import hashlib
 import math
 import os
+from typing import cast
 
 import pyarrow as pa
 import pytest
@@ -34,3 +35,9 @@ def vxf(tmpdir_factory: pytest.TempPathFactory, request: pytest.FixtureRequest)
         a = vx.array(pa.table(columns))  # pyright: ignore[reportCallIssue, reportUnknownArgumentType, reportArgumentType]
         vx.io.write(a, str(fname))
     return vx.open(str(fname))
+
+
+@pytest.fixture(scope="session", params=[10_000, 2_000_000], ids=["small", "large"])
+def array_fixture(request: pytest.FixtureRequest) -> vx.Array:
+    size = cast(int, request.param)
+    return vx.array(pa.table({"x": list(range(size))}))

vortex-python/benchmark/conftest.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+import pickle
+
+import pytest
+from pytest_benchmark.fixture import BenchmarkFixture  # pyright: ignore[reportMissingTypeStubs]
+
+import vortex as vx
+
+
+@pytest.mark.parametrize("protocol", [4, 5], ids=lambda p: f"p{p}")  # pyright: ignore[reportAny]
+@pytest.mark.parametrize("operation", ["dumps", "loads", "roundtrip"])
+@pytest.mark.benchmark(disable_gc=True)
+def test_pickle(
+    benchmark: BenchmarkFixture,
+    array_fixture: vx.Array,
+    protocol: int,
+    operation: str,
+):
+    benchmark.group = f"pickle_p{protocol}"
+
+    if operation == "dumps":
+        benchmark(lambda: pickle.dumps(array_fixture, protocol=protocol))
+    elif operation == "loads":
+        pickled_data = pickle.dumps(array_fixture, protocol=protocol)
+        benchmark(lambda: pickle.loads(pickled_data))  # pyright: ignore[reportAny]
+    elif operation == "roundtrip":
+        benchmark(lambda: pickle.loads(pickle.dumps(array_fixture, protocol=protocol)))  # pyright: ignore[reportAny]

vortex-python/benchmark/test_serialization.py

@@ -71,7 +71,13 @@
     scalar,
 )
 from ._lib.serde import ArrayContext, ArrayParts  # pyright: ignore[reportMissingModuleSource]
-from .arrays import Array, PyArray, array
+from .arrays import (
+    Array,
+    PyArray,
+    _unpickle_array,  # pyright: ignore[reportPrivateUsage]
+    _unpickle_array_p5,  # pyright: ignore[reportPrivateUsage]
+    array,
+)
 from .file import VortexFile, open
 from .scan import RepeatedScan
 
@@ -156,6 +162,9 @@
     "Registry",
     "ArrayContext",
     "ArrayParts",
+    # Pickle
+    "_unpickle_array",
+    "_unpickle_array_p5",
     # File
     "VortexFile",
     "open",

vortex-python/python/vortex/__init__.py

@@ -5,6 +5,7 @@ from typing import final
 
 import pyarrow as pa
 
+from .arrays import Array
 from .dtype import DType
 
 @final
@@ -26,3 +27,8 @@ class ArrayParts:
 @final
 class ArrayContext:
     def __len__(self) -> int: ...
+
+def decode_ipc_array(array_bytes: bytes, dtype_bytes: bytes) -> Array: ...
+def decode_ipc_array_buffers(
+    array_buffers: list[bytes | memoryview], dtype_buffers: list[bytes | memoryview]
+) -> Array: ...

vortex-python/python/vortex/_lib/serde.pyi

@@ -11,7 +11,12 @@
 
 import vortex._lib.arrays as _arrays  # pyright: ignore[reportMissingModuleSource]
 from vortex._lib.dtype import DType  # pyright: ignore[reportMissingModuleSource]
-from vortex._lib.serde import ArrayContext, ArrayParts  # pyright: ignore[reportMissingModuleSource]
+from vortex._lib.serde import (  # pyright: ignore[reportMissingModuleSource]
+    ArrayContext,
+    ArrayParts,
+    decode_ipc_array,
+    decode_ipc_array_buffers,
+)
 
 try:
     import pandas
@@ -466,3 +471,21 @@ def decode(cls, parts: ArrayParts, ctx: ArrayContext, dtype: DType, len: int) ->
         current array. Implementations of this function should validate this information, and then construct
         a new array.
         """
+
+
+def _unpickle_array(array_bytes: bytes, dtype_bytes: bytes) -> Array:  # pyright: ignore[reportUnusedFunction]
+    """Unpickle a Vortex array from IPC-encoded bytes.
+
+    This is an internal function used by the pickle module.
+    """
+    return decode_ipc_array(array_bytes, dtype_bytes)
+
+
+def _unpickle_array_p5(array_buffers: list[bytes | memoryview], dtype_buffers: list[bytes | memoryview]) -> Array:  # pyright: ignore[reportUnusedFunction]
+    """Unpickle a Vortex array from out-of-band PickleBuffers.
+
+    This is an internal function used by the pickle module. When
+    pickle protocol 5 is supported, this methods will be called on unpickle,
+    saving one extra copy operation for array buffers.
+    """
+    return decode_ipc_array_buffers(array_buffers, dtype_buffers)

vortex-python/python/vortex/arrays.py

@@ -10,15 +10,18 @@ pub(crate) mod py;
 mod range_to_sequence;
 
 use arrow_array::{Array as ArrowArray, ArrayRef as ArrowArrayRef};
+use pyo3::IntoPyObjectExt;
 use pyo3::exceptions::{PyTypeError, PyValueError};
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyList, PyRange, PyRangeMethods};
+use pyo3_bytes::PyBytes;
 use vortex::arrays::ChunkedVTable;
 use vortex::arrow::IntoArrowArray;
 use vortex::compute::{Operator, compare, take};
 use vortex::dtype::{DType, Nullability, PType, match_each_integer_ptype};
 use vortex::error::VortexError;
 use vortex::{Array, ArrayRef, ToCanonical};
+use vortex_ipc::messages::{EncoderMessage, MessageEncoder};
 
 use crate::arrays::native::PyNativeArray;
 use crate::arrays::py::{PyPythonArray, PythonArray};
@@ -653,4 +656,82 @@ impl PyArray {
             .map(|buffer| buffer.to_vec())
             .collect())
     }
+
+    /// Support for Python's pickle protocol.
+    ///
+    /// This method serializes the array using Vortex IPC format and returns
+    /// the data needed for pickle to reconstruct the array.
+    fn __reduce__<'py>(
+        slf: &'py Bound<'py, Self>,
+    ) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyAny>)> {
+        let py = slf.py();
+        let array = PyArrayRef::extract_bound(slf.as_any())?.into_inner();
+
+        let mut encoder = MessageEncoder::default();
+        let buffers = encoder.encode(EncoderMessage::Array(&*array));
+
+        // concat all buffers
+        let mut serialized = Vec::new();
+        for buf in buffers.iter() {
+            serialized.extend_from_slice(buf);
+        }
+
+        let dtype_buffers = encoder.encode(EncoderMessage::DType(array.dtype()));
+        let mut dtype_bytes = Vec::new();
+        for buf in dtype_buffers.iter() {
+            dtype_bytes.extend_from_slice(buf);
+        }
+
+        let vortex_module = PyModule::import(py, "vortex")?;
+        let unpickle_fn = vortex_module.getattr("_unpickle_array")?;
+
+        let args = (serialized, dtype_bytes).into_pyobject(py)?;
+        Ok((unpickle_fn, args.into_any()))
+    }
+
+    /// Support for Python's pickle protocol with protocol version awareness.
+    ///
+    /// When protocol >= 5, this uses PickleBuffer for out-of-band buffer transfer,
+    /// which avoids copying large data buffers.
+    fn __reduce_ex__<'py>(
+        slf: &'py Bound<'py, Self>,
+        protocol: i32,
+    ) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyAny>)> {
+        let py = slf.py();
+
+        if protocol < 5 {
+            return Self::__reduce__(slf);
+        }
+
+        let array = PyArrayRef::extract_bound(slf.as_any())?.into_inner();
+
+        let mut encoder = MessageEncoder::default();
+        let array_buffers = encoder.encode(EncoderMessage::Array(&*array));
+        let dtype_buffers = encoder.encode(EncoderMessage::DType(array.dtype()));
+
+        let pickle_module = PyModule::import(py, "pickle")?;
+        let pickle_buffer_class = pickle_module.getattr("PickleBuffer")?;
+
+        let mut pickle_buffers = Vec::new();
+        for buf in array_buffers.into_iter() {
+            // PyBytes wraps bytes::Bytes and implements the buffer protocol
+            // This allows PickleBuffer to reference the data without copying
+            let py_bytes = PyBytes::new(buf).into_py_any(py)?;
+            let pickle_buffer = pickle_buffer_class.call1((py_bytes,))?;
+            pickle_buffers.push(pickle_buffer);
+        }
+
+        let mut dtype_pickle_buffers = Vec::new();
+        for buf in dtype_buffers.into_iter() {
+            let py_bytes = PyBytes::new(buf).into_py_any(py)?;
+            let pickle_buffer = pickle_buffer_class.call1((py_bytes,))?;
+            dtype_pickle_buffers.push(pickle_buffer);
+        }
+
+        let vortex_module = PyModule::import(py, "vortex")?;
+        let unpickle_fn = vortex_module.getattr("_unpickle_array_p5")?;
+
+        let args = (pickle_buffers, dtype_pickle_buffers).into_pyobject(py)?;
+        Ok((unpickle_fn, args.into_any()))
+    }
 }