Possible to parallelize sklearn.impute.KNNImputer?

[khughitt]

Greetings!

Do you think it might be possible to parallelize the algorithm for sklearn.impute.KNNImputer in the future?

scikit-learn's implementation of sklearn.neighbors.KNeighborsClassifier accepts an n_jobs parameter to achieve this, but the corresponding imputation function does not and can be quite slow for large datasets.

[TomDLT]

Thanks for the suggestion!

It seems like it is only a matter of adding self.n_jobs and passing it to pairwise_distances_chunked.

[khughitt]

@TomDLT Great! If that's the case, I can spend a little time to familiarize myself with the code and submit a PR (if you don't already haven't gotten to it first.)

Thanks for the quick response and suggestion!

[TomDLT]

You are very welcome to submit a PR.

Note that I expect the code update to be simple, but the KNNImputer code is reasonably complex, and we will need a confirmation that the pairwise_distances_chunked is operating in each chunk independently in KNNImputer. @thomasjpfan @jnothman

[khughitt]

@TomDLT Indeed. I tried simply passing along n_jobs, but that doesn't seem to make much different and most of the work still occurs in a single thread:

https://github.com/khughitt/scikit-learn/commit/3e65862ad2296dda2fa65afecc363dd0b7a5a784

It looks like for the KNN imputer, most of the heavy-lifting is being handled by the reduce_func, whereas n_jobs only affects the distance calculations, so my guess is that the logic in process_chunk() (which gets passed to reduce_chunk) would need to be modified?

Another possible alternative would be to modify the transform() function so that parallelization occurs there and pairwise_distances_chunked() is called multiple times.

[TomDLT]

Ok so the bottleneck is not in the pairwise distances but in the reduce function, which will be harder to parallelize. (I think your fix can still be useful.)

I don't think we should call pairwise_distances_chunked multiple times in parallel, since it would kind of duplicate its purpose, but we could parallelize chunks inside the pairwise_distances_chunked function itself. However I think parallelizing operations in the reduce function would make more sense.

process_chunk seems to process each column independently, so we could try do that in parallel?
This will definitely require more work since it requires digging more in depth in the KNNImputer code.

[khughitt]

Sounds good. I'll wait to hear back from others before submitting a PR with the previous changes.

It could be a bit tricky: since pairwise_distances_chunked(), passing n_jobs to the reduce_func inside of it would lead to nested parallelization, and possibly using more cores than desired..

One approach might be to modify the logic a bit so that pairwise_distances_chunked() can determine which level makes the most sense to parallelize. For example, we could add another kwargs to that function along the lines of parallelize_reduce_func which defaults to False.

Then, inside the definition for process_chunk() in _knn, modify the column iteration loop at line 250:

for col in range(X.shape[1]):
    ...

to adopt a similar approach to how parallelization is handled in pairwise.py:_parallel_pairwise():

    Parallel(backend="threading", n_jobs=n_jobs)(
        fd(func, ret, s, X, Y[s], **kwds)
        for s in gen_even_slices(_num_samples(Y), effective_n_jobs(n_jobs)))

If n_effective_jobs > 1?

[jnothman]

I'll guessing that the best opportunity for parallelism is thread-based parallelism over columns... The reduce func would benefit from profiling. Cython may be a part of the solution???

[TomDLT]

it would lead to nested parallelization

I don't think it would, since pairwise_distances_chunked does not parallelize over chunks. On each chunk it parallelizes inside the pairwise_distances call, and it would parallelize later inside the reduce_func call if we implement that.

Cython may be a part of the solution???

I agree, currently the code has two nested python loops, one over line chunks, and one over columns in process_chunk. We could try to parallelize the second loop, either with multiprocessing, or with threading if we can release the GIL in a Cython version of process_chunk.

[khughitt]

Some profiling results for a 1000 x 100 random matrix with ~10% missing values:

import numpy as np
import pandas as pd
from sklearn.impute import KNNImputer
import profile

np.random.seed(1)

X = np.random.random([1000, 100])
X = pd.DataFrame(X).mask(X < 0.1)

imputer = KNNImputer(n_neighbors=5)

profile.run('imputer.fit_transform(X)')

Results:

