Vector endian swap helper (w/ new package structure)

bin/target_driver.sh

@@ -11,4 +11,5 @@ git fetch origin
 git checkout "$version"
 git pull origin "$version"
 cd ..
-cp driver/tests/unit/common/codec/packstream/v1/test_packstream.py tests/v1/from_driver/test_packstream.py
+cp driver/tests/unit/common/codec/packstream/v1/test_packstream.py tests/codec/packstream/v1/from_driver/test_packstream.py
+cp driver/tests/unit/common/test_vector.py tests/vector/from_driver/test_vector.py

changelog.d/45.feature.md

@@ -0,0 +1,3 @@
+Add extension for the `Vector` type<ISSUES_LIST>.
+* Speed up endian conversion (byte flipping).
+* Speed up conversion from and to native python types.

pyproject.toml

@@ -63,7 +63,7 @@ build-backend = "maturin"
 
 [tool.maturin]
 features = ["pyo3/extension-module", "pyo3/generate-import-lib"]
-module-name = "neo4j._codec.packstream._rust"
+module-name = "neo4j._rust"
 exclude = [
     "/.editorconfig",
     ".gitignore",

requirements-dev.txt

@@ -8,6 +8,7 @@ isort>=5.11.5  # TODO: 6.0 - bump when support for Python 3.7 is dropped
 tox>=4.8.0  # TODO: 6.0 - bump when support for Python 3.7 is dropped
 pytest>=7.4.4  # TODO: 6.0 - bump when support for Python 3.7 is dropped
 pytest-benchmark>=4.0.0
+pytest-mock>=3.14.1
 
 # for Python driver's TestKit backend
 freezegun>=1.5.1

src/codec.rs

@@ -0,0 +1,33 @@
+// Copyright (c) "Neo4j"
+// Neo4j Sweden AB [https://neo4j.com]
+//
+// Licensed 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
+//
+//     https://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.
+
+mod packstream;
+
+use pyo3::prelude::*;
+
+use crate::register_package;
+
+pub(super) fn init_module(m: &Bound<PyModule>, name: &str) -> PyResult<()> {
+    let py = m.py();
+
+    m.gil_used(false)?;
+    register_package(m, name)?;
+
+    let mod_packstream = PyModule::new(py, "packstream")?;
+    m.add_submodule(&mod_packstream)?;
+    packstream::init_module(&mod_packstream, format!("{name}.packstream").as_str())?;
+
+    Ok(())
+}

src/codec/packstream.rs

@@ -0,0 +1,104 @@
+// Copyright (c) "Neo4j"
+// Neo4j Sweden AB [https://neo4j.com]
+//
+// Licensed 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
+//
+//     https://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.
+
+mod v1;
+
+use pyo3::basic::CompareOp;
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use pyo3::types::{PyBytes, PyTuple};
+use pyo3::IntoPyObjectExt;
+
+use crate::register_package;
+
+pub(super) fn init_module(m: &Bound<PyModule>, name: &str) -> PyResult<()> {
+    let py = m.py();
+
+    m.gil_used(false)?;
+    register_package(m, name)?;
+
+    let mod_v1 = PyModule::new(py, "v1")?;
+    m.add_submodule(&mod_v1)?;
+    v1::init_module(&mod_v1, format!("{name}.v1").as_str())?;
+
+    m.add_class::<Structure>()?;
+
+    Ok(())
+}
+
+#[pyclass]
+#[derive(Debug)]
+pub struct Structure {
+    tag: u8,
+    #[pyo3(get)]
+    fields: Vec<PyObject>,
+}
+
+#[pymethods]
+impl Structure {
+    #[new]
+    #[pyo3(signature = (tag, *fields))]
+    #[pyo3(text_signature = "(tag, *fields)")]
+    fn new(tag: &[u8], fields: Vec<PyObject>) -> PyResult<Self> {
+        if tag.len() != 1 {
+            return Err(PyErr::new::<PyValueError, _>("tag must be a single byte"));
+        }
+        let tag = tag[0];
+        Ok(Self { tag, fields })
+    }
+
+    #[getter(tag)]
+    fn read_tag<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
+        PyBytes::new(py, &[self.tag])
+    }
+
+    #[getter(fields)]
+    fn read_fields<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
+        PyTuple::new(py, &self.fields)
+    }
+
+    fn eq(&self, other: &Self, py: Python<'_>) -> PyResult<bool> {
+        if self.tag != other.tag || self.fields.len() != other.fields.len() {
+            return Ok(false);
+        }
+        for (a, b) in self
+            .fields
+            .iter()
+            .map(|e| e.bind(py))
+            .zip(other.fields.iter().map(|e| e.bind(py)))
+        {
+            if !a.eq(b)? {
+                return Ok(false);
+            }
+        }
+        Ok(true)
+    }
+
+    fn __richcmp__(&self, other: &Self, op: CompareOp, py: Python<'_>) -> PyResult<PyObject> {
+        Ok(match op {
+            CompareOp::Eq => self.eq(other, py)?.into_py_any(py)?,
+            CompareOp::Ne => (!self.eq(other, py)?).into_py_any(py)?,
+            _ => py.NotImplemented(),
+        })
+    }
+
+    fn __hash__(&self, py: Python<'_>) -> PyResult<isize> {
+        let mut fields_hash = 0;
+        for field in &self.fields {
+            fields_hash += field.bind(py).hash()?;
+        }
+        Ok(fields_hash.wrapping_add(self.tag.into()))
+    }
+}

