---
name: pyo3-zero-copy-bridge
description: High-performance Rust and Python zero-copy FFI bridge generator for AI inference engines and numeric extensions. Enforces GIL release via py.allow_threads, NumPy/Arrow contiguous memory sharing, and PyO3 0.21+ Bound lifetime boundary safety.
---

# pyo3-zero-copy-bridge: Zero-Copy Rust/Python FFI Architecture for AI Runtimes

This skill guides AI agents (Cursor, Claude Code, Antigravity) in generating high-throughput, low-latency FFI glue code connecting Python AI applications with high-performance Rust kernels.

General LLM code generators frequently introduce hidden buffer copies, hold the Global Interpreter Lock (GIL) during heavy numerical loops, or construct unsafe pointers that cause segmentation faults when Python's garbage collector frees underlying arrays.

## Core Engineering Principles

### 1. Releasing the Python GIL (`Python::allow_threads`)
- **Hazard**: Executing SIMD vector operations, matrix multiplication, tokenization, or model inference while holding the GIL blocks all other Python worker threads in the process.
- **Rule**: Never run compute-intensive Rust loops while holding `Python<'py>`.
- **Implementation**:
  1. Extract Python array metadata, validate memory layout, and convert Python handles into raw slices or Rust structs inside the GIL scope.
  2. Invoke `py.allow_threads(move || { ... })` to detach execution from the Python interpreter.
  3. Inside the closure, parallelize across CPU cores using Rayon or offload to accelerator runtimes without contention.
  4. Return pure Rust data structures or pre-allocated output buffers back across the boundary.

### 2. NumPy Zero-Copy Memory Sharing
- **Hazard**: Using naive vector conversions like `array.to_vec()?` or converting to `Vec<f32>` triggers full heap reallocation and memory copy, causing latency spikes and cache thrashing on multi-megabyte tensors.
- **Rule**: Direct slice access via `numpy::PyReadonlyArrayDyn` or `numpy::PyReadwriteArrayDyn`.
- **Preconditions**:
  - Memory must be contiguous (C-contiguous or standard strided layout).
  - Data types must match exactly (e.g. `f32` vs `f64`).
  - Use `.as_slice()` or `.as_array()` which returns `ndarray::ArrayView` without allocating. If non-contiguous, explicitly fail or request a contiguous copy on the Python side before crossing FFI.

### 3. Apache Arrow C Data Interface (`arrow-rs` + `pyarrow`)
- **Hazard**: Serializing tabular embeddings or model feature batches into JSON, Python dictionaries, or pickle objects kills serialization throughput.
- **Rule**: Pass tabular data via the Apache Arrow C Data Interface (`FFI_ArrowArray` and `FFI_ArrowSchema`).
- **Mechanism**:
  - Python creates or receives PyArrow `RecordBatch` or `Table`.
  - Python exports array pointer addresses via `_export_to_c`.
  - Rust imports the pointers using `arrow::ffi::ArrowArrayStreamReader` or `arrow::ffi::from_ffi`, reconstructing `arrow::record_batch::RecordBatch` with zero memory copies.

### 4. Lifetime & Memory Boundary Safety (PyO3 0.21+ `Bound<'py, T>`)
- **Hazard**: Storing raw pointers (`*const f32`) extracted from Python NumPy buffers into long-lived Rust structs. If Python garbage collects the NumPy array object while Rust is reading the slice, memory becomes invalid (use-after-free).
- **Rule**:
  - Use PyO3's modern `Bound<'py, T>` API instead of legacy `&T` or raw pointers.
  - When storing Python references in Rust structs, wrap them in `Py<T>` and verify refcounts.
  - Slices borrowed from `PyReadonlyArray` must never outlive the `'py` GIL lifetime token unless the underlying memory is owned by Rust (e.g. via `PyArray::from_vec`).

---

## Agent System Prompt Directives

When writing PyO3 extension modules or AI inference bindings, you must follow these directives:

```markdown
1. Always use PyO3 0.21+ Bound<'py, T> APIs. Do not use deprecated raw pointer extraction.
2. For any compute loop taking more than 50 microseconds, release the GIL using py.allow_threads(|| ...).
3. NumPy arrays must be passed via numpy::PyReadonlyArray or PyReadwriteArray. Verify memory contiguity with .as_slice() or .as_array() and avoid .to_vec().
4. For tabular and column-oriented AI data, use Apache Arrow C Data Interface (arrow::ffi) instead of custom serialization.
5. Provide a production-ready Cargo.toml specifying [lib] crate-type = ["cdylib"] and abi3 stable Python ABI features (abi3-py310).
```

---

## Production Architecture Reference

### Zero-Copy NumPy Embedding Normalization with GIL Release

