Update PyPI workflow to build and publish Rust wheels across platforms

Build prebuilt wheels for linux (x86_64, aarch64), macOS (x86_64,
aarch64), and Windows (x64) for Python 3.10-3.13. Users get the
native extension without needing Rust installed. Also adds
vector-db-benchmark-rs as an optional dependency (pip install
vector-benchmark[rust]).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: fcostaoliveira

.github/workflows/publish-pypi.yml

@@ -5,15 +5,75 @@ on:
   workflow_dispatch:  # Allow manual triggering
 
 jobs:
-  pypi:
-    name: Publish to PyPI
+  # Build the Rust native extension wheels for each platform/arch
+  build-rust-wheels:
+    name: Build Rust wheels (${{ matrix.os }}, ${{ matrix.target }})
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          # Linux x86_64
+          - os: ubuntu-latest
+            target: x86_64
+          # Linux aarch64
+          - os: ubuntu-latest
+            target: aarch64
+          # macOS x86_64
+          - os: macos-13
+            target: x86_64
+          # macOS aarch64 (Apple Silicon)
+          - os: macos-14
+            target: aarch64
+          # Windows x86_64
+          - os: windows-latest
+            target: x64
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.10'
+
+    - name: Build wheels
+      uses: PyO3/maturin-action@v1
+      with:
+        target: ${{ matrix.target }}
+        args: --release --out dist --interpreter 3.10 3.11 3.12 3.13
+        manylinux: auto
+        working-directory: rust
+
+    - name: Upload wheels
+      uses: actions/upload-artifact@v4
+      with:
+        name: wheels-${{ matrix.os }}-${{ matrix.target }}
+        path: rust/dist/*.whl
+
+  # Build the Rust sdist (fallback for platforms without prebuilt wheels)
+  build-rust-sdist:
+    name: Build Rust sdist
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Build sdist
+      uses: PyO3/maturin-action@v1
+      with:
+        command: sdist
+        args: --out dist
+        working-directory: rust
+
+    - name: Upload sdist
+      uses: actions/upload-artifact@v4
+      with:
+        name: rust-sdist
+        path: rust/dist/*.tar.gz
+
+  # Build the main Python package
+  build-python:
+    name: Build Python package
     runs-on: ubuntu-latest
-    environment:
-      name: pypi
-      url: https://pypi.org/p/vector-benchmark
-    permissions:
-      contents: read
-      id-token: write
     steps:
     - uses: actions/checkout@v4
 
@@ -37,7 +97,51 @@ jobs:
         pip install twine
         twine check dist/*
 
-    - name: Publish to PyPI
+    - name: Upload Python dist
+      uses: actions/upload-artifact@v4
+      with:
+        name: python-dist
+        path: dist/*
+
+  # Publish everything to PyPI
+  publish:
+    name: Publish to PyPI
+    runs-on: ubuntu-latest
+    needs: [build-rust-wheels, build-rust-sdist, build-python]
+    environment:
+      name: pypi
+      url: https://pypi.org/p/vector-benchmark
+    permissions:
+      contents: read
+      id-token: write
+    steps:
+    - name: Download all artifacts
+      uses: actions/download-artifact@v4
+      with:
+        path: all-dist
+
+    - name: Collect dist files
+      run: |
+        mkdir -p dist
+        # Rust wheels and sdist
+        find all-dist/wheels-* -name '*.whl' -exec cp {} dist/ \; 2>/dev/null || true
+        find all-dist/rust-sdist -name '*.tar.gz' -exec cp {} dist/ \; 2>/dev/null || true
+        # Python package
+        cp all-dist/python-dist/* dist/
+        echo "Packages to publish:"
+        ls -la dist/
+
+    - name: Publish Rust extension to PyPI
+      uses: PyO3/maturin-action@v1
+      with:
+        command: upload
+        args: --non-interactive --skip-existing dist/vector_db_benchmark_rs*
+
+    - name: Publish Python package to PyPI
       env:
         POETRY_PYPI_TOKEN_PYPI: ${{ secrets.PYPI_TOKEN }}
-      run: poetry publish
+      run: |
+        pip install poetry
+        # Publish only the Python package (already built)
+        pip install twine
+        twine upload --skip-existing dist/vector_benchmark* dist/vector-benchmark*

pyproject.toml

@@ -53,7 +53,10 @@ psycopg = {extras = ["binary"], version = "^3.1.17"}
 pgvector = "^0.2.4"
 ml-dtypes = "^0.4.0"
 boto3 = "^1.39.4"
+vector-db-benchmark-rs = {version = ">=0.1.0", optional = true}
 
+[tool.poetry.extras]
+rust = ["vector-db-benchmark-rs"]
 
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^2.20.0"