Timing DataFrame Filters

from dataclasses import dataclass, field
from contextlib import contextmanager
from time import perf_counter

@dataclass
class Timer:
    start: float = None
    end: float = None
    
    @property
    def elapsed(self):
        if self.start is None or self.end is None:
            raise ValueError('Timer must have both end and start')
        return self.end - self.start

@dataclass
class TimerManager:
    registry: list = field(default_factory=list)
    
    @contextmanager
    def time(self, description):
        timer =  Timer(start=perf_counter())
        yield timer
        timer.end = perf_counter()
        self.registry.append((description, timer))

from pandas import Timestamp
pd_start = Timestamp(2000, 4, 1)
pd_stop = Timestamp(2000, 4, 2)

timer = TimerManager()

# 1.
with timer.time('pd.loc[start <= date < end]'):
    pd_df.loc[
        (pd_start <= pd_df['date']) & (pd_df['date'] < pd_stop)
    ]
    
# 2.
with timer.time('pd.loc[date.between(…)]'):
    pd_df.loc[
        pd_df['date'].between(pd_start, pd_stop, inclusive='left')
    ]
    
# 3.
with timer.time('pd.loc[date] {cold}'):
    pd_df_idx.loc['2000-06-01']
    
# 4. 
# this performs MUCH faster than the above, it is a direct output
# caching because there is still a speed difference even with different inputs
with timer.time('pd.loc[date] {warm}'):
    pd_df_idx.loc['2000-04-01']
    
for desc, t in timer.registry:
    if desc.startswith('pd'):
        print(f'{desc: >35} → {t.elapsed*1000:.6f}µs')

from polars import datetime
pl_start = datetime(2000, 4, 1)
pl_stop = datetime(2000, 4, 2)

with timer.time('pl.filter(start <= date < stop)'):
    pl_df.filter(
        (pl_start <= col('date')) & (col('date') < pl_stop)
    )
    
with timer.time('pl.lazy.filter(start <= date < stop)'):
    pl_lazy.filter(
        (pl_start <= col('date')) & (col('date') < pl_stop)
    ).collect()
    
with timer.time('pl.filter(date.date == start)'):
    pl_df.filter(
        (col('date').dt.date() == pl_start.dt.date())
    )
    
with timer.time('pl.lazy.filter(date.date == start)'):
    pl_lazy.filter(
        (col('date').dt.date() == pl_start.dt.date())
    ).collect()
    
with timer.time('pl[idx_start:idx_stop]'):
    pl_df[
        slice(
            *pl_df.select(
                idx_start=col('date').search_sorted(pl_start),
                idx_stop =col('date').search_sorted(pl_stop)
            ).row(0)
        )
    ]
    
with timer.time('pl[lazy.idx_start:lazy.idx_stop]'):
    pl_df[
        slice(
            *pl_lazy.select(
                idx_start=col('date').search_sorted(pl_start),
                idx_stop=col('date').search_sorted(pl_stop)
            )
            .collect()
            .row(0)
        )
    ]
    
for desc, t in timer.registry:
    if desc.startswith('pl'):
        print(f'{desc: >35} → {t.elapsed*1000:.6f}µs')

df = (
    DataFrame(timer.registry, columns=['description', 'timer'])
    .assign(
        elapsed_us=lambda d: 1_000 * d['timer'].map(lambda t: t.elapsed),
        package=lambda d: d['description'].str.extract('^(pl|pd)'),
    )
    .drop(columns='timer')
    .sort_values('elapsed_us', ascending=False)
)

df

%matplotlib inline

from matplotlib.pyplot import rc, setp
from flexitext import flexitext

rc('figure', figsize=(10, 6), facecolor='white')
rc('font', size=12)
rc('axes.spines', top=False, right=False, left=False)


palette = {
    'pl': '#1F77B4FF',
    'pd': '#FF7F0EFF',
}

ax = df.plot.barh(
    x='description', y='elapsed_us', legend=False, width=.8,
    color=df['package'].map(palette),
)
ax.set_ylabel('')
ax.yaxis.set_tick_params(length=0)
ax.bar_label(ax.containers[0], fmt='{:.2f}', padding=5)
ax.set_xlabel(r'Duration (µs)')
setp(ax.get_yticklabels(), va='center')

ax.figure.tight_layout()
left_x = min(text.get_tightbbox().x0 for text in ax.get_yticklabels())
x,_ = ax.transAxes.inverted().transform([left_x, 0])

annot = flexitext(
    s='<size:xx-large,weight:semibold>Time Elapsed for Datetime Filtering in'
     f' <color:{palette["pl"]}>Polars <size:medium>{polars.__version__}</></>'
     f' vs <color:{palette["pd"]}>pandas <size:medium>{pandas.__version__}</></>'
      ' </>',
    x=x, y=1.02, va='bottom', ha='left',
);