Merge pull request #514 from wudidapaopao/fix_import_torch

Isolate pybind11 internals to resolve torch import conflict

Author: auxten

.github/workflows/build_linux_arm64_wheels-gh.yml

@@ -376,6 +376,17 @@ jobs:
             pyenv shell --unset
           done
         continue-on-error: false
+      - name: Test torch compatibility on Python 3.11
+        run: |
+          export PATH="$HOME/.pyenv/bin:$PATH"
+          eval "$(pyenv init -)"
+          pyenv shell 3.11
+          python -m pip install dist/*.whl --force-reinstall
+          python -m pip install pytest
+          python -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+          python -m pytest tests/test_torch_compatibility.py -v
+          pyenv shell --unset
+        continue-on-error: false
       - name: Test DataStore on Python 3.11
         run: |
           ulimit -c unlimited

.github/workflows/build_linux_x86_wheels.yml

@@ -361,6 +361,17 @@ jobs:
             pyenv shell --unset
           done
         continue-on-error: false
+      - name: Test torch compatibility on Python 3.11
+        run: |
+          export PATH="$HOME/.pyenv/bin:$PATH"
+          eval "$(pyenv init -)"
+          pyenv shell 3.11
+          python -m pip install dist/*.whl --force-reinstall
+          python -m pip install pytest
+          python -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+          python -m pytest tests/test_torch_compatibility.py -v
+          pyenv shell --unset
+        continue-on-error: false
       - name: Test DataStore on Python 3.11
         run: |
           ulimit -c unlimited

.github/workflows/build_macos_arm64_wheels.yml

@@ -353,6 +353,17 @@ jobs:
             pyenv shell --unset
           done
         continue-on-error: false
+      - name: Test torch compatibility on Python 3.11
+        run: |
+          export PATH="$HOME/.pyenv/bin:$PATH"
+          eval "$(pyenv init -)"
+          pyenv shell 3.11
+          python -m pip install dist/*.whl --force-reinstall --no-cache-dir
+          python -m pip install pytest
+          python -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+          python -m pytest tests/test_torch_compatibility.py -v
+          pyenv shell --unset
+        continue-on-error: false
       - name: Test DataStore on Python 3.11
         run: |
           export PATH="$HOME/.pyenv/bin:$PATH"

.github/workflows/build_macos_x86_wheels.yml

@@ -353,6 +353,17 @@ jobs:
             pyenv shell --unset
           done
         continue-on-error: false
+      - name: Test torch compatibility on Python 3.11
+        run: |
+          export PATH="$HOME/.pyenv/bin:$PATH"
+          eval "$(pyenv init -)"
+          pyenv shell 3.11
+          python -m pip install dist/*.whl --force-reinstall
+          python -m pip install pytest
+          python -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+          python -m pytest tests/test_torch_compatibility.py -v
+          pyenv shell --unset
+        continue-on-error: false
       - name: Test DataStore on Python 3.11
         run: |
           ulimit -c unlimited

contrib/pybind11-cmake/CMakeLists.txt

@@ -60,6 +60,7 @@ endif()
 target_include_directories(pybind11_stubs PRIVATE ${PYBIND11_INCLUDE_DIR})
 target_include_directories(pybind11_stubs PRIVATE ${Python_INCLUDE_DIRS})
 target_compile_definitions(pybind11_stubs PUBLIC
+    PYBIND11_INTERNALS_KIND="_chdb"
     PYBIND11_NONLIMITEDAPI_LIB_SUFFIX_FOR_MODULE="${PYBIND11_LIB_SUFFIX}"
     Py_LIMITED_API=0x03090000
 )

tests/test_torch_compatibility.py

@@ -0,0 +1,39 @@
+#!python3
+
+import importlib.util
+import unittest
+
+HAS_TORCH = importlib.util.find_spec("torch") is not None
+
+
+@unittest.skipUnless(HAS_TORCH, "torch is not installed")
+class TestTorchCompatibility(unittest.TestCase):
+    """Test that chdb and torch can coexist in the same process."""
+
+    def test_import_torch_then_chdb(self):
+        import torch
+        import chdb
+
+        # torch: basic tensor operation
+        t = torch.tensor([1.0, 2.0, 3.0])
+        self.assertEqual(t.sum().item(), 6.0)
+
+        # chdb: basic SQL query
+        res = chdb.query("SELECT 1 + 2 AS result", "CSV")
+        self.assertIn("3", str(res))
+
+    def test_import_chdb_then_torch(self):
+        import chdb
+        import torch
+
+        # chdb: basic SQL query
+        res = chdb.query("SELECT 'ok' AS status", "CSV")
+        self.assertIn("ok", str(res))
+
+        # torch: basic tensor operation
+        t = torch.tensor([[1, 2], [3, 4]])
+        self.assertEqual(t.shape, (2, 2))
+
+
+if __name__ == "__main__":
+    unittest.main()