Take over scikit-learn#16831

wdevazelhes · wdevazelhes · commit c60a3e911734 · 2021-03-14T14:26:02.000+01:00
diff --git a/sklearn/preprocessing/_csr_polynomial_expansion.pyx b/sklearn/preprocessing/_csr_polynomial_expansion.pyx
@@ -9,7 +9,9 @@ from numpy cimport ndarray
 cimport numpy as np
 
 np.import_array()
-ctypedef np.int32_t INDEX_T
+ctypedef fused INDEX_T:
+    np.int32_t
+    np.int64_t
 
 ctypedef fused DATA_T:
     np.float32_t
@@ -120,7 +122,7 @@ def _csr_polynomial_expansion(ndarray[DATA_T, ndim=1] data,
 
     cdef INDEX_T expanded_index = 0, row_starts, row_ends, i, j, k, \
                  i_ptr, j_ptr, k_ptr, num_cols_in_row,  \
-                 expanded_column
+                 expanded_column, col
 
     with nogil:
         expanded_indptr[0] = indptr[0]
diff --git a/sklearn/preprocessing/_data.py b/sklearn/preprocessing/_data.py
@@ -1774,8 +1774,12 @@ def transform(self, X):
                 to_stack.append(np.ones(shape=(n_samples, 1), dtype=X.dtype))
             to_stack.append(X)
             for deg in range(2, self.degree+1):
-                Xp_next = _csr_polynomial_expansion(X.data, X.indices,
-                                                    X.indptr, X.shape[1],
+                # use np.int64 for index datatype to prevent overflow
+                # in case X has a large dimension
+                Xp_next = _csr_polynomial_expansion(X.data,
+                                                    X.indices.astype(np.int64),
+                                                    X.indptr.astype(np.int64),
+                                                    np.int64(X.shape[1]),
                                                     self.interaction_only,
                                                     deg)
                 if Xp_next is None:
