Accelerating K-NN: A Guide to Python Concurrency and Parallelism (Part 1)

The Bottleneck

As practitioners, we value K-NN for its simplicity, but its brute-force nature makes it notoriously slow on large datasets. This series will explore how to smash this performance barrier using Python’s concurrency and parallelism tools. First, let’s establish a shared vocabulary for the concepts we’ll be discussing.

A Bit of Jargon 🧐

To navigate the world of performance optimization, we need to be precise with our language. Here are the core concepts:

• Concurrency: The ability to manage multiple tasks at once. A single-core system achieves this by interleaving tasks, a process known as multitasking. The tasks don’t run simultaneously, but all make progress.
• Parallelism: The ability to execute multiple tasks simultaneously. This requires a system with multiple execution units, like a multi-core CPU or a GPU. Parallelism is a specific form of concurrency.
• Process: An instance of a running program with its own private memory space, managed by the operating system. Processes are isolated from each other and are ideal for leveraging multiple CPU cores. Communication between them requires serialization (pickling) of data, which has an overhead.
• Thread: An execution unit that lives inside a process. Threads within the same process share memory, which makes data sharing easy but also risks data corruption without proper synchronization (e.g., using locks).
• Coroutine: A special type of function that can be paused and resumed, allowing other tasks to run. Coroutines enable concurrency within a single thread, managed by an event loop. They are lightweight and excel at handling thousands of I/O-bound operations.
• Global Interpreter Lock (GIL): A mutex (a type of lock) that protects access to Python objects, preventing multiple native threads from executing Python bytecode at the same time within a single process. Because of the GIL, multithreading in CPython cannot achieve true parallelism for CPU-bound Python code. However, threads are still highly effective for I/O-bound tasks, as the GIL is released during I/O wait times.

With this foundation, let’s examine why K-NN is a prime candidate for these techniques.

Brute Force and Its Complexity

K-NN classifies a new data point by finding the ‘K’ closest labeled points in its training set and taking a majority vote. “Closeness” is typically defined by a distance metric like the Euclidean distance. For a brute-force K-NN, the time complexity of the prediction phase is:

O(M⋅N⋅D)

Where:

• M is the number of query points to classify.
• N is the number of data points in the training set.
• D is the number of dimensions (features) of the data.

This formula tells us that for every new point (M), the algorithm computes its distance to every single point in the training data (N), and each distance calculation involves operations across all dimensions (D). This multiplicative effect makes the naive implementation impractical for many real-world applications.

Let’s break down the prediction for a single query point, x_q:

1. Distance Calculation: Compute the distance d(x_q,x_i) for all training points x_i. This is a pure CPU-bound task. Crucially, each calculation is independent of the others, making this step embarrassingly parallel.
2. Neighbor Identification: Sort the distances to find the K smallest values.
3. Majority Vote: Find the most common label among those K neighbors to make a prediction.

Step 1 consumes the vast majority of the runtime. This is the target for our optimization efforts.

Prerequisites & Installation

To follow along with the code examples run the command below in your terminal.

pip install numpy scikit-learn matplotlib

A Baseline Implementation

Before we can improve, we must measure. Let’s create a baseline single-threaded K-NN in pure Python and NumPy to establish our benchmark. The code iterates through each test point and then through every training point to find its neighbors.

import numpy as np
import time
from collections import Counter
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

# --- KNN Functions ---
def euclidean_distance(p1, p2):
    """Calculate the Euclidean distance between two points."""
    return np.sqrt(np.sum((p1 - p2) ** 2))

def predict_single(x_query, X_train, y_train, k):
    """Predict the label for a single query point."""
    distances = [euclidean_distance(x_query, x_t) for x_t in X_train]
    k_nearest_indices = np.argsort(distances)[:k]
    k_nearest_labels = [y_train[i] for i in k_nearest_indices]
    return Counter(k_nearest_labels).most_common(1)[0][0]

