don't copy

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

Author: onursatici

vortex-python/python/vortex/__init__.py

@@ -75,7 +75,6 @@
     Array,
     PyArray,
     _unpickle_array,  # pyright: ignore[reportPrivateUsage]
-    _unpickle_array_p5,  # pyright: ignore[reportPrivateUsage]
     array,
 )
 from .file import VortexFile, open
@@ -164,7 +163,6 @@
     "ArrayParts",
     # Pickle
     "_unpickle_array",
-    "_unpickle_array_p5",
     # File
     "VortexFile",
     "open",

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

@@ -1,6 +1,7 @@
 #  SPDX-License-Identifier: Apache-2.0
 #  SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+from collections.abc import Sequence
 from typing import final
 
 import pyarrow as pa
@@ -30,5 +31,5 @@ class ArrayContext:
 
 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_buffers: Sequence[bytes | memoryview], dtype_buffers: Sequence[bytes | memoryview]
 ) -> Array: ...

vortex-python/python/vortex/arrays.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import abc
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from typing import TYPE_CHECKING, Any
 
 import pyarrow
@@ -14,7 +14,6 @@
 from vortex._lib.serde import (  # pyright: ignore[reportMissingModuleSource]
     ArrayContext,
     ArrayParts,
-    decode_ipc_array,
     decode_ipc_array_buffers,
 )
 
@@ -473,19 +472,13 @@ def decode(cls, parts: ArrayParts, ctx: ArrayContext, dtype: DType, len: int) ->
         """
 
 
-def _unpickle_array(array_bytes: bytes, dtype_bytes: bytes) -> Array:  # pyright: ignore[reportUnusedFunction]
-    """Unpickle a Vortex array from IPC-encoded bytes.
+def _unpickle_array(array_buffers: Sequence[bytes | memoryview], dtype_buffers: Sequence[bytes | memoryview]) -> Array:  # pyright: ignore[reportUnusedFunction]
+    """Unpickle a Vortex array from IPC-encoded buffer lists.
 
-    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 for both protocol 4 and 5.
 
-    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.
+    For protocol 4, receives list[bytes] from __reduce__.
+    For protocol 5, receives list[PickleBuffer/memoryview] from __reduce_ex__.
+    Both use decode_ipc_array_buffers which concatenates the buffers during deserialization.
     """
     return decode_ipc_array_buffers(array_buffers, dtype_buffers)

vortex-python/src/arrays/mod.rs

@@ -670,29 +670,23 @@ impl PyArray {
         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);
-        }
+        // Return buffers as a list instead of concatenating
+        let array_buffers: Vec<Vec<u8>> = buffers.iter().map(|b| b.to_vec()).collect();
 
         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 dtype_buffers: Vec<Vec<u8>> = dtype_buffers.iter().map(|b| b.to_vec()).collect();
 
         let vortex_module = PyModule::import(py, "vortex")?;
         let unpickle_fn = vortex_module.getattr("_unpickle_array")?;
 
-        let args = (serialized, dtype_bytes).into_pyobject(py)?;
+        let args = (array_buffers, dtype_buffers).into_pyobject(py)?;
         Ok((unpickle_fn, args.into_any()))
     }
 
-    /// Support for Python's pickle protocol with protocol version awareness.
+    /// Support for Python's pickle protocol for protocol >= 5
     ///
-    /// When protocol >= 5, this uses PickleBuffer for out-of-band buffer transfer,
-    /// which avoids copying large data buffers.
+    /// uses PickleBuffer for out-of-band buffer transfer,
+    /// which potentially avoids copying large data buffers.
     fn __reduce_ex__<'py>(
         slf: &'py Bound<'py, Self>,
         protocol: i32,
@@ -729,7 +723,7 @@ impl PyArray {
         }
 
         let vortex_module = PyModule::import(py, "vortex")?;
-        let unpickle_fn = vortex_module.getattr("_unpickle_array_p5")?;
+        let unpickle_fn = vortex_module.getattr("_unpickle_array")?;
 
         let args = (pickle_buffers, dtype_pickle_buffers).into_pyobject(py)?;
         Ok((unpickle_fn, args.into_any()))