I think the greatest pain point of the current implementation is that in order to compute the output index we must compute squares or cubes of indices (e.g. i,j,k) and while the output index may fit within int{32,64}, the intermediate square/cube calculation may overflow, throwing the entire thing off. This effectively limits the valid range of input indices.

I'm open to any suggestions for how to circumvent this issue.

As NumPy does not defines int128_t and does define npy_int128 to be identical to int64_t, I think we can to compute squared or cubes of indices try using np.uint64_t-typed variables (which are unsigned long long int-typed variables)

Alternatively, we could try to rework the formula not to have huge values.

Alright I ended up predicting when an intermediate overflow would occur, and then deferring to python in those cases (to make use of arbitrary-precision PyLongs). This only occurs for the cases where the indices are too large to safely compute the output index, and so won't kick into effect unless the input data has billions of features in the best case, and 500k in the worst case. Even when it kicks in, the performance drop only affects exactly the elements with too-large indices, which I think is reasonable to assume would be the minority of the data.

from sklearn.preprocessing import PolynomialFeatures
from scipy import sparse

X = sparse.random(100_000, 1000, random_state=0)
pf = PolynomialFeatures(interaction_only=False, include_bias=True, degree=2)
%timeit -n9 pf.fit_transform(X)

Cython only: 774 ms ± 7.6 ms per loop (mean ± std. dev. of 7 runs, 9 loops each)
PR:          784 ms ± 4.3 ms per loop (mean ± std. dev. of 7 runs, 9 loops each)

A suggestion to mention that using np.int64 for indices and indptr should work