def knn_predict_baseline(X_test, X_train, y_train, k):
    """Run KNN prediction for a set of test points."""
    predictions = [predict_single(x_query, X_train, y_train, k) for x_query in X_test]
    return np.array(predictions)

# --- Initial Performance Measurement ---
n_samples = 5000
n_features = 5
X, y = make_classification(n_samples=n_samples, n_features=n_features, n_informative=3, n_classes=3, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

k = 5
start_time = time.time()
y_pred_baseline = knn_predict_baseline(X_test, X_train, y_train, k)
end_time = time.time()
baseline_duration = end_time - start_time

print(f"Training set size: {X_train.shape}")
print(f"Test set size: {X_test.shape}")
print(f"\nBaseline K-NN execution time for {X_train.shape[0]} samples: {baseline_duration:.4f} seconds")

On an Apple M2 machine, the execution time is significant.

Baseline K-NN execution time: 10.001 seconds

Visualizing the Bottleneck

The O(M⋅N⋅D) complexity implies that runtime should scale quadratically with the number of samples N and M increasing while D is constant. Let’s verify this empirically by running our baseline function on datasets of increasing size and plotting the results.

sample_sizes = [2000, 4000, 6000, 8000, 10000]
execution_times = []
print("\n--- Measuring runtime for different dataset sizes ---")
for n_samples in sample_sizes:
    # Generate data
    X, y = make_classification(n_samples=n_samples, n_features=20, n_informative=10, n_classes=3, random_state=42)
    # Use 80% for training, 20% for testing
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # Time the prediction
    start = time.time()
    knn_predict_baseline(X_test, X_train, y_train, k=5)
    end = time.time()

    duration = end - start
    execution_times.append(duration)
    print(f"Size: {X_train.shape[0]:<5} samples | Time: {duration:.4f} seconds")
# --- Plotting the results ---
plt.figure(figsize=(10, 6))
plt.plot(np.array(sample_sizes) * 0.8, execution_times, marker='o', linestyle='-', color='r')
plt.title('KNN Baseline Performance: Runtime vs. Dataset Size', fontsize=16)
plt.xlabel('Number of Training Samples (N)', fontsize=12)
plt.ylabel('Execution Time (seconds)', fontsize=12)
plt.grid(True)
plt.show()

The output confirms our proposition: the relationship appears quadratic.

This graph is the problem we need to solve. A system whose performance degrades so predictably and steeply is not scalable.

Our Optimization Toolkit 🛠️

To attack this problem, Python provides a powerful set of libraries.

• Multiprocessing: The primary tool for CPU-bound parallelism. It bypasses the GIL by creating new processes, allowing Python code to run on multiple cores simultaneously.
• Threading: Best for I/O-bound concurrency. Threads are lightweight but are constrained by the GIL for CPU-bound tasks.
• concurrent.futures: A high-level API over the previous two, offering a simple way to manage pools of processes or threads.
• Joblib: A library optimized for parallelizing loops in scientific computing. It’s particularly efficient with NumPy arrays.
• Asyncio: A framework for single-threaded concurrency using coroutines, ideal for managing thousands of non-blocking I/O operations.

In Part 2, we will use these tools on our K-NN problem. We’ll start with process-based parallelism to improve our models capability. As a preview, using 2 cores for our baseline implementation returns a 1.60X performance improvement.

from joblib import Parallel, delayed

# --- Performance Measurement ---
start_time = time.time()
y_pred_joblib = Parallel(n_jobs=2)(
    delayed(predict_single)(x, X_train, y_train, 5) for x in X_test
)

end_time = time.time()

duration = end_time - start_time
print(f"\nJoblib execution time: {duration:.4f} seconds") #Joblib execution time: 6.2325 seconds
print(f"Speedup vs. Baseline: {baseline_duration  / duration:.2f}x") #Speedup vs. Baseline: 1.60x

Baseline K-NN execution time: 6.2325 seconds