How to Use PySpark, OR-Tools, Pyomo, and Gurobi for Large-Scale Optimization: A Simple Guide

The Vehicle Routing Problem (VRP).

Scenario: A distribution company has a central depot and a fleet of vehicles to deliver goods to various customers at different locations. Each vehicle has a maximum capacity (e.g., weight or volume), and each customer has a specific demand. The goal is to find the optimal routes for each vehicle so that all customers are served, vehicle capacities are not exceeded, and the total distance traveled is minimized.

To make this concrete, let’s build a small vehicle routing example. First, make sure you have the necessary libraries installed:

pip install pyspark pyomo pandas

# For the solver, we'll use CBC, which is open-source and works well with Pyomo.
# On Debian/Ubuntu:
sudo apt-get install -y coinor-cbc

# For other systems (like Windows or macOS), you might need to install it via conda
# or download binaries.

Step 1: Preparing Data with PySpark

First, let’s simulate our customer data. We’ll have their location (x, y coordinates) and their demand. Although the data is small for this example, in a real-world scenario, it could come from a large database or data lake, which is where PySpark shines.

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, IntegerType, FloatType
import pandas as pd

# Initialize the SparkSession
spark = SparkSession.builder.appName("RouteOptimizationData").getOrCreate()

# Customer data: ID, Coordinates (x, y), and Demand
data = [
    (0, 50, 50, 0),  # Depot
    (1, 30, 40, 15),
    (2, 60, 20, 10),
    (3, 80, 50, 20),
    (4, 40, 80, 12),
    (5, 70, 70, 18)
]

# Define the schema for the DataFrame
schema = StructType([
    StructField("id", IntegerType()),
    StructField("x", FloatType()),
    StructField("y", FloatType()),
    StructField("demand", IntegerType())
])

# Create the Spark DataFrame
customer_df = spark.createDataFrame(data, schema)

print("Customer data loaded with PySpark:")
customer_df.show()

# In a real scenario, you would perform complex transformations here.
# For our model, we'll collect the data into a Pandas DataFrame.
customer_data_pd = customer_df.toPandas()

# Vehicle parameters
num_vehicles = 2
vehicle_capacity = 40

spark.stop()

What this code does:

• It starts a SparkSession, the entry point for any Spark functionality.
• It creates a DataFrame with our customer data, including the depot (ID 0).
• In a real project, this is the stage where you would handle gigabytes or terabytes of data to clean and prepare it.
• We collect the data into a format (pandas.DataFrame) that Pyomo can easily use for modeling.

Step 2: Modeling the Problem with Pyomo

Now, we’ll use Pyomo to build the mathematical model of our routing problem. We’ll calculate a distance matrix and then define our variables, objective, and constraints.

import pyomo.environ as pyo
import math

# --- Prepare data for the model ---
locations = customer_data_pd.set_index('id')[['x', 'y']].to_dict('index')
demands = customer_data_pd.set_index('id')['demand'].to_dict()
customers = [i for i in demands if i != 0] # Exclude the depot
depot = 0
N = len(locations) # Total number of locations (depot + customers)

# Function to calculate Euclidean distance
def distance(p1, p2):
    return math.sqrt((p1['x'] - p2['x'])**2 + (p1['y'] - p2['y'])**2)

# Create the distance matrix
dist_matrix = {(i, j): distance(locations[i], locations[j]) for i in locations for j in locations}

# --- Build the Pyomo Model ---
model = pyo.ConcreteModel()

# --- Sets ---
model.C = pyo.Set(initialize=customers)      # Set of customers
model.V = pyo.Set(initialize=range(num_vehicles)) # Set of vehicles
model.N = pyo.Set(initialize=locations.keys())   # Set of all nodes (customers + depot)

# --- Variables ---
# x[i, j, k] is 1 if vehicle k travels from node i to node j
model.x = pyo.Var(model.N, model.N, model.V, within=pyo.Binary)

# --- Objective Function ---
# Minimize the total distance traveled by all vehicles
def obj_rule(model):
    return sum(dist_matrix[i, j] * model.x[i, j, k] for i in model.N for j in model.N for k in model.V)
model.objective = pyo.Objective(rule=obj_rule, sense=pyo.minimize)

# --- Constraints ---
# 1. Each customer must be visited exactly once by one vehicle
def serve_once_rule(model, c):
    return sum(model.x[i, c, k] for i in model.N for k in model.V) == 1
model.serve_once = pyo.Constraint(model.C, rule=serve_once_rule)

# 2. The capacity of each vehicle cannot be exceeded
def capacity_rule(model, k):
    return sum(demands[c] * sum(model.x[i, c, k] for i in model.N) for c in model.C) <= vehicle_capacity
model.capacity = pyo.Constraint(model.V, rule=capacity_rule)