Percent missing values: 10.14%
         49328 function calls (48509 primitive calls) in 1.971 seconds

   Ordered by: standard name

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
       70    0.000    0.000    0.000    0.000 :0(_abc_instancecheck)
       18    0.000    0.000    0.000    0.000 :0(_abc_subclasscheck)
       18    0.000    0.000    0.000    0.000 :0(_filters_mutated)
      212    0.001    0.000    0.005    0.000 :0(all)
      208    0.001    0.000    0.002    0.000 :0(any)
        2    0.000    0.000    0.000    0.000 :0(append)
      102    0.000    0.000    0.000    0.000 :0(arange)
      100    0.097    0.001    0.097    0.001 :0(argpartition)
1177/1173    0.004    0.000    0.004    0.000 :0(array)
        2    0.000    0.000    0.000    0.000 :0(c_einsum)
        2    0.000    0.000    0.000    0.000 :0(callable)
        1    0.000    0.000    0.029    0.029 :0(clip)
        9    0.000    0.000    0.000    0.000 :0(copy)
       20    0.000    0.000    0.000    0.000 :0(empty)
       10    0.000    0.000    0.000    0.000 :0(endswith)
       16    0.000    0.000    0.000    0.000 :0(ensure_int64)
        1    0.000    0.000    1.970    1.970 :0(exec)
        2    0.000    0.000    0.000    0.000 :0(fill)
      200    0.000    0.000    0.000    0.000 :0(format)
      119    0.000    0.000    0.000    0.000 :0(get)
        1    0.000    0.000    0.000    0.000 :0(get_ident)
8605/8603    0.009    0.000    0.009    0.000 :0(getattr)
     1416    0.002    0.000    0.002    0.000 :0(geterrobj)
  504/484    0.001    0.000    0.007    0.000 :0(hasattr)
       60    0.000    0.000    0.000    0.000 :0(hash)
  927/723    1.386    0.001    1.522    0.002 :0(implement_array_function)
        6    0.000    0.000    0.000    0.000 :0(insert)
       16    0.000    0.000    0.000    0.000 :0(is_iterator)
       52    0.000    0.000    0.000    0.000 :0(is_list_like)
        2    0.000    0.000    0.000    0.000 :0(is_scalar)
     5282    0.007    0.000    0.008    0.000 :0(isinstance)
     1608    0.002    0.000    0.002    0.000 :0(issubclass)
      108    0.000    0.000    0.000    0.000 :0(items)
  614/542    0.001    0.000    0.001    0.000 :0(len)
        2    0.000    0.000    0.003    0.001 :0(map_infer)
      101    0.000    0.000    0.000    0.000 :0(min)
      201    0.001    0.000    0.001    0.000 :0(nonzero)
      101    0.000    0.000    0.000    0.000 :0(ravel)
      816    0.008    0.000    0.008    0.000 :0(reduce)
        6    0.000    0.000    0.000    0.000 :0(remove)
        6    0.000    0.000    0.000    0.000 :0(rpartition)
      708    0.001    0.000    0.001    0.000 :0(seterrobj)
        1    0.001    0.001    0.001    0.001 :0(setprofile)
      202    0.000    0.000    0.000    0.000 :0(startswith)
      300    0.001    0.000    0.005    0.000 :0(sum)
      200    0.001    0.000    0.001    0.000 :0(take)
       16    0.000    0.000    0.000    0.000 :0(take_1d_object_object)
        4    0.000    0.000    0.000    0.000 :0(transpose)
     3220    0.004    0.000    0.004    0.000 :0(update)
        1    0.000    0.000    0.000    0.000 :0(vars)