```rust
use pyo3::prelude::*;
use numpy::{PyArrayDyn, PyReadonlyArrayDyn};
use rayon::prelude::*;

#[pyfunction]
#[pyo3(name = "normalize_embeddings")]
pub fn normalize_embeddings<'py>(
    py: Python<'py>,
    input_array: PyReadonlyArrayDyn<'py, f32>,
) -> PyResult<Bound<'py, PyArrayDyn<f32>>> {
    // 1. Validate contiguous memory layout without copying
    let input_slice = input_array.as_slice().map_err(|e| {
        PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
            "Input array must be contiguous in memory: {e}"
        ))
    })?;

    let shape = input_array.shape().to_vec();
    let total_elements = input_slice.len();

    // 2. Allocate uninitialized or zeroed vector in Rust heap
    let mut output_vec = vec![0.0f32; total_elements];

    // 3. Release GIL for parallel computation
    py.allow_threads(|| {
        let chunk_size = shape.last().copied().unwrap_or(total_elements);

        output_vec
            .par_chunks_mut(chunk_size)
            .zip(input_slice.par_chunks(chunk_size))
            .for_each(|(out_chunk, in_chunk)| {
                let norm_sq: f32 = in_chunk.iter().map(|&x| x * x).sum();
                let inv_norm = if norm_sq > 1e-12 {
                    1.0 / norm_sq.sqrt()
                } else {
                    0.0
                };
                for i in 0..chunk_size {
                    out_chunk[i] = in_chunk[i] * inv_norm;
                }
            });
    });

    // 4. Wrap Rust vector into NumPy array (transfers ownership to Python GC without copying)
    let numpy_array = PyArrayDyn::from_vec(py, output_vec);
    let reshaped = numpy_array.reshape(shape)?;
    Ok(reshaped)
}

#[pymodule]
fn fast_ai_bridge(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(normalize_embeddings, m)?)?;
    Ok(())
}
```

### Zero-Copy Arrow RecordBatch Exchange via C Data Interface

```rust
use pyo3::prelude::*;
use pyo3::ffi::Py_uintptr_t;
use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema};
use arrow::record_batch::RecordBatch;

#[pyfunction]
pub fn process_arrow_batch(
    _py: Python<'_>,
    array_ptr: Py_uintptr_t,
    schema_ptr: Py_uintptr_t,
) -> PyResult<usize> {
    // Reconstruct Arrow structures directly from C Data Interface pointers
    let ffi_array = array_ptr as *mut FFI_ArrowArray;
    let ffi_schema = schema_ptr as *mut FFI_ArrowSchema;

    let (array_data, data_type) = unsafe {
        let ffi_array_struct = std::ptr::read(ffi_array);
        let ffi_schema_struct = std::ptr::read(ffi_schema);
        arrow::ffi::from_ffi(ffi_array_struct, &ffi_schema_struct)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?
    };

    let record_batch = RecordBatch::from(arrow::array::make_array(array_data));
    let rows = record_batch.num_rows();

    // Process record batch with zero heap reallocations
    Ok(rows)
}
```

---

## Toolchain & Build Scaffolding

### Cargo.toml
```toml
[package]
name = "fast_ai_bridge"
version = "0.1.0"
edition = "2021"

[lib]
name = "fast_ai_bridge"
crate-type = ["cdylib"]

[dependencies]
pyo3 = { version = "0.22", features = ["extension-module", "abi3-py310"] }
numpy = "0.22"
arrow = { version = "53.0", features = ["ffi"] }
rayon = "1.10"
ndarray = "0.16"

[profile.release]
opt-level = 3
lto = "fat"
codegen-units = 1
panic = "abort"
```

### pyproject.toml
```toml
[build-system]
requires = ["maturin>=1.5,<2.0"]
build-backend = "maturin"

[project]
name = "fast-ai-bridge"
version = "0.1.0"
requires-python = ">=3.10"
dependencies = [
    "numpy>=1.24",
    "pyarrow>=14.0",
]

[tool.maturin]
features = ["pyo3/extension-module"]
module-name = "fast_ai_bridge"
```

### Build & Verification Commands
```bash
# Develop mode: compile and install directly into active virtual environment
maturin develop --release

# Distribution wheel: produce abi3 stable wheels compatible with any Python >= 3.10
maturin build --release --strip
```

---

## Verification & Safety Checklist

- [ ] `py.allow_threads` wraps all heavy CPU loops, tokenizers, and model inference passes.
- [ ] NumPy arrays accessed through `.as_slice()` or `.as_array()` rather than `.to_vec()` or cloning.
- [ ] Array memory layout validated as contiguous before taking raw slice references.
- [ ] Using modern PyO3 `Bound<'py, T>` types; no deprecated raw pointers or unchecked lifetime conversions.
- [ ] `Cargo.toml` configured with `abi3-py310` to avoid compiling separate wheels for every Python patch version.
- [ ] Release profile includes `lto = "fat"` and `opt-level = 3` for SIMD auto-vectorization across FFI boundaries.