feat/perf: add performant async implementation

Author: winstxnhdw

py-gxhash/Cargo.toml

@@ -13,5 +13,9 @@ name = "gxhash"
 crate-type = ["cdylib"]
 
 [dependencies]
-pyo3 = "0.22.0"
+pyo3 = "0.24.0"
 gxhash = { path = "..", features = ["hybrid"], default-features = false }
+pyo3-async-runtimes = { version = "0.24.0", features = ["tokio-runtime"] }
+tokio = "1.44.1"
+memmap2 = "0.9.5"
+libc = "0.2.171"

py-gxhash/README.md

@@ -1,6 +1,5 @@
 # py-gxhash
 
 ```bash
-uv venv --seed
-uv run maturin develop
+uv sync --reinstall
 ```

py-gxhash/gxhash.pyi

@@ -1,4 +1,9 @@
-def gxhash32(input_bytes: bytes, seed: int) -> int:
+from typing import Protocol
+
+class File(Protocol):
+    def fileno(self) -> int: ...
+
+def gxhash32(file: File, seed: int) -> int:
     """
     Summary
     -------
@@ -7,28 +12,31 @@ def gxhash32(input_bytes: bytes, seed: int) -> int:
 
     Parameters
     ----------
-    input_bytes (bytes): input bytes to hash
+    file (File)
+        file-like object
 
-    seed (int): seed for the hash function
+    seed (int)
+        seed for the hash function
 
 
     Returns
     -------
-    hash (int): u32 hash of the input bytes
+    hash (int)
+        u32 hash of the input bytes
 
 
     Example
     -------
     ```python
-    import gxhash
-
-    input_bytes =  bytes([42] * 1000)
+    file = TemporaryFile()
+    file.write(bytes([42] * 1000))
+    file.seek(0)
     seed = 1234
-    print(f"Hash is {gxhash.gxhash32(input_bytes, seed)}!")
+    print(f"Hash is {gxhash.gxhash32(file, seed)}!")
     ```
     """
 
-def gxhash32_nogil(input_bytes: bytes, seed: int) -> int:
+async def gxhash32_async(file: File, seed: int) -> int:
     """
     Summary
     -------
@@ -37,28 +45,31 @@ def gxhash32_nogil(input_bytes: bytes, seed: int) -> int:
 
     Parameters
     ----------
-    input_bytes (bytes): input bytes to hash
+    file (File)
+        file-like object
 
-    seed (int): seed for the hash function
+    seed (int)
+        seed for the hash function
 
 
     Returns
     -------
-    hash (int): u32 hash of the input bytes
+    hash (Awaitable[int])
+        u32 hash of the input bytes
 
 
     Example
     -------
     ```python
-    import gxhash
-
-    input_bytes =  bytes([42] * 1000)
+    file = TemporaryFile()
+    file.write(bytes([42] * 1000))
+    file.seek(0)
     seed = 1234
-    print(f"Hash is {gxhash.gxhash32_nogil(input_bytes, seed)}!")
+    print(f"Hash is {gxhash.gxhash32_async(file, seed)}!")
     ```
     """
 
-def gxhash64(input_bytes: bytes, seed: int) -> int:
+def gxhash64(file: File, seed: int) -> int:
     """
     Summary
     -------
@@ -67,28 +78,31 @@ def gxhash64(input_bytes: bytes, seed: int) -> int:
 
     Parameters
     ----------
-    input_bytes (bytes): input bytes to hash
+    file (File)
+        file-like object
 
-    seed (int): seed for the hash function
+    seed (int)
+        seed for the hash function
 
 
     Returns
     -------
-    hash (int): u64 hash of the input bytes
+    hash (int)
+        u64 hash of the input bytes
 
 
     Example
     -------
     ```python
-    import gxhash
-
-    input_bytes =  bytes([42] * 1000)
+    file = TemporaryFile()
+    file.write(bytes([42] * 1000))
+    file.seek(0)
     seed = 1234
-    print(f"Hash is {gxhash.gxhash64(input_bytes, seed)}!")
+    print(f"Hash is {gxhash.gxhash64(file, seed)}!")
     ```
     """
 
