How to use custom item identifiers in NearestNeighbors of sckit-learn

I use sckit-learn NearestNeighbors as a recommendation engine of movies, with Movie Lens database. The recommendation engine is an item-item based recommander (neighbors of one item are other items).

I'm putting a pickled version of the algorithm behind an API Flask application.

I work with Movie Lens IDs (identifiers) of movies, that are not contiguous. Ex : movie 1, movie 2, movie 3, movie 7, movie 11, etc (there are no movies 5, 6, 8, 9, 10). Those IDs are stored as the Int64Index of a Pandas Dataframe.

For allowing retrieval of proper movie neighbors, I need two things :

• to retrieve the correct movie vector based on its Movie Lens ID
• the NearestNeighbors algorithm must return those custom IDs as Neighbors, instead of the contiguous IDs of numpy array

Otherwise, I won't be able to make the relationship between the numpy IDs and the Movie Lens IDs.

Is there a way to do that ? Or do you think there is another approch that can fit my use case ?

I looked up the source code of NearestNeighbors.fit() method, and it looks like the Dataframe gets turned into a regular numpy array at some point, and "forget" about the custom ID.

I'm surprised there is no such an option, for that common use case. Maybe I am missing something.

Here is the code I use so far :

# The traning part :
# ==================

import pandas as pd
from sklearn.neighbors import NearestNeighbors
from scipy.sparse import csr_matrix
import pickle

df = pd.read_csv('movie_lens_ratings.csv')

# Here the custom movie IDs are stored as the index of the pivot DF :
df_pivot = df.pivot_table(index='movieId', columns='userId', values='rating').fillna(0)

sparse = csr_matrix(df_pivot.values)

knn = NearestNeighbors(metric='cosine', algorithm='brute', n_neighbors=5)
knn.fit(df_pivot)

pickle.dump(knn, open('knn_movie_lens.pkl', 'wb'))

# The API part :
# ==================

class KnnRecommender:
    def __init__(self, movie_lens_id):
        self.movie_lens_id = movie_lens_id

    def recommend(self):
        model: NearestNeighbors = pickle.load(open(f'knn_movie_lens.pkl', 'rb'))
        # Here I fetch the training data. It's a regular numpy array,
        # without any custom indices :
        data: csr_matrix = model._fit_X

        # This will return wrong movie vector, since self.movie_lens_id does not
        # match the numpy indices :
        movie_rating_vector: np.ndarray = data.getrow(self.movie_lens_id).toarray()

        # Here, the neighbors are contiguous numpy indices. I cannot use them
        # to retrieve proper movies from my database :
        distances, neighbors = model.kneighbors(movie_rating_vector)

        return neighbors

Thank you