INTPYTHON-392 Add support for free-threaded Python (#311)

blink1073 · web-flow · commit 011bab206c5b · 2025-05-09T21:12:20.000-05:00
diff --git a/.github/workflows/dist-python.yml b/.github/workflows/dist-python.yml
@@ -39,14 +39,17 @@ jobs:
         - [ubuntu-24.04, manylinux_aarch64]
         - [macos-14, macosx_*]
         - [windows-2019, win_amd64]
-        python: ["cp39", "cp310", "cp311", "cp312", "cp313"]
+        python: ["cp39", "cp310", "cp311", "cp312", "cp313", "cp313t"]
         exclude:
           - buildplat: [macos-14, macosx_*]
             python: "cp39"
+          - buildplat: [windows-2019, win_amd64]
+            python: "cp313t"
         include:
           - buildplat: [macos-13, macosx_*]
             python: "cp39"
 
+
     steps:
       - name: Checkout pymongoarrow
         uses: actions/checkout@v4
@@ -85,13 +88,15 @@ jobs:
         env:
           MACOS_TEST_SKIP: "*arm64"
           CIBW_BUILD: cp39-macosx_*
+          CIBW_ENABLE: cpython-freethreading
           MACOSX_DEPLOYMENT_TARGET: "10.14"
         run: python -m cibuildwheel --output-dir wheelhouse
 
       - name: Build wheels
         if: ${{ matrix.buildplat[0] != 'macos-11' }}
         env:
           CIBW_BUILD: ${{ matrix.python }}-${{ matrix.buildplat[1] }}
+          CIBW_ENABLE: cpython-freethreading
           MACOSX_DEPLOYMENT_TARGET: "12.0"
         run: python -m cibuildwheel --output-dir wheelhouse
 
diff --git a/.github/workflows/test-python.yml b/.github/workflows/test-python.yml
@@ -37,7 +37,10 @@ jobs:
     strategy:
       matrix:
         os: ["ubuntu-latest", "macos-latest", "windows-latest"]
-        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13", "3.13t"]
+        exclude:
+          - os: "windows-latest"
+            python-version: "3.13t"
       fail-fast: false
     name: CPython ${{ matrix.python-version }}-${{ matrix.os }}
     steps:
@@ -86,6 +89,8 @@ jobs:
       - name: Ensure imports with no test deps
         run: just import-check
       - name: Run the tests
+        env:
+          UV_PYTHON: ${{matrix.python-version}}
         run: just test
 
   docs:
diff --git a/bindings/python/CHANGELOG.md b/bindings/python/CHANGELOG.md
@@ -6,6 +6,7 @@
 # Changes in Version 1.8.0 (2025/MM/YY)
 
 - Add support for PyArrow 20.0.
+- Add support for free-threaded python on Linux and MacOS.
 
 # Changes in Version 1.7.2 (2025/04/23)
 
diff --git a/bindings/python/justfile b/bindings/python/justfile
@@ -14,7 +14,7 @@ import-check:
     uv run python -c "from pymongoarrow.lib import libbson_version"
 
 benchmark *args:
-    uv sync --dev --extra test
+    uv sync --dev --extra test --extra test-polars
     uv run asv run -e --python=$(uv run python -c "import sys;print(sys.executable)") {{args}}
 
 install:
@@ -23,7 +23,7 @@ install:
     uv run pre-commit install
 
 test *args:
-    uv sync --extra test
+    uv sync --extra test --extra test-polars || uv sync --extra test
     uv run pytest {{args}}
 
 lint:
diff --git a/bindings/python/pymongoarrow/lib.pyx b/bindings/python/pymongoarrow/lib.pyx
@@ -15,6 +15,7 @@
 # Cython compiler directives
 # distutils: language=c++
 # cython: language_level=3
+# cython: freethreading_compatible = True
 
 # Stdlib imports
 import sys
diff --git a/bindings/python/pymongoarrow/libarrow.pxd b/bindings/python/pymongoarrow/libarrow.pxd
@@ -15,6 +15,7 @@
 # Cython compiler directives
 # cython: language_level=3
 # distutils: language=c++
+# cython: freethreading_compatible = True
 from libcpp.vector cimport vector
 from libc.stdint cimport int32_t, uint8_t
 from pyarrow.lib cimport *
diff --git a/bindings/python/pymongoarrow/libbson.pxd b/bindings/python/pymongoarrow/libbson.pxd
@@ -15,6 +15,7 @@
 # Cython compiler directives
 # cython: language_level=3
 # distutils: language=c
+# cython: freethreading_compatible = True
 from libc.stdint cimport int32_t, int64_t, uint8_t, uint32_t, uint64_t
 
 
diff --git a/bindings/python/pyproject.toml b/bindings/python/pyproject.toml
@@ -51,7 +51,8 @@ Source = "https://github.com/mongodb-labs/mongo-arrow/tree/main/bindings/python"
 Tracker = "https://jira.mongodb.org/projects/INTPYTHON/issues"
 
 [project.optional-dependencies]
-test = ["pytz", "pytest", "polars"]
+test = ["pytz", "pytest"]
+test-polars = ["polars"]
 
 [tool.setuptools]
 zip-safe = false
@@ -103,6 +104,7 @@ faulthandler_timeout = 1500
 xfail_strict = true
 filterwarnings = [
   "error",
+  "module:The global interpreter lock:RuntimeWarning", # from pandas
 ]
 
 [tool.ruff]
diff --git a/bindings/python/test/test_arrow.py b/bindings/python/test/test_arrow.py
@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import concurrent.futures
 import io
 import json
 import tempfile
+import threading
 import unittest
 import unittest.mock as mock
 from datetime import date, datetime
@@ -244,8 +246,8 @@ def test_aggregate_omits_id_if_not_in_schema(self):
     def round_trip(self, data, schema, coll=None):
         if coll is None:
             coll = self.coll
-        self.coll.drop()
-        res = write(self.coll, data)
+        coll.drop()
+        res = write(coll, data)
         self.assertEqual(len(data), res.raw_result["insertedCount"])
         self.assertEqual(data, find_arrow_all(coll, {}, schema=schema))
         return res
@@ -1052,6 +1054,29 @@ def test_empty_embedded_array(self):
         assert pmapatable2.to_pylist()[0] == doc2
         write_table(pmapatable2, io.BytesIO())
 
+    def test_threading(self):
+        schema, data = self._create_data()
+
+        def run_test():
+            client = client_context.get_client(
+                event_listeners=[self.getmore_listener, self.cmd_listener]
+            )
+            name = f"test-{threading.current_thread().name}"
+            coll = client.pymongoarrow_test.get_collection(
+                name, write_concern=WriteConcern(w="majority")
+            )
+            coll.drop()
+            self.round_trip(data, Schema(schema), coll=coll)
+            client.close()
+
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            futures = []
+            for i in range(5):
+                futures.append(executor.submit(run_test))
+            concurrent.futures.wait(futures)
+            for future in futures:
+                future.result()
+
 
 class TestArrowExplicitApi(ArrowApiTestMixin, unittest.TestCase):
     def run_find(self, *args, **kwargs):
diff --git a/bindings/python/test/test_pandas.py b/bindings/python/test/test_pandas.py
@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # from datetime import datetime, timedelta
+import concurrent.futures
 import datetime
 import tempfile
+import threading
 import unittest
 import unittest.mock as mock
 import warnings
@@ -266,18 +268,22 @@ def inner(i):
         raw_data["nested"] = [inner(i) for i in range(3)]
         return pd.DataFrame(data=raw_data).astype(schema)
 
-    def test_auto_schema(self):
+    def _check_auto_schema(self, coll):
+        coll.drop()
         data = self._create_nested_data()
-        self.coll.drop()
-        res = write(self.coll, data)
+        res = write(coll, data)
         self.assertEqual(len(data), res.raw_result["insertedCount"])
         for func in [find_pandas_all, aggregate_pandas_all]:
-            out = func(self.coll, {} if func == find_pandas_all else []).drop(columns=["_id"])
+            out = func(coll, {} if func == find_pandas_all else []).drop(columns=["_id"])
             for name in data.columns:
                 val = out[name]
                 if str(val.dtype) == "object":
                     val = val.astype(data[name].dtype)
                 pd.testing.assert_series_equal(data[name], val)
+        coll.drop()
+
+    def test_auto_schema(self):
+        self._check_auto_schema(self.coll)
 
     def test_auto_schema_heterogeneous(self):
         vals = [1, "2", True, 4]
@@ -345,6 +351,26 @@ def test_exclude_none(self):
         col_data = list(self.coll.find({}))
         assert "b" not in col_data[3]
 
+    def test_threading(self):
+        def run_test():
+            client = client_context.get_client(
+                event_listeners=[self.getmore_listener, self.cmd_listener]
+            )
+            name = f"test-{threading.current_thread().name}"
+            coll = client.pymongoarrow_test.get_collection(
+                name, write_concern=WriteConcern(w="majority")
+            )
+            self._check_auto_schema(coll)
+            client.close()
+
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            futures = []
+            for i in range(5):
+                futures.append(executor.submit(run_test))
+            concurrent.futures.wait(futures)
+            for future in futures:
+                future.result()
+
 
 class TestBSONTypes(PandasTestBase):
     @classmethod
diff --git a/bindings/python/test/test_polars.py b/bindings/python/test/test_polars.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import concurrent.futures
+import threading
 import unittest
 import unittest.mock as mock
 import uuid
@@ -20,9 +22,8 @@
 from test.utils import AllowListEventListener
 
 import bson
-import polars as pl
 import pyarrow as pa
-from polars.testing import assert_frame_equal
+import pytest
 from pyarrow import int32, int64
 from pymongo import DESCENDING, WriteConcern
 from pymongo.collection import Collection
@@ -38,6 +39,15 @@
     ObjectIdType,
 )
 
