From b011dcb0a3788e410806f1e033c372e26d547922 Mon Sep 17 00:00:00 2001
From: nicaurvi <nyecarr@gmail.com>
Date: Tue, 9 May 2023 16:14:46 -0400
Subject: [PATCH] Enable Windows python bindings (#343)

* Use int64 for counter to fix windows compilation error

* Fix windows python bindings by adding install_lib command to move windows build output into python package

* Update to use Path instead of os

* Change batch_insert num_inserts signature to signed type for OpenMP compatibility

* Update num_inserts to int32_t per PR request

---------

Co-authored-by: Nick Caurvina <nicaurvi@microsoft.com>
---
 .gitignore                      |  2 ++
 python/src/diskann_bindings.cpp |  4 +--
 setup.py                        | 43 ++++++++++++++++++++++++++++++---
 3 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/.gitignore b/.gitignore
index d7a8e474..2a098de5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -371,3 +371,5 @@ wheelhouse/*
 dist/*
 venv*/**
 *.swp
+
+gperftools
diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp
index 97143fd0..1d24616a 100644
--- a/python/src/diskann_bindings.cpp
+++ b/python/src/diskann_bindings.cpp
@@ -194,7 +194,7 @@ template <class T> struct DynamicInMemIndex
     }
 
     auto batch_insert(py::array_t<T, py::array::c_style | py::array::forcecast> &vectors,
-                      py::array_t<IdT, py::array::c_style | py::array::forcecast> &ids, const size_t num_inserts,
+                      py::array_t<IdT, py::array::c_style | py::array::forcecast> &ids, const int32_t num_inserts,
                       const int num_threads = 0)
     {
         if (num_threads == 0)
@@ -204,7 +204,7 @@ template <class T> struct DynamicInMemIndex
         py::array_t<int> insert_retvals(num_inserts);
 
 #pragma omp parallel for schedule(dynamic, 1)
-        for (size_t i = 0; i < num_inserts; i++)
+        for (int32_t i = 0; i < num_inserts; i++)
         {
             insert_retvals.mutable_data()[i] = _index->insert_point(vectors.data(i), *(ids.data(i)));
         }
diff --git a/setup.py b/setup.py
index d309825d..ff5bed18 100644
--- a/setup.py
+++ b/setup.py
@@ -3,12 +3,14 @@
 
 import os
 import re
+import shutil
 import subprocess
 import sys
 from pathlib import Path
 
-from setuptools import Extension, find_packages, setup
+from setuptools import Extension, setup
 from setuptools.command.build_ext import build_ext
+from setuptools.command.install_lib import install_lib
 
 # Convert distutils Windows platform specifiers to CMake -A arguments
 PLAT_TO_CMAKE = {
@@ -44,7 +46,7 @@ class CMakeBuild(build_ext):
             f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}{os.sep}",
             f"-DPYTHON_EXECUTABLE={sys.executable}",
             f"-DCMAKE_BUILD_TYPE={cfg}",  # not used on MSVC, but no harm
-            f"-DVERSION_INFO={self.distribution.get_version()}" # commented out because we want this set in the CMake file
+            f"-DVERSION_INFO={self.distribution.get_version()}"  # commented out, we want this set in the CMake file
         ]
         build_args = []
         # Adding CMake arguments set as environment variable
@@ -125,10 +127,43 @@ class CMakeBuild(build_ext):
         )
 
 
+class InstallCMakeLibs(install_lib):
+    def run(self):
+        """
+        Windows only copy from the x64/Release directory and place them in the package
+        """
+
+        self.announce("Moving library files", level=3)
+
+        self.skip_build = True
+
+        # we only need to move the windows build output
+        windows_build_output_dir = Path('.') / 'x64' / 'Release'
+
+        if windows_build_output_dir.exists():
+            libs = [
+                os.path.join(windows_build_output_dir, _lib) for _lib in
+                os.listdir(windows_build_output_dir) if
+                os.path.isfile(os.path.join(windows_build_output_dir, _lib)) and
+                os.path.splitext(_lib)[1] in [".dll", '.lib', '.pyd', '.exp']
+            ]
+
+            for lib in libs:
+                shutil.move(
+                    lib,
+                    os.path.join(self.build_dir, 'diskannpy', os.path.basename(lib))
+                )
+
+        super().run()
+
+
 setup(
     ext_modules=[CMakeExtension("diskannpy._diskannpy", ".")],
-    cmdclass={"build_ext": CMakeBuild},
+    cmdclass={
+        "build_ext": CMakeBuild,
+        'install_lib': InstallCMakeLibs
+    },
     zip_safe=False,
-    package_dir = {"diskannpy": "python/src"},
+    package_dir={"diskannpy": "python/src"},
     exclude_package_data={"diskannpy": ["diskann_bindings.cpp"]}
 )