-def gxhash64_nogil(input_bytes: bytes, seed: int) -> int:
+async def gxhash64_async(file: File, seed: int) -> int:
     """
     Summary
     -------
@@ -97,28 +111,31 @@ def gxhash64_nogil(input_bytes: bytes, seed: int) -> int:
 
     Parameters
     ----------
-    input_bytes (bytes): input bytes to hash
+    file (File)
+        file-like object
 
-    seed (int): seed for the hash function
+    seed (int)
+        seed for the hash function
 
 
     Returns
     -------
-    hash (int): u64 hash of the input bytes
+    hash (Awaitable[int])
+        u64 hash of the input bytes
 
 
     Example
     -------
     ```python
-    import gxhash
-
-    input_bytes =  bytes([42] * 1000)
+    file = TemporaryFile()
+    file.write(bytes([42] * 1000))
+    file.seek(0)
     seed = 1234
-    print(f"Hash is {gxhash.gxhash64_nogil(input_bytes, seed)}!")
+    print(f"Hash is {gxhash.gxhash64_async(file, seed)}!")
     ```
     """
 
-def gxhash128(input_bytes: bytes, seed: int) -> int:
+def gxhash128(file: File, seed: int) -> int:
     """
     Summary
     -------
@@ -127,28 +144,31 @@ def gxhash128(input_bytes: bytes, seed: int) -> int:
 
     Parameters
     ----------
-    input_bytes (bytes): input bytes to hash
+    file (File)
+        file-like object
 
-    seed (int): seed for the hash function
+    seed (int)
+        seed for the hash function
 
 
     Returns
     -------
-    hash (int): u128 hash of the input bytes
+    hash (int)
+        u128 hash of the input bytes
 
 
     Example
     -------
     ```python
-    import gxhash
-
-    input_bytes =  bytes([42] * 1000)
+    file = TemporaryFile()
+    file.write(bytes([42] * 1000))
+    file.seek(0)
     seed = 1234
-    print(f"Hash is {gxhash.gxhash128(input_bytes, seed)}!")
+    print(f"Hash is {gxhash.gxhash128(file, seed)}!")
     ```
     """
 
-def gxhash128_nogil(input_bytes: bytes, seed: int) -> int:
+async def gxhash128_async(file: File, seed: int) -> int:
     """
     Summary
     -------
@@ -157,23 +177,26 @@ def gxhash128_nogil(input_bytes: bytes, seed: int) -> int:
 
     Parameters
     ----------
-    input_bytes (bytes): input bytes to hash
+    file (File)
+        file-like object
 
-    seed (int): seed for the hash function
+    seed (int)
+        seed for the hash function
 
 
     Returns
     -------
-    hash (int): u128 hash of the input bytes
+    hash (Awaitable[int])
+        u128 hash of the input bytes
 
 
     Example
     -------
     ```python
-    import gxhash
-
-    input_bytes =  bytes([42] * 1000)
+    file = TemporaryFile()
+    file.write(bytes([42] * 1000))
+    file.seek(0)
     seed = 1234
-    print(f"Hash is {gxhash.gxhash128_nogil(input_bytes, seed)}!")
+    print(f"Hash is {gxhash.gxhash128_async(file, seed)}!")
     ```
     """

py-gxhash/src/lib.rs

@@ -1,42 +1,61 @@
 use pyo3::prelude::*;
+use pyo3_async_runtimes::tokio::future_into_py;
+use std::os::fd::FromRawFd;
+
+fn get_file_descriptor(py: Python, file: PyObject) -> Result<i32, PyErr> {
+    file.call_method0(py, pyo3::intern!(py, "fileno"))?.extract(py)
+}
+
+fn gxhash<T>(hasher: fn(&[u8], i64) -> T, file_descriptor: i32, seed: i64) -> PyResult<T> {
+    let file = unsafe { std::fs::File::from_raw_fd(libc::dup(file_descriptor)) };
+    let mmap = unsafe { memmap2::Mmap::map(&file).unwrap() };
+    drop(file);
+    Ok(hasher(&mmap, seed))
+}
 
 #[pyfunction]