1600/1100    0.004    0.000    0.030    0.000 :0(view)
      201    0.000    0.000    0.000    0.000 :0(zeros)
        1    0.000    0.000    0.000    0.000 <__array_function__ internals>:2(all)
      101    0.000    0.000    0.003    0.000 <__array_function__ internals>:2(any)
      100    0.000    0.000    0.098    0.001 <__array_function__ internals>:2(argpartition)
      100    0.000    0.000    0.001    0.000 <__array_function__ internals>:2(can_cast)
        1    0.000    0.000    0.030    0.030 <__array_function__ internals>:2(clip)
      102    0.000    0.000    0.001    0.000 <__array_function__ internals>:2(copyto)
        3    0.000    0.000    1.384    0.461 <__array_function__ internals>:2(dot)
        2    0.000    0.000    0.000    0.000 <__array_function__ internals>:2(einsum)
      101    0.000    0.000    0.004    0.000 <__array_function__ internals>:2(flatnonzero)
        4    0.000    0.000    0.000    0.000 <__array_function__ internals>:2(may_share_memory)
        2    0.000    0.000    0.001    0.001 <__array_function__ internals>:2(ndim)
      201    0.001    0.000    0.004    0.000 <__array_function__ internals>:2(nonzero)
      101    0.000    0.000    0.002    0.000 <__array_function__ internals>:2(ravel)
        2    0.000    0.000    0.000    0.000 <__array_function__ internals>:2(result_type)
      100    0.000    0.000    0.002    0.000 <__array_function__ internals>:2(shape)
        6    0.000    0.000    0.000    0.000 <__array_function__ internals>:2(sum)
      220    0.001    0.000    0.001    0.000 <frozen importlib._bootstrap>:1017(_handle_fromlist)
        6    0.000    0.000    0.000    0.000 <frozen importlib._bootstrap>:389(parent)
        1    0.000    0.000    1.970    1.970 <string>:1(<module>)
        9    0.000    0.000    0.000    0.000 __init__.py:1004(is_scalar_nan)
        2    0.000    0.000    0.000    0.000 __init__.py:729(gen_batches)
        1    0.000    0.000    0.000    0.000 __init__.py:965(get_chunk_n_rows)
  332/328    0.001    0.000    0.002    0.000 _asarray.py:14(asarray)
      107    0.000    0.000    0.000    0.000 _asarray.py:86(asanyarray)
        2    0.000    0.000    0.000    0.000 _base.py:114(_more_tags)
        1    0.000    0.000    0.000    0.000 _base.py:53(_check_weights)
      100    0.000    0.000    0.000    0.000 _base.py:64(_get_weights)
        1    0.000    0.000    0.000    0.000 _base.py:75(_fit_indicator)
        1    0.000    0.000    0.000    0.000 _base.py:85(_transform_indicator)
        1    0.000    0.000    0.000    0.000 _base.py:99(_concatenate_indicator)
        7    0.000    0.000    0.000    0.000 _config.py:14(get_config)
      250    0.000    0.000    0.000    0.000 _dtype.py:24(_kind_name)
      250    0.000    0.000    0.002    0.000 _dtype.py:307(_name_includes_bit_suffix)
      250    0.001    0.000    0.004    0.000 _dtype.py:321(_name_get)
      100    0.003    0.000    0.206    0.002 _knn.py:113(_calc_impute)
        1    0.000    0.000    0.012    0.012 _knn.py:159(fit)
        1    0.001    0.001    1.959    1.959 _knn.py:194(transform)
        1    0.024    0.024    0.243    0.243 _knn.py:239(process_chunk)
        4    0.000    0.000    0.000    0.000 _mask.py:7(_get_mask)
        1    0.000    0.000    0.029    0.029 _methods.py:106(_clip)
      300    0.001    0.000    0.004    0.000 _methods.py:45(_sum)
      208    0.000    0.000    0.001    0.000 _methods.py:53(_any)
      200    0.000    0.000    0.004    0.000 _methods.py:56(_all)
        2    0.001    0.000    0.002    0.001 _methods.py:72(_clip_dep_is_scalar_nan)
        2    0.000    0.000    0.000    0.000 _methods.py:82(_clip_dep_is_byte_swapped)
        1    0.027    0.027    0.027    0.027 _methods.py:87(_clip_dep_invoke_with_casting)
        1    0.000    0.000    0.000    0.000 _parallel_backends.py:185(in_main_thread)
        1    0.000    0.000    0.000    0.000 _parallel_backends.py:280(__init__)
        1    0.000    0.000    0.000    0.000 _parallel_backends.py:34(__init__)
        1    0.000    0.000    0.000    0.000 _parallel_backends.py:501(effective_n_jobs)
      708    0.002    0.000    0.003    0.000 _ufunc_config.py:132(geterr)
      708    0.004    0.000    0.010    0.000 _ufunc_config.py:32(seterr)
      304    0.000    0.000    0.000    0.000 _ufunc_config.py:429(__init__)
      304    0.001    0.000    0.005    0.000 _ufunc_config.py:433(__enter__)
      304    0.001    0.000    0.005    0.000 _ufunc_config.py:438(__exit__)
       18    0.000    0.000    0.000    0.000 abc.py:100(__subclasscheck__)
       70    0.000    0.000    0.000    0.000 abc.py:96(__instancecheck__)
        8    0.000    0.000    0.005    0.001 accessor.py:17(__init__)
        8    0.000    0.000    0.005    0.001 accessor.py:183(__get__)
        8    0.000    0.000    0.005    0.001 accessor.py:220(_validate)
       16    0.000    0.000    0.000    0.000 accessor.py:222(<genexpr>)
       16    0.000    0.000    0.000    0.000 algorithms.py:1457(_get_take_nd_function)
       16    0.000    0.000    0.001    0.000 algorithms.py:1586(take_nd)
       10    0.000    0.000    0.000    0.000 base.py:1032(__iter__)
       15    0.000    0.000    0.000    0.000 base.py:1189(isspmatrix)
        2    0.000    0.000    0.000    0.000 base.py:1835(is_object)
        2    0.000    0.000    0.000    0.000 base.py:1875(is_categorical)
       98    0.000    0.000    0.002    0.000 base.py:256(is_dtype)
        1    0.000    0.000    0.000    0.000 base.py:340(_more_tags)
        1    0.000    0.000    0.000    0.000 base.py:343(_get_tags)
        1    0.000    0.000    0.000    0.000 base.py:354(_check_n_features)
        1    0.000    0.000    0.011    0.011 base.py:383(_validate_data)
        2    0.000    0.000    0.000    0.000 base.py:4114(_can_hold_identifiers_and_holds_name)
      112    0.000    0.000    0.001    0.000 base.py:413(find)
       18    0.000    0.000    0.000    0.000 base.py:5551(ensure_index)
       18    0.000    0.000    0.000    0.000 base.py:5648(maybe_extract_name)
        1    0.000    0.000    1.970    1.970 base.py:663(fit_transform)
       20    0.000    0.000    0.000    0.000 blocks.py:124(__init__)
       20    0.000    0.000    0.000    0.000 blocks.py:135(_check_ndim)
       34    0.000    0.000    0.000    0.000 blocks.py:201(internal_values)
        4    0.000    0.000    0.000    0.000 blocks.py:213(get_values)
       23    0.000    0.000    0.000    0.000 blocks.py:233(mgr_locs)
       18    0.000    0.000    0.000    0.000 blocks.py:2368(__init__)
       20    0.000    0.000    0.000    0.000 blocks.py:237(mgr_locs)
        2    0.000    0.000    0.000    0.000 blocks.py:256(make_block_same_class)
       18    0.000    0.000    0.002    0.000 blocks.py:2650(get_block_type)
       18    0.000    0.000    0.003    0.000 blocks.py:2696(make_block)
        2    0.000    0.000    0.000    0.000 blocks.py:2723(_extend_blocks)
       21    0.000    0.000    0.000    0.000 blocks.py:315(dtype)
        2    0.000    0.000    0.000    0.000 blocks.py:514(astype)
        2    0.000    0.000    0.000    0.000 blocks.py:670(copy)
       18    0.000    0.000    0.000    0.000 cast.py:1179(maybe_castable)
       16    0.000    0.000    0.001    0.000 cast.py:1301(maybe_cast_to_datetime)
       16    0.000    0.000    0.000    0.000 cast.py:1584(construct_1d_ndarray_preserving_na)
        4    0.000    0.000    0.000    0.000 common.py:1025(is_datetime_or_timedelta_dtype)
        4    0.000    0.000    0.000    0.000 common.py:1180(needs_i8_conversion)
       16    0.000    0.000    0.000    0.000 common.py:1296(is_float_dtype)
        4    0.000    0.000    0.000    0.000 common.py:1330(is_bool_dtype)
       94    0.000    0.000    0.001    0.000 common.py:1460(is_extension_array_dtype)
        8    0.000    0.000    0.000    0.000 common.py:1565(_get_dtype)
       94    0.000    0.000    0.001    0.000 common.py:1600(_is_dtype_type)
       20    0.000    0.000    0.000    0.000 common.py:1733(pandas_dtype)
       76    0.000    0.000    0.000    0.000 common.py:178(classes)
       76    0.000    0.000    0.000    0.000 common.py:180(<lambda>)
       18    0.000    0.000    0.000    0.000 common.py:183(classes_and_not_datetimelike)
       18    0.000    0.000    0.000    0.000 common.py:188(<lambda>)
       56    0.000    0.000    0.001    0.000 common.py:194(is_object_dtype)
      218    0.001    0.000    0.003    0.000 common.py:224(is_sparse)
       32    0.000    0.000    0.000    0.000 common.py:348(is_datetime64_dtype)
       38    0.000    0.000    0.001    0.000 common.py:381(is_datetime64tz_dtype)
       16    0.000    0.000    0.000    0.000 common.py:422(is_timedelta64_dtype)
       22    0.000    0.000    0.001    0.000 common.py:456(is_period_dtype)
       18    0.000    0.000    0.000    0.000 common.py:492(is_interval_dtype)
       20    0.000    0.000    0.000    0.000 common.py:530(is_categorical_dtype)
        2    0.000    0.000    0.000    0.000 common.py:608(is_dtype_equal)
       18    0.000    0.000    0.000    0.000 common.py:696(is_integer_dtype)
       36    0.000    0.000    0.000    0.000 construction.py:338(extract_array)
       18    0.000    0.000    0.004    0.000 construction.py:389(sanitize_array)
       18    0.000    0.000    0.002    0.000 construction.py:516(_try_cast)
       18    0.000    0.000    0.000    0.000 construction.py:575(is_empty_data)
        2    0.000    0.000    0.000    0.000 contextlib.py:353(__init__)
        2    0.000    0.000    0.000    0.000 contextlib.py:356(__enter__)
        2    0.000    0.000    0.000    0.000 contextlib.py:359(__exit__)
      100    0.003    0.000    0.024    0.000 core.py:1011(__call__)
      100    0.005    0.000    0.027    0.000 core.py:1158(__call__)
      300    0.001    0.000    0.001    0.000 core.py:1284(_replace_dtype_fields_recursive)
      300    0.001    0.000    0.001    0.000 core.py:1316(_replace_dtype_fields)
      300    0.001    0.000    0.002    0.000 core.py:1330(make_mask_descr)
     1100    0.002    0.000    0.003    0.000 core.py:1363(getmask)
      100    0.001    0.000    0.003    0.000 core.py:1425(getmaskarray)
      200    0.001    0.000    0.002    0.000 core.py:1645(make_mask_none)
      100    0.000    0.000    0.002    0.000 core.py:1818(_check_mask_axis)
      200    0.005    0.000    0.023    0.000 core.py:2819(__new__)
      700    0.013    0.000    0.024    0.000 core.py:2949(_update_from)
      500    0.005    0.000    0.027    0.000 core.py:2975(__array_finalize__)
     1100    0.002    0.000    0.002    0.000 core.py:3406(dtype)
      600    0.001    0.000    0.001    0.000 core.py:3422(shape)
      100    0.001    0.000    0.003    0.000 core.py:3434(__setmask__)
      600    0.001    0.000    0.002    0.000 core.py:3636(_get_data)
      100    0.001    0.000    0.003    0.000 core.py:3730(filled)
      200    0.001    0.000    0.001    0.000 core.py:4027(_delegate_binop)
      100    0.001    0.000    0.025    0.000 core.py:4165(__mul__)
      100    0.001    0.000    0.028    0.000 core.py:4189(__truediv__)
      100    0.001    0.000    0.001    0.000 core.py:434(_check_fill_value)
      200    0.002    0.000    0.006    0.000 core.py:4440(count)
      100    0.001    0.000    0.016    0.000 core.py:5038(sum)
      100    0.001    0.000    0.074    0.001 core.py:5214(mean)
      200    0.002    0.000    0.003    0.000 core.py:645(get_masked_subclass)
      100    0.001    0.000    0.011    0.000 core.py:6549(array)
      200    0.000    0.000    0.000    0.000 core.py:659(<listcomp>)
      400    0.001    0.000    0.002    0.000 core.py:672(getdata)
      100    0.001    0.000    0.013    0.000 core.py:7903(asarray)
      100    0.002    0.000    0.006    0.000 core.py:850(__call__)
       18    0.000    0.000    0.000    0.000 dtypes.py:1113(is_dtype)
       22    0.000    0.000    0.000    0.000 dtypes.py:901(is_dtype)
       10    0.000    0.000    0.000    0.000 einsumfunc.py:989(_einsum_dispatcher)
        2    0.000    0.000    0.000    0.000 einsumfunc.py:997(einsum)
        1    0.121    0.121    0.121    0.121 extmath.py:119(safe_sparse_dot)
        2    0.000    0.000    0.000    0.000 extmath.py:50(row_norms)
        6    0.000    0.000    0.001    0.000 extmath.py:686(_safe_accumulator_op)
      100    0.001    0.000    0.092    0.001 extras.py:533(average)
      101    0.000    0.000    0.000    0.000 fromnumeric.py:1701(_ravel_dispatcher)
      101    0.000    0.000    0.001    0.000 fromnumeric.py:1705(ravel)
      201    0.000    0.000    0.000    0.000 fromnumeric.py:1812(_nonzero_dispatcher)
      201    0.000    0.000    0.002    0.000 fromnumeric.py:1816(nonzero)
      100    0.000    0.000    0.000    0.000 fromnumeric.py:1911(_shape_dispatcher)
      100    0.000    0.000    0.001    0.000 fromnumeric.py:1915(shape)
        1    0.000    0.000    0.000    0.000 fromnumeric.py:2029(_clip_dispatcher)
        1    0.000    0.000    0.030    0.030 fromnumeric.py:2033(clip)
        6    0.000    0.000    0.000    0.000 fromnumeric.py:2100(_sum_dispatcher)
        6    0.000    0.000    0.000    0.000 fromnumeric.py:2105(sum)
      101    0.000    0.000    0.000    0.000 fromnumeric.py:2245(_any_dispatcher)
      101    0.000    0.000    0.002    0.000 fromnumeric.py:2249(any)
        1    0.000    0.000    0.000    0.000 fromnumeric.py:2333(_all_dispatcher)
        1    0.000    0.000    0.000    0.000 fromnumeric.py:2337(all)
        2    0.000    0.000    0.000    0.000 fromnumeric.py:3071(_ndim_dispatcher)
        2    0.000    0.000    0.001    0.000 fromnumeric.py:3075(ndim)
      302    0.001    0.000    0.129    0.000 fromnumeric.py:52(_wrapfunc)
      108    0.001    0.000    0.002    0.000 fromnumeric.py:70(_wrapreduction)
      108    0.000    0.000    0.000    0.000 fromnumeric.py:71(<dictcomp>)
      100    0.000    0.000    0.000    0.000 fromnumeric.py:752(_argpartition_dispatcher)
      100    0.000    0.000    0.098    0.001 fromnumeric.py:756(argpartition)
        2    0.000    0.000    0.000    0.000 function.py:48(__call__)
      426    0.001    0.000    0.001    0.000 generic.py:10(_check)
        2    0.000    0.000    0.001    0.000 generic.py:11544(logical_func)
        4    0.000    0.000    0.000    0.000 generic.py:1777(__array__)
       20    0.000    0.000    0.000    0.000 generic.py:195(__init__)
        4    0.000    0.000    0.000    0.000 generic.py:232(attrs)
       16    0.000    0.000    0.000    0.000 generic.py:249(_validate_dtype)
        2    0.000    0.000    0.000    0.000 generic.py:365(_get_axis_number)
       18    0.000    0.000    0.000    0.000 generic.py:447(_info_axis)
        4    0.000    0.000    0.000    0.000 generic.py:5086(__finalize__)
    34/30    0.000    0.000    0.003    0.000 generic.py:5114(__getattr__)
       24    0.000    0.000    0.000    0.000 generic.py:5132(__setattr__)
        4    0.000    0.000    0.000    0.000 generic.py:5188(_protect_consolidate)
        4    0.000    0.000    0.000    0.000 generic.py:5199(_consolidate_inplace)
        4    0.000    0.000    0.000    0.000 generic.py:5202(f)
        4    0.000    0.000    0.000    0.000 generic.py:5259(values)
        4    0.000    0.000    0.000    0.000 generic.py:5336(_values)
       16    0.000    0.000    0.010    0.001 generic.py:5341(dtypes)
        2    0.000    0.000    0.000    0.000 generic.py:5384(astype)
        1    0.000    0.000    0.000    0.000 getlimits.py:366(__new__)
        2    0.000    0.000    0.000    0.000 inference.py:263(is_dict_like)
        4    0.000    0.000    0.000    0.000 inference.py:289(<genexpr>)
       60    0.000    0.000    0.000    0.000 inference.py:322(is_hashable)
       16    0.000    0.000    0.000    0.000 inspect.py:2845(parameters)
        1    0.000    0.000    0.000    0.000 inspect.py:487(getmro)
        3    0.000    0.000    0.000    0.000 inspect.py:72(isclass)
       20    0.000    0.000    0.000    0.000 managers.py:1499(__init__)
        2    0.000    0.000    0.000    0.000 managers.py:1520(from_blocks)
       16    0.000    0.000    0.000    0.000 managers.py:153(blknos)
       18    0.000    0.000    0.003    0.000 managers.py:1531(from_array)
       36    0.000    0.000    0.000    0.000 managers.py:1542(_block)
        2    0.000    0.000    0.000    0.000 managers.py:1569(dtype)
       34    0.000    0.000    0.000    0.000 managers.py:1580(internal_values)
        2    0.000    0.000    0.000    0.000 managers.py:202(shape)
        6    0.000    0.000    0.000    0.000 managers.py:204(<genexpr>)
        4    0.000    0.000    0.000    0.000 managers.py:223(_is_single_block)
        1    0.000    0.000    0.000    0.000 managers.py:228(_rebuild_blknos_and_blklocs)
        6    0.000    0.000    0.000    0.000 managers.py:249(items)
       16    0.000    0.000    0.002    0.000 managers.py:253(get_dtypes)
       16    0.000    0.000    0.000    0.000 managers.py:254(<listcomp>)
        6    0.000    0.000    0.000    0.000 managers.py:296(__len__)
        2    0.000    0.000    0.000    0.000 managers.py:356(apply)
        2    0.000    0.000    0.000    0.000 managers.py:375(<dictcomp>)
        2    0.000    0.000    0.000    0.000 managers.py:564(astype)
        4    0.000    0.000    0.000    0.000 managers.py:647(is_consolidated)
        1    0.000    0.000    0.000    0.000 managers.py:655(_consolidate_check)
        1    0.000    0.000    0.000    0.000 managers.py:656(<listcomp>)
        4    0.000    0.000    0.000    0.000 managers.py:779(as_array)
        4    0.000    0.000    0.000    0.000 managers.py:932(consolidate)
      102    0.000    0.000    0.000    0.000 multiarray.py:1043(copyto)
        4    0.000    0.000    0.000    0.000 multiarray.py:1340(may_share_memory)
      100    0.000    0.000    0.000    0.000 multiarray.py:468(can_cast)
        2    0.000    0.000    0.000    0.000 multiarray.py:634(result_type)
        3    0.000    0.000    0.000    0.000 multiarray.py:706(dot)
        2    0.000    0.000    0.000    0.000 nanops.py:168(_get_fill_value)
        2    0.000    0.000    0.000    0.000 nanops.py:193(_maybe_get_mask)
        2    0.000    0.000    0.000    0.000 nanops.py:237(_get_values)
        2    0.000    0.000    0.000    0.000 nanops.py:322(_na_ok_dtype)
        2    0.000    0.000    0.000    0.000 nanops.py:396(nanany)
       18    0.000    0.000    0.000    0.000 numeric.py:150(is_all_dates)
        2    0.000    0.000    0.000    0.000 numeric.py:237(inferred_type)
        2    0.000    0.000    0.000    0.000 numeric.py:268(full)
      101    0.000    0.000    0.000    0.000 numeric.py:587(_flatnonzero_dispatcher)
      101    0.000    0.000    0.004    0.000 numeric.py:591(flatnonzero)
      412    0.001    0.000    0.001    0.000 numerictypes.py:286(issubclass_)
      206    0.001    0.000    0.002    0.000 numerictypes.py:360(issubdtype)
        1    0.000    0.000    1.703    1.703 pairwise.py:1351(_parallel_pairwise)
        1    0.000    0.000    0.000    0.000 pairwise.py:1420(_check_chunk_size)
        1    0.000    0.000    0.000    0.000 pairwise.py:1441(_precompute_metric_params)
        2    0.000    0.000    1.947    0.973 pairwise.py:1468(pairwise_distances_chunked)
        1    0.000    0.000    1.703    1.703 pairwise.py:1631(pairwise_distances)
        1    0.098    0.098    0.220    0.220 pairwise.py:201(euclidean_distances)
        1    0.068    0.068    1.703    1.703 pairwise.py:323(nan_euclidean_distances)
        3    0.000    0.000    0.000    0.000 pairwise.py:39(_return_float_dtype)
        2    0.000    0.000    0.002    0.001 pairwise.py:63(check_pairwise_arrays)
        1    0.000    0.000    0.000    0.000 parallel.py:373(effective_n_jobs)
        1    0.000    0.000    0.000    0.000 parallel.py:76(get_active_backend)
        1    0.000    0.000    0.000    0.000 process.py:198(daemon)
        1    0.000    0.000    0.000    0.000 process.py:37(current_process)
        1    0.000    0.000    1.971    1.971 profile:0(imputer.fit_transform(X))
        0    0.000             0.000          profile:0(profiler)
        2    0.000    0.000    0.000    0.000 range.py:316(dtype)
       64    0.000    0.000    0.000    0.000 range.py:683(__len__)
       20    0.000    0.000    0.009    0.000 series.py:201(__init__)
        4    0.000    0.000    0.000    0.000 series.py:381(_constructor)
       18    0.000    0.000    0.000    0.000 series.py:398(_set_axis)
        2    0.000    0.000    0.004    0.002 series.py:4071(apply)
        2    0.000    0.000    0.001    0.000 series.py:4211(_reduce)
        2    0.000    0.000    0.000    0.000 series.py:427(dtype)
       24    0.000    0.000    0.000    0.000 series.py:442(name)
       24    0.000    0.000    0.000    0.000 series.py:492(name)
       34    0.000    0.000    0.000    0.000 series.py:540(_values)
        6    0.000    0.000    0.000    0.000 series.py:595(__len__)
        1    0.000    0.000    0.000    0.000 threading.py:1306(current_thread)
        1    0.000    0.000    0.000    0.000 validation.py:1015(<listcomp>)
        8    0.000    0.000    0.000    0.000 validation.py:180(_num_samples)
        6    0.000    0.000    0.000    0.000 validation.py:390(_ensure_no_complex_data)
        6    0.001    0.000    0.024    0.004 validation.py:397(check_array)
      202    0.000    0.000    0.000    0.000 validation.py:532(<genexpr>)
     16/7    0.000    0.000    1.726    0.247 validation.py:59(inner_f)
       16    0.000    0.000    0.000    0.000 validation.py:71(<dictcomp>)
        6    0.000    0.000    0.001    0.000 validation.py:76(_assert_all_finite)
        1    0.000    0.000    0.000    0.000 validation.py:955(check_is_fitted)
        6    0.000    0.000    0.000    0.000 warnings.py:165(simplefilter)
        6    0.000    0.000    0.000    0.000 warnings.py:181(_add_filter)
        6    0.000    0.000    0.000    0.000 warnings.py:437(__init__)
        6    0.000    0.000    0.000    0.000 warnings.py:458(__enter__)
        6    0.000    0.000    0.000    0.000 warnings.py:477(__exit__)

[khughitt]

cProfiler output in a format usable by snakeviz:

test_knn.profile.gz

[jnothman]

What is the shape of the data your had issues with?

[jnothman]

I think it's very hard to see the bottlenecks when more of the runtime is taken by import than by the execution in question.

[jnothman]

Yes, a lot of time is spent in _calc_impute, which is called by process_chunk, which is called by pairwise_distances_chunked. process_chunk takes a greater fraction of the time with an increase n_features. And then there's the expected O(n^2) scaling with respect to n_samples. From the looks of it large amounts of time are spent in _calc_impute (mostly in argpartition), but also a lot in nan_euclidean_distances. nan_euclidean_distance might easily be made more performant with cythonisation. I'm not sure whether we can easily improve on the argpartition cost. I've not looked at the effect of n_neighbors.