# 3. Every vehicle must leave the depot
def leave_depot_rule(model, k):
    return sum(model.x[depot, j, k] for j in model.C) == 1
model.leave_depot = pyo.Constraint(model.V, rule=leave_depot_rule)

# 4. Every vehicle must return to the depot
def return_depot_rule(model, k):
    return sum(model.x[i, depot, k] for i in model.C) == 1
model.return_depot = pyo.Constraint(model.V, rule=return_depot_rule)

# 5. Flow conservation (if a vehicle enters a customer node, it must leave it)
def flow_rule(model, c, k):
    in_flow = sum(model.x[i, c, k] for i in model.N)
    out_flow = sum(model.x[c, j, k] for j in model.N)
    return in_flow == out_flow
model.flow = pyo.Constraint(model.C, model.V, rule=flow_rule)

print("Pyomo model built successfully!")

What this code does:

• It calculates the distances between all points.
• It defines Sets, which are the building blocks of the model (customers, vehicles, nodes).
• It defines the Decision Variables: x[i, j, k] will be 1 if vehicle k goes from i to j, and 0 otherwise.
• It establishes the Objective: to minimize the sum of all distances traveled.
• It adds the Constraints, which are the business rules (visit each customer, respect capacity, etc.).

Step 3: Solving the Model

Now, we pass our Pyomo model to a solver. We’ll use CBC, an open-source solver for mixed-integer programming. If you had Gurobi installed, you could simply change the solver’s name.

# --- Solve ---
# solver = pyo.SolverFactory('gurobi') # Example if you had Gurobi
# solver = pyo.SolverFactory('ortools') # Example if you wanted to use OR-Tools
solver = pyo.SolverFactory('cbc') # Using the open-source CBC solver
results = solver.solve(model, tee=True) # tee=True shows the solver's log

# Check the solution status
if (results.solver.status == pyo.SolverStatus.ok) and \
   (results.solver.termination_condition == pyo.TerminationCondition.optimal):
    print("Optimal solution found!")
else:
    print("No optimal solution found. Status:", results.solver.termination_condition)

Step 4: Analyzing the Results

Finally, we extract and display the routes found by the solver in a human-readable format.

# --- Present the Results ---
print("\n--- Optimized Routes ---")
total_distance = pyo.value(model.objective)
print(f"Total Distance Traveled: {total_distance:.2f}\n")

for k in model.V:
    route = []
    current_location = depot

    while True:
        route.append(current_location)
        found_next = False
        # Find the next location in the route for vehicle k
        for j in model.N:
            if current_location != j and pyo.value(model.x[current_location, j, k]) > 0.9:
                current_location = j
                found_next = True
                break
        if not found_next or current_location == depot:
            # If no next location is found or we returned to the depot
            if depot not in route:
                 route.append(depot)
            break

    # Calculate route demand and distance
    route_demand = sum(demands.get(city, 0) for city in route)
    route_dist = sum(dist_matrix[route[i], route[i+1]] for i in range(len(route)-1))

    # Print only if the vehicle was used
    if len(route) > 2: # Route is more than just [depot, depot]
        print(f"Vehicle {k}:")
        print(f"  Route: {' -> '.join(map(str, route))}")
        print(f"  Demand Served: {route_demand} (Capacity: {vehicle_capacity})")
        print(f"  Route Distance: {route_dist:.2f}\n")

What this code does:

• It extracts the value of the objective function (the minimum total distance).
• It iterates through each vehicle to reconstruct the route found by the solver, following the x[i, j, k] variables that were set to 1.
• It prints the route, total demand served, and distance for each vehicle, providing a clear action plan.

Conclusion: Turning Data into Intelligent Decisions

In this article, we journeyed from large-scale data preparation with PySpark, to elegant mathematical modeling with Pyomo, and finally, to solving complex problems with powerful solvers like CBC, OR-Tools, or Gurobi. As we saw in our hands-on vehicle routing example, combining these tools creates a robust and scalable workflow for tackling real-world challenges.

The true power lies not in any single tool, but in their synergy. PySpark handles data volumes that would overwhelm traditional systems. Pyomo provides the flexibility to describe any business problem in a structured language. And the solvers do the heavy lifting of finding the best mathematical solution.

If you are just getting started, don’t be intimidated. Begin with smaller problems, like the one we demonstrated, and experiment with each component. Swap out the solver, add new constraints to your model, and see how the solutions change. By mastering this toolkit, you are not just writing code — you are building a decision engine capable of driving cost savings, boosting efficiency, and providing a significant competitive edge to any business.

The field of optimization is vast and constantly evolving. The skills you’ve begun to develop here are more valuable than ever. Keep exploring, keep learning, and most importantly, keep applying this knowledge to solve problems that matter.