+try:
+    import polars as pl
+    from polars.testing import assert_frame_equal
+except ImportError:
+    pl = None
+
+if pl is None:
+    pytest.skip("Requires polars.", allow_module_level=True)
+
 
 class PolarsTestBase(unittest.TestCase):
     @classmethod
@@ -65,14 +75,13 @@ def setUpClass(cls):
     def setUp(self):
         """Insert simple use case data."""
         self.coll.drop()
-        self.coll.insert_many(
-            [
-                {"_id": 1, "data": 10},
-                {"_id": 2, "data": 20},
-                {"_id": 3, "data": 30},
-                {"_id": 4},
-            ]
-        )
+        self.data = [
+            {"_id": 1, "data": 10},
+            {"_id": 2, "data": 20},
+            {"_id": 3, "data": 30},
+            {"_id": 4},
+        ]
+        self.coll.insert_many(self.data)
         self.cmd_listener.reset()
         self.getmore_listener.reset()
 
@@ -115,15 +124,17 @@ def test_find_simple(self):
         self.assertEqual(find_cmd.command_name, "find")
         self.assertEqual(find_cmd.command["projection"], {"_id": True, "data": True})
 
-    def test_aggregate_simple(self):
+    def _check_aggregation_simple(self, coll):
         expected = pl.DataFrame(
             data={
                 "_id": pl.Series(values=[1, 2, 3, 4], dtype=pl.Int32),
                 "data": pl.Series(values=[20, 40, 60, None], dtype=pl.Int64),
             }
         )
+        coll.drop()
+        coll.insert_many(self.data)
         projection = {"_id": True, "data": {"$multiply": [2, "$data"]}}
-        table = aggregate_polars_all(self.coll, [{"$project": projection}], schema=self.schema)
+        table = aggregate_polars_all(coll, [{"$project": projection}], schema=self.schema)
         self.assertTrue(table.equals(expected))
 
         agg_cmd = self.cmd_listener.results["started"][-1]
@@ -132,6 +143,9 @@ def test_aggregate_simple(self):
         self.assertEqual(agg_cmd.command["pipeline"][0]["$project"], projection)
         self.assertEqual(agg_cmd.command["pipeline"][1]["$project"], {"_id": True, "data": True})
 
+    def test_aggregate_simple(self):
+        self._check_aggregation_simple(self.coll)
+
     @mock.patch.object(Collection, "insert_many", side_effect=Collection.insert_many, autospec=True)
     def test_write_batching(self, mock):
         data = pl.DataFrame(data={"_id": pl.Series(values=range(100040), dtype=pl.Int64)})
@@ -395,3 +409,23 @@ def test_bson_types(self):
                 assert dfpl["value"].dtype == data_type["ptype"]
             except pl.exceptions.ComputeError:
                 assert isinstance(table["value"].type, pa.ExtensionType)
+
+    def test_threading(self):
+        def run_test():
+            client = client_context.get_client(
+                event_listeners=[self.getmore_listener, self.cmd_listener]
+            )
+            name = f"test-{threading.current_thread().name}"
+            coll = client.pymongoarrow_test.get_collection(
+                name, write_concern=WriteConcern(w="majority")
+            )
+            self._check_aggregation_simple(coll)
+            client.close()
+
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            futures = []
+            for i in range(5):
+                futures.append(executor.submit(run_test))
+            concurrent.futures.wait(futures)
+            for future in futures:
+                future.result()
diff --git a/bindings/python/uv.lock b/bindings/python/uv.lock