-fn gxhash32(input_bytes: &[u8], seed: i64) -> PyResult<u32> {
-    Ok(gxhash::gxhash32(input_bytes, seed))
+fn gxhash32(py: Python, file: PyObject, seed: i64) -> PyResult<u32> {
+    let file_descriptor = get_file_descriptor(py, file)?;
+    gxhash(gxhash::gxhash32, file_descriptor, seed)
 }
 
 #[pyfunction]
-fn gxhash32_nogil(py: Python, input_bytes: &[u8], seed: i64) -> PyResult<u32> {
-    py.allow_threads(|| Ok(gxhash::gxhash32(input_bytes, seed)))
+fn gxhash32_async(py: Python, file: PyObject, seed: i64) -> PyResult<Bound<PyAny>> {
+    let file_descriptor = get_file_descriptor(py, file)?;
+    future_into_py(py, async move { gxhash(gxhash::gxhash32, file_descriptor, seed) })
 }
 
 #[pyfunction]
-fn gxhash64(input_bytes: &[u8], seed: i64) -> PyResult<u64> {
-    Ok(gxhash::gxhash64(input_bytes, seed))
+fn gxhash64(py: Python, file: PyObject, seed: i64) -> PyResult<u64> {
+    let file_descriptor = get_file_descriptor(py, file)?;
+    gxhash(gxhash::gxhash64, file_descriptor, seed)
 }
 
 #[pyfunction]
-fn gxhash64_nogil(py: Python, input_bytes: &[u8], seed: i64) -> PyResult<u64> {
-    py.allow_threads(|| Ok(gxhash::gxhash64(input_bytes, seed)))
+fn gxhash64_async(py: Python, file: PyObject, seed: i64) -> PyResult<Bound<PyAny>> {
+    let file_descriptor = get_file_descriptor(py, file)?;
+    future_into_py(py, async move { gxhash(gxhash::gxhash64, file_descriptor, seed) })
 }
 
 #[pyfunction]
-fn gxhash128(input_bytes: &[u8], seed: i64) -> PyResult<u128> {
-    Ok(gxhash::gxhash128(input_bytes, seed))
+fn gxhash128(py: Python, file: PyObject, seed: i64) -> PyResult<u128> {
+    let file_descriptor = get_file_descriptor(py, file)?;
+    gxhash(gxhash::gxhash128, file_descriptor, seed)
 }
 
 #[pyfunction]
-fn gxhash128_nogil(py: Python, input_bytes: &[u8], seed: i64) -> PyResult<u128> {
-    py.allow_threads(|| Ok(gxhash::gxhash128(input_bytes, seed)))
+fn gxhash128_async(py: Python, file: PyObject, seed: i64) -> PyResult<Bound<PyAny>> {
+    let file_descriptor = get_file_descriptor(py, file)?;
+    future_into_py(py, async move { gxhash(gxhash::gxhash128, file_descriptor, seed) })
 }
 
 #[pymodule(name = "gxhash")]
 fn pygxhash(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(gxhash32, m)?)?;
-    m.add_function(wrap_pyfunction!(gxhash32_nogil, m)?)?;
+    m.add_function(wrap_pyfunction!(gxhash32_async, m)?)?;
     m.add_function(wrap_pyfunction!(gxhash64, m)?)?;
-    m.add_function(wrap_pyfunction!(gxhash64_nogil, m)?)?;
+    m.add_function(wrap_pyfunction!(gxhash64_async, m)?)?;
     m.add_function(wrap_pyfunction!(gxhash128, m)?)?;
-    m.add_function(wrap_pyfunction!(gxhash128_nogil, m)?)?;
+    m.add_function(wrap_pyfunction!(gxhash128_async, m)?)?;
     Ok(())
 }