I am not sure it is simple to write a vectorized and performant code that replicates the desired logic.
However, it is relatively simple to write it in a way that it is easy to accelerate with Numba.
Firstly, let us write your code as a (pure) function of the dataframe, returning the values to eventually put in df["Available to Promise"].
Eventually, it is easy to inglobate its result into the original dataframe with:
df["Available to Promise"] = calc_avail_OP(df)
The OP's code, save for exception handling and printing (and incorporation into the original dataframe as just discussed) is equivalent to the following:
import numpy as np
import pandas as pd
def calc_avail_OP(df):
temp_df = df[["Product", "Inventory"]].drop_duplicates()
temp_df["free"] = df["Inventory"]
result = np.zeros(len(df), dtype=df["Inventory"].dtype)
for i, row in df.iterrows():
available = (
temp_df[row["Product"] == temp_df["Product"]]["free"]
.reset_index(drop=True)
.iloc[0]
)
if available - row["Quantity"] >= 0:
result[i] = available - row["Quantity"]
a = (
temp_df.loc[row["Product"] == temp_df["Product"]]
.reset_index()["index"]
.iloc[0]
)
temp_df.at[a, "free"] = available - row["Quantity"]
else:
result[i] = available
return result
Now, if the input is sorted so that the unique products appear consecutively, the same can be achieved with a few scalar temporary variables on native NumPy objects, and this can be effectively accelerated with Numba:
import numba as nb
@nb.njit
def _calc_avail_nb(products, quantities, stocks):
n = len(products)
avails = np.empty(n, dtype=stocks.dtype)
last_product = products[0]
avail = stocks[0]
for i in range(n):
if products[i] != last_product:
last_product = products[i]
avail = stocks[i]
qty = quantities[i]
if avail >= qty:
avail -= qty
avails[i] = avail
return avails
def calc_avail_nb(df):
return _calc_avail_nb(
df["Product"].to_numpy(dtype="U"),
df["Quantity"].to_numpy(),
df["Inventory"].to_numpy()
)
If the input is not guaranteed to be sorted, one could keep track of inventory information with a dict():
import numba as nb
@nb.njit
def _calc_avail_dict_nb(products, quantities, stocks):
inventory = {products[0]: stocks[0]}
n = len(products)
avails = np.empty(n, dtype=stocks.dtype)
for i in range(n):
product = products[i]
avail = inventory.setdefault(products[i], stocks[i])
qty = quantities[i]
if avail >= qty:
avail -= qty
inventory[products[i]] = avail
avails[i] = avail
return avails
def calc_avail_dict_nb(df):
return _calc_avail_dict_nb(
df["Product"].to_numpy(dtype="U"),
df["Quantity"].to_numpy(),
df["Inventory"].to_numpy()
)
The following text include a comparison with some approaches from other answers:
def stock(val):
s = val
q = yield
while True:
s = s - q if s >= q else s
q = yield s
def exaust_stock(df):
st = stock(df.iloc[0]['Inventory']).send
st(None)
return df['Quantity'].apply(st)
def calc_avail_gen(df):
return (
df
.groupby('Product')
.apply(exaust_stock)
.reset_index(level=0, drop=True)
.to_numpy()
)
@nb.njit
def _calc_avail_grouped_nb(quant, inv):
stock = inv[0]
n = len(quant)
out = np.zeros((n,), dtype=np.int_)
for i in range(n):
if stock > 0 and quant[i] <= stock:
stock -= quant[i]
out[i] = stock
else:
out[i] = stock
return out
def calc_avail_grouped_nb(df):
return (
df
.groupby('Product')
.apply(lambda x: _calc_avail_grouped_nb(x['Quantity'].to_numpy(), x['Inventory'].to_numpy()))
.explode()
.to_numpy(dtype=np.int_)
)
The test indicate that while they do provide the same results, calc_avail_nb() and calc_avail_dict_nb() provide a speed increase of ~200x on the test input.
data = {
'Customer':[1,2,3,4,5,6,7,8],
'Product':['A','A','A','A','B','B','B','B'],
'Quantity':[100,1000,700,50,20,50,40,30],
'Inventory':[800,800,800,800,100,100,100,100]
}
df = pd.DataFrame(data)
funcs = calc_avail_OP, calc_avail_nb, calc_avail_dict_nb, calc_avail_gen, calc_avail_grouped_nb
base = funcs[0](df)
timings = {}
n = len(df)
timings[n] = []
for func in funcs:
res = func(df)
is_good = np.allclose(base, res)
timed = %timeit -n 8 -r 8 -q -o func(df)
is_good = True
timing = timed.best * 1e6
timings[n].append(timing if is_good else None)
print(f"{func.__name__:>24} {is_good!s:5} {timing:10.3f} µs {timings[n][0] / timing:5.1f}x")
# calc_avail_OP True 11699.373 µs 1.0x
# calc_avail_nb True 52.821 µs 221.5x
# calc_avail_dict_nb True 57.198 µs 204.5x
# calc_avail_gen True 3360.806 µs 3.5x
# calc_avail_grouped_nb True 1099.665 µs 10.6x
Similar tests on larger inputs seem to point to an even larger speed gain.
The timings are computed with the following:
import string
import random
def gen_df(n, m=None, max_stock=None):
if not m:
m = 2 + n // 16
if not max_stock:
max_stock = n
k = n.bit_length()
inventory = {
"".join(
random.choices(string.ascii_letters, k=random.randint(1, 2 + k))
): random.randint(max_stock // 2, max_stock)
for _ in range(m)
}
products = random.choices(list(inventory.keys()), k=n)
return pd.DataFrame(
{
"Customer": np.random.randint(1, int(1.1 * max_stock), n),
"Product": products,
"Quantity": np.random.randint(1, int(1.1 * max_stock), n),
"Inventory": [inventory[product] for product in products],
}
)
np.random.seed(0)
random.seed(0)
timings = {}
for i in range(3, 18, 3):
n = 2 ** i
print(f"i={i}, n={n}")
df = gen_df(n)
base = funcs[0](df)
timings[n] = []
for func in funcs:
res = func(df)
is_good = np.allclose(base, res)
timed = %timeit -n 1 -r 1 -q -o func(df)
is_good = True
timing = timed.best * 1e3
timings[n].append(timing if is_good else None)
print(f"{func.__name__:>24} {is_good!s:5} {timing:10.3f} ms {timings[n][0] / timing:5.1f}x")
and plotted with:
import pandas as pd
import matplotlib.pyplot as plt
df = pd.DataFrame(data=timings, index=[func.__name__ for func in funcs]).transpose()
df.plot(marker='o', xlabel='Input Size / #', ylabel='Best timing / µs', figsize=(6, 4))
fig = plt.gcf()
fig.patch.set_facecolor('white')
df = pd.DataFrame(data=timings, index=[func.__name__ for func in funcs]).transpose()
df = df[[funcs[0].__name__]].to_numpy() / df
df.plot(marker='o', xlabel='Input Size / #', ylabel='Speed increase / %x', figsize=(6, 4))
fig = plt.gcf()
fig.patch.set_facecolor('white')
to obtain, respectively:
and