src/v1.rs → src/codec/packstream/v1.rs

@@ -19,6 +19,8 @@ mod unpack;
 use pyo3::prelude::*;
 use pyo3::wrap_pyfunction;
 
+use crate::register_package;
+
 const TINY_STRING: u8 = 0x80;
 const TINY_LIST: u8 = 0x90;
 const TINY_MAP: u8 = 0xA0;
@@ -44,7 +46,10 @@ const BYTES_8: u8 = 0xCC;
 const BYTES_16: u8 = 0xCD;
 const BYTES_32: u8 = 0xCE;
 
-pub(crate) fn register(m: &Bound<PyModule>) -> PyResult<()> {
+pub(crate) fn init_module(m: &Bound<PyModule>, name: &str) -> PyResult<()> {
+    m.gil_used(false)?;
+    register_package(m, name)?;
+
     m.add_function(wrap_pyfunction!(unpack::unpack, m)?)?;
     m.add_function(wrap_pyfunction!(pack::pack, m)?)?;
 

src/v1/pack.rs → src/codec/packstream/v1/pack.rs

@@ -22,12 +22,12 @@ use pyo3::sync::GILOnceCell;
 use pyo3::types::{PyBytes, PyDict, PyString, PyType};
 use pyo3::{intern, IntoPyObjectExt};
 
+use super::super::Structure;
 use super::{
     BYTES_16, BYTES_32, BYTES_8, FALSE, FLOAT_64, INT_16, INT_32, INT_64, INT_8, LIST_16, LIST_32,
     LIST_8, MAP_16, MAP_32, MAP_8, NULL, STRING_16, STRING_32, STRING_8, TINY_LIST, TINY_MAP,
     TINY_STRING, TINY_STRUCT, TRUE,
 };
-use crate::Structure;
 
 #[derive(Debug)]
 struct TypeMappings {

src/v1/unpack.rs → src/codec/packstream/v1/unpack.rs

@@ -19,12 +19,12 @@ use pyo3::sync::with_critical_section;
 use pyo3::types::{IntoPyDict, PyByteArray, PyBytes, PyDict, PyList, PyTuple};
 use pyo3::{intern, IntoPyObjectExt};
 
+use super::super::Structure;
 use super::{
     BYTES_16, BYTES_32, BYTES_8, FALSE, FLOAT_64, INT_16, INT_32, INT_64, INT_8, LIST_16, LIST_32,
     LIST_8, MAP_16, MAP_32, MAP_8, NULL, STRING_16, STRING_32, STRING_8, TINY_LIST, TINY_MAP,
     TINY_STRING, TINY_STRUCT, TRUE,
 };
-use crate::Structure;
 
 #[pyfunction]
 #[pyo3(signature = (bytes, idx, hydration_hooks=None))]

src/lib.rs

@@ -13,26 +13,23 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-pub mod v1;
+mod codec;
+mod vector;
 
-use pyo3::basic::CompareOp;
-use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
-use pyo3::types::{PyBytes, PyTuple};
-use pyo3::IntoPyObjectExt;
 
 #[pymodule(gil_used = false)]
 #[pyo3(name = "_rust")]
-fn packstream(m: &Bound<PyModule>) -> PyResult<()> {
+fn init_module(m: &Bound<PyModule>) -> PyResult<()> {
     let py = m.py();
 
-    m.add_class::<Structure>()?;
+    let mod_codec = PyModule::new(py, "codec")?;
+    m.add_submodule(&mod_codec)?;
+    codec::init_module(&mod_codec, "codec")?;
 
-    let mod_v1 = PyModule::new(py, "v1")?;
-    mod_v1.gil_used(false)?;
-    v1::register(&mod_v1)?;
-    m.add_submodule(&mod_v1)?;
-    register_package(&mod_v1, "v1")?;
+    let mod_vector = PyModule::new(py, "vector")?;
+    m.add_submodule(&mod_vector)?;
+    vector::init_module(&mod_vector, "vector")?;
 
     Ok(())
 }
@@ -41,7 +38,7 @@ fn packstream(m: &Bound<PyModule>) -> PyResult<()> {
 // https://github.com/PyO3/pyo3/issues/1517#issuecomment-808664021
 fn register_package(m: &Bound<PyModule>, name: &str) -> PyResult<()> {
     let py = m.py();
-    let module_name = format!("neo4j._codec.packstream._rust.{name}").into_pyobject(py)?;
+    let module_name = format!("neo4j._rust.{name}").into_pyobject(py)?;
 
     py.import("sys")?
         .getattr("modules")?
@@ -50,68 +47,3 @@ fn register_package(m: &Bound<PyModule>, name: &str) -> PyResult<()> {
 
     Ok(())
 }
-
-#[pyclass]
-#[derive(Debug)]
-pub struct Structure {
-    tag: u8,
-    #[pyo3(get)]
-    fields: Vec<PyObject>,
-}
-
-#[pymethods]
-impl Structure {
-    #[new]
-    #[pyo3(signature = (tag, *fields))]
-    #[pyo3(text_signature = "(tag, *fields)")]
-    fn new(tag: &[u8], fields: Vec<PyObject>) -> PyResult<Self> {
-        if tag.len() != 1 {
-            return Err(PyErr::new::<PyValueError, _>("tag must be a single byte"));
-        }
-        let tag = tag[0];
-        Ok(Self { tag, fields })
-    }
-
-    #[getter(tag)]
-    fn read_tag<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
-        PyBytes::new(py, &[self.tag])
-    }
-
-    #[getter(fields)]
-    fn read_fields<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
-        PyTuple::new(py, &self.fields)
-    }
-
-    fn eq(&self, other: &Self, py: Python<'_>) -> PyResult<bool> {
-        if self.tag != other.tag || self.fields.len() != other.fields.len() {
-            return Ok(false);
-        }
-        for (a, b) in self
-            .fields
-            .iter()
-            .map(|e| e.bind(py))
-            .zip(other.fields.iter().map(|e| e.bind(py)))
-        {
-            if !a.eq(b)? {
-                return Ok(false);
-            }
-        }
-        Ok(true)
-    }
-
-    fn __richcmp__(&self, other: &Self, op: CompareOp, py: Python<'_>) -> PyResult<PyObject> {
-        Ok(match op {
-            CompareOp::Eq => self.eq(other, py)?.into_py_any(py)?,
-            CompareOp::Ne => (!self.eq(other, py)?).into_py_any(py)?,
-            _ => py.NotImplemented(),
-        })
-    }
-
-    fn __hash__(&self, py: Python<'_>) -> PyResult<isize> {
-        let mut fields_hash = 0;
-        for field in &self.fields {
-            fields_hash += field.bind(py).hash()?;
-        }
-        Ok(fields_hash.wrapping_add(self.tag.into()))
-    }
-}

src/vector.rs

@@ -0,0 +1,41 @@
+// Copyright (c) "Neo4j"
+// Neo4j Sweden AB [https://neo4j.com]
+//
+// Licensed 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
+//
+//     https://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.
+
+mod native_conversion;
+mod swap_endian;
+
+use crate::register_package;
+use pyo3::prelude::*;
+
+pub(super) fn init_module(m: &Bound<PyModule>, name: &str) -> PyResult<()> {
+    m.gil_used(false)?;
+    register_package(m, name)?;
+
+    m.add_function(wrap_pyfunction!(swap_endian::swap_endian, m)?)?;
+    m.add_function(wrap_pyfunction!(native_conversion::vec_f64_from_native, m)?)?;
+    m.add_function(wrap_pyfunction!(native_conversion::vec_f64_to_native, m)?)?;
+    m.add_function(wrap_pyfunction!(native_conversion::vec_f32_from_native, m)?)?;
+    m.add_function(wrap_pyfunction!(native_conversion::vec_f32_to_native, m)?)?;
+    m.add_function(wrap_pyfunction!(native_conversion::vec_i64_from_native, m)?)?;
+    m.add_function(wrap_pyfunction!(native_conversion::vec_i64_to_native, m)?)?;
+    m.add_function(wrap_pyfunction!(native_conversion::vec_i32_from_native, m)?)?;
+    m.add_function(wrap_pyfunction!(native_conversion::vec_i32_to_native, m)?)?;
+    m.add_function(wrap_pyfunction!(native_conversion::vec_i16_from_native, m)?)?;
+    m.add_function(wrap_pyfunction!(native_conversion::vec_i16_to_native, m)?)?;
+    m.add_function(wrap_pyfunction!(native_conversion::vec_i8_from_native, m)?)?;
+    m.add_function(wrap_pyfunction!(native_conversion::vec_i8_to_native, m)?)?;
+
+    Ok(())
+}