[SWIG][mmlspark] allow allocating more than int max array (#2859)

microsoft · Apr 19, 2020 · 0d3e204 · 0d3e204
1 parent 995a597
commit 0d3e204
Showing 1 changed file with 30 additions and 5 deletions.
diff --git a/swig/lightgbmlib.i b/swig/lightgbmlib.i
@@ -233,11 +233,6 @@
 %pointer_cast(int32_t *, void *, int32_t_to_voidp_ptr)
 %pointer_cast(int64_t *, void *, int64_t_to_voidp_ptr)
 
-%array_functions(double, doubleArray)
-%array_functions(float, floatArray)
-%array_functions(int, intArray)
-%array_functions(long, longArray)
-
 /* Custom pointer manipulation template */
 %define %pointer_manipulation(TYPE, NAME)
 %{
@@ -278,6 +273,36 @@ TYPE *NAME##_handle();
 
 %enddef
 
+%define %long_array_functions(TYPE,NAME)
+%{
+  static TYPE *new_##NAME(int64_t nelements) { %}
+  %{  return new TYPE[nelements](); %}
+  %{}
+
+  static void delete_##NAME(TYPE *ary) { %}
+  %{  delete [] ary; %}
+  %{}
+
+  static TYPE NAME##_getitem(TYPE *ary, int64_t index) {
+    return ary[index];
+  }
+  static void NAME##_setitem(TYPE *ary, int64_t index, TYPE value) {
+    ary[index] = value;
+  }
+  %}
+
+TYPE *new_##NAME(int64_t nelements);
+void delete_##NAME(TYPE *ary);
+TYPE NAME##_getitem(TYPE *ary, int64_t index);
+void NAME##_setitem(TYPE *ary, int64_t index, TYPE value);
+
+%enddef
+
+%long_array_functions(double, doubleArray)
+%long_array_functions(float, floatArray)
+%long_array_functions(int, intArray)
+%long_array_functions(long, longArray)
+
 %pointer_manipulation(void*, voidpp)
 
 /* Allow dereferencing of void** to void* */