feat: Use Cython to speed up

.github/workflows/publish_dev_package.yml

@@ -6,38 +6,88 @@ on:
       - master
 
 jobs:
-  build-n-publish:
-    name: Build and publish Python 🐍 distributions 📦 to TestPyPI
-    runs-on: ubuntu-latest
-    environment: testpypi
-    permissions:
-      id-token: write
+  build_wheels:
+    name: Build wheels on ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest, macos-13, macos-14]
+
+    steps:
+    - name: Check out from Git
+      uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+
+    - name: Get history and tags for SCM versioning
+      run: |
+        git fetch --prune --unshallow || true
+        git fetch --depth=1 origin +refs/tags/*:refs/tags/* || true
+
+    - name: Build wheels
+      uses: pypa/[email protected]
+      env:
+        CIBW_BUILD_VERBOSITY: 1
+
+    - name: Upload wheels
+      uses: actions/upload-artifact@v4
+      with:
+        name: wheels-${{ matrix.os }}
+        path: ./wheelhouse/*.whl
 
+  build_sdist:
+    name: Build source distribution
+    runs-on: ubuntu-latest
+
     steps:
     - name: Check out from Git
       uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+
     - name: Get history and tags for SCM versioning
       run: |
-          git fetch --prune --unshallow
-          git fetch --depth=1 origin +refs/tags/*:refs/tags/*
-    - name: Set up Python 3.13
+        git fetch --prune --unshallow || true
+        git fetch --depth=1 origin +refs/tags/*:refs/tags/* || true
+
+    - name: Set up Python
       uses: actions/setup-python@v5
       with:
-        python-version: 3.13
-    - name: Install pypa/build
-      run: >-
-        python -m
-        pip install
-        build
-        --user
-    - name: Build a binary wheel and a source tarball
-      run: >-
-        python -m
-        build
-        --sdist
-        --wheel
-        --outdir dist/
-        .
+        python-version: '3.13'
+
+    - name: Build sdist
+      run: |
+        python -m pip install build
+        python -m build --sdist --outdir dist/
+
+    - name: Upload sdist
+      uses: actions/upload-artifact@v4
+      with:
+        name: sdist
+        path: dist/*.tar.gz
+
+  upload_testpypi:
+    name: Upload to TestPyPI
+    needs: [build_wheels, build_sdist]
+    runs-on: ubuntu-latest
+    environment: testpypi
+    permissions:
+      id-token: write
+
+    steps:
+    - name: Download all artifacts
+      uses: actions/download-artifact@v4
+      with:
+        pattern: wheels-*
+        merge-multiple: true
+        path: dist/
+
+    - name: Download sdist
+      uses: actions/download-artifact@v4
+      with:
+        name: sdist
+        path: dist/
+
     - name: Publish distribution 📦 to Test PyPI
       uses: pypa/gh-action-pypi-publish@release/v1
       with:

.github/workflows/publish_on_release.yml

@@ -5,38 +5,88 @@ on:
     types: [published]
 
 jobs:
-  build-n-publish:
-    name: Build and publish Python 🐍 distributions 📦 to PyPI
-    runs-on: ubuntu-latest
-    environment: pypi
-    permissions:
-      id-token: write
+  build_wheels:
+    name: Build wheels on ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest, macos-13, macos-14]
+
+    steps:
+    - name: Check out from Git
+      uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+
+    - name: Get history and tags for SCM versioning
+      run: |
+        git fetch --prune --unshallow || true
+        git fetch --depth=1 origin +refs/tags/*:refs/tags/* || true
+
+    - name: Build wheels
+      uses: pypa/[email protected]
+      env:
+        CIBW_BUILD_VERBOSITY: 1
+
+    - name: Upload wheels
+      uses: actions/upload-artifact@v4
+      with:
+        name: wheels-${{ matrix.os }}
+        path: ./wheelhouse/*.whl
 
+  build_sdist:
+    name: Build source distribution
+    runs-on: ubuntu-latest
+
     steps:
     - name: Check out from Git
       uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+
     - name: Get history and tags for SCM versioning
       run: |
-          git fetch --prune --unshallow
-          git fetch --depth=1 origin +refs/tags/*:refs/tags/*
-    - name: Set up Python 3.13
+        git fetch --prune --unshallow || true
+        git fetch --depth=1 origin +refs/tags/*:refs/tags/* || true
+
+    - name: Set up Python
       uses: actions/setup-python@v5
       with:
-        python-version: 3.13
-    - name: Install pypa/build
-      run: >-
-        python -m
-        pip install
-        build
-        --user
-    - name: Build a binary wheel and a source tarball
-      run: >-
-        python -m
-        build
-        --sdist
-        --wheel
-        --outdir dist/
-        .
+        python-version: '3.13'
+
+    - name: Build sdist
+      run: |
+        python -m pip install build
+        python -m build --sdist --outdir dist/
+
+    - name: Upload sdist
+      uses: actions/upload-artifact@v4
+      with:
+        name: sdist
+        path: dist/*.tar.gz
+
+  upload_pypi:
+    name: Upload to PyPI
+    needs: [build_wheels, build_sdist]
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+
+    steps:
+    - name: Download all artifacts
+      uses: actions/download-artifact@v4
+      with:
+        pattern: wheels-*
+        merge-multiple: true
+        path: dist/
+
+    - name: Download sdist
+      uses: actions/download-artifact@v4
+      with:
+        name: sdist
+        path: dist/
+
     - name: Publish distribution 📦 to PyPI
       if: startsWith(github.ref, 'refs/tags')
       uses: pypa/gh-action-pypi-publish@release/v1

MANIFEST.in

@@ -0,0 +1,13 @@
+include README.md
+include LICENSE
+include pyproject.toml
+include setup.py
+include _version_helper.py
+
+recursive-include pymilvus *.pyx *.pxd
+recursive-include pymilvus *.py
+recursive-exclude pymilvus *.pyc
+recursive-exclude pymilvus __pycache__
+
+global-exclude *.so
+global-exclude *.c

pymilvus/client/_fast_extract.pyx

@@ -0,0 +1,113 @@
+# cython: language_level=3
+# cython: boundscheck=False
+# cython: wraparound=False
+# cython: cdivision=True
+
+from cpython.dict cimport PyDict_SetItem
+from cpython.list cimport PyList_GET_ITEM, PyList_GET_SIZE
+from cpython.object cimport PyObject
+
+
+cpdef void assign_scalar_fast(
+    list entity_rows,
+    str field_name,
+    list data_list,
+    list valid_data,
+    bint has_valid
+):
+    """Fast Cython implementation of scalar field assignment.
+
+    Args:
+        entity_rows: List of dictionaries to populate
+        field_name: Name of the field to set
+        data_list: List of values to assign (already converted from protobuf)
+        valid_data: List of validity flags (empty if no validity checks)
+        has_valid: Whether to check validity
+    """
+    cdef Py_ssize_t i
+    cdef Py_ssize_t row_count = PyList_GET_SIZE(entity_rows)
+    cdef PyObject* field_name_obj = <PyObject*>field_name
+    cdef dict row
+    cdef object value
+
+    if has_valid:
+        for i in range(row_count):
+            row = <dict>PyList_GET_ITEM(entity_rows, i)
+            if <bint>PyList_GET_ITEM(valid_data, i):
+                value = <object>PyList_GET_ITEM(data_list, i)
+            else:
+                value = None
+            PyDict_SetItem(row, field_name, value)
+    else:
+        for i in range(row_count):
+            row = <dict>PyList_GET_ITEM(entity_rows, i)
+            value = <object>PyList_GET_ITEM(data_list, i)
+            PyDict_SetItem(row, field_name, value)
+
+
+cpdef void assign_array_fast(
+    list entity_rows,
+    str field_name,
+    list array_data,
+    list valid_data,
+    bint has_valid,
+    int element_type
+):
+    """Fast Cython implementation of array field assignment.
+
+    Args:
+        entity_rows: List of dictionaries to populate
+        field_name: Name of the field to set
+        array_data: List of array proto objects
+        valid_data: List of validity flags
+        has_valid: Whether to check validity
+        element_type: DataType of array elements
+    """
+    cdef Py_ssize_t i
+    cdef Py_ssize_t row_count = PyList_GET_SIZE(entity_rows)
+    cdef dict row
+    cdef object array_item
+    cdef object value
+
+    # element_type values from DataType enum
+    # INT64=5, BOOL=1, INT8/16/32=2/3/4, FLOAT=10, DOUBLE=11, STRING/VARCHAR=20/21
+
+    if has_valid:
+        for i in range(row_count):
+            row = <dict>PyList_GET_ITEM(entity_rows, i)
+            if not <bint>PyList_GET_ITEM(valid_data, i):
+                PyDict_SetItem(row, field_name, None)
+            else:
+                array_item = <object>PyList_GET_ITEM(array_data, i)
+                # Extract based on element type
+                if element_type == 5:  # INT64
+                    value = array_item.long_data.data
+                elif element_type == 1:  # BOOL
+                    value = array_item.bool_data.data
+                elif element_type in (2, 3, 4):  # INT8/16/32
+                    value = array_item.int_data.data
+                elif element_type == 10:  # FLOAT
+                    value = array_item.float_data.data
+                elif element_type == 11:  # DOUBLE
+                    value = array_item.double_data.data
+                else:  # STRING/VARCHAR
+                    value = array_item.string_data.data
+                PyDict_SetItem(row, field_name, value)
+    else:
+        for i in range(row_count):
+            row = <dict>PyList_GET_ITEM(entity_rows, i)
+            array_item = <object>PyList_GET_ITEM(array_data, i)
+            # Extract based on element type
+            if element_type == 5:  # INT64
+                value = array_item.long_data.data
+            elif element_type == 1:  # BOOL
+                value = array_item.bool_data.data
+            elif element_type in (2, 3, 4):  # INT8/16/32
+                value = array_item.int_data.data
+            elif element_type == 10:  # FLOAT
+                value = array_item.float_data.data
+            elif element_type == 11:  # DOUBLE
+                value = array_item.double_data.data
+            else:  # STRING/VARCHAR
+                value = array_item.string_data.data
+            PyDict_SetItem(row, field_name, value)