--- title: "From O(n) to O(1): How We Made DataFrame Lookups 23x Faster" description: "We optimize APIs. We scale databases. We architect systems for speed.\nBut sometimes the biggest performance wins hide in plain sight — in a single function called thousands of times.\n" date: 2026-01-08T00:00:00Z tags: - "python" - "pandas" - "performance" - "optimization" - "dataframes" - "engineering" categories: - "python" - "performance" - "engineering" author: ["Pedro Dias"] canonical: https://itspedrodias.com/posts/data-processing/ --- # From O(n) to O(1): How We Made DataFrame Lookups 23x Faster Recently, I was working on a data processing pipeline that needed to replicate configurations across multiple entities. The script worked fine during development, but when we scaled to production data volumes, execution time became a bottleneck. What should have taken seconds was taking minutes. Here's how I identified the problem and fixed it with a simple but powerful optimization. ## The Problem The pipeline processes a CSV file containing thousands of rows and needs to find matching records between different subsets of data. The original code used pandas' [`query()` method](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.query.html) to find matches: ```python MATCHING_COLUMNS = [ 'entity_id', 'external_id', 'name', 'category', 'value', 'label', 'status', 'phase' ] def find_matching_row(source_row, target_df, logger): """Original implementation - O(n) per lookup.""" query_parts = [] for col in MATCHING_COLUMNS: val = source_row[col] if pd.isna(val): query_parts.append(f"`{col}`.isnull()") elif isinstance(val, str): escaped_val = val.replace("'", "\\'") query_parts.append(f"`{col}` == '{escaped_val}'") else: query_parts.append(f"`{col}` == {val}") query = " and ".join(query_parts) try: matching_rows = target_df.query(query) except Exception as e: logger.error(f"Query failed: {e}") return None if matching_rows.empty: return None return matching_rows.iloc[0] ``` This looks reasonable, right? It dynamically builds a query string and uses pandas' built-in [`query()` method](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.query.html). But there's a hidden performance trap here. ## Why This Is Slow Every call to `df.query()`: 1. **Parses the query string** - pandas needs to interpret the string expression 2. **Evaluates against every row** - it scans the entire DataFrame ([O(n) complexity](https://en.wikipedia.org/wiki/Big_O_notation)) 3. **Handles edge cases** - escaping, null checks, type coercion When you call this function once, it's barely noticeable. But in my case, I was calling it **thousands of times** in nested loops: ```python for entity in entities: # ~30 entities for config in configs: # ~50 configs per entity for item in items: # ~10 items per config find_matching_row(item, target_df, logger) # Called 15,000 times! ``` With a DataFrame of 5,000+ rows, each query scans all rows. That's **75 million row comparisons** just for the lookups. ## The Solution: Multi-Index Lookups The fix was to **pre-build an index** and use [O(1) hash-based lookups](https://en.wikipedia.org/wiki/Big_O_notation) instead of O(n) scans. ### Step 1: Create an Indexed DataFrame ```python def prepare_lookup_index(df): """ Create a multi-index DataFrame for fast O(1) lookups. Returns tuple of (indexed_df, sentinel_value). """ df_indexed = df.copy() # Handle NaN values with a sentinel (can't use NaN in index) sentinel = '__NAN_SENTINEL__' for col in MATCHING_COLUMNS: if col in df_indexed.columns: df_indexed[col] = df_indexed[col].fillna(sentinel) # Set multi-index for O(1) lookups # See pandas MultiIndex docs: https://pandas.pydata.org/docs/reference/api/pandas.MultiIndex.html df_indexed = df_indexed.set_index(MATCHING_COLUMNS, drop=False) df_indexed.index = pd.MultiIndex.from_tuples( df_indexed.index, names=MATCHING_COLUMNS ) return df_indexed, sentinel ``` ### Step 2: Build Lookup Keys ```python def build_lookup_key(source_row, sentinel): """Build a tuple key for multi-index lookup.""" key = [] for col in MATCHING_COLUMNS: val = source_row[col] if pd.isna(val): key.append(sentinel) else: key.append(val) return tuple(key) ``` ### Step 3: Use Direct Index Access ```python def find_matching_row_fast(source_row, target_indexed_df, sentinel, logger): """ Optimized implementation - O(1) per lookup. Uses hash-based index lookup via pandas `.loc[]` instead of query scanning. See: https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.loc.html """ lookup_key = build_lookup_key(source_row, sentinel) try: matching_rows = target_indexed_df.loc[lookup_key] except KeyError: return None except Exception as e: logger.error(f"Lookup failed: {e}") return None # Handle single match (Series) vs multiple matches (DataFrame) if isinstance(matching_rows, pd.Series): return matching_rows if matching_rows.empty: return None return matching_rows.iloc[0] ``` ### Step 4: Prepare Index Once, Use Many Times The key insight is to **prepare the index once** before the loop, not inside it: ```python # Prepare index ONCE before processing target_indexed, sentinel = prepare_lookup_index(target_df) for entity in entities: for config in configs: for item in items: # Now each lookup is O(1) instead of O(n) find_matching_row_fast(item, target_indexed, sentinel, logger) ``` ## Bonus: Dictionary Lookups for Simple Keys For lookups based on just a few columns (like ID fields), a plain Python dictionary is even faster: ```python def prepare_dict_index(df): """ Create a dictionary for O(1) lookups by composite key. """ lookup_dict = {} for idx, row in df.iterrows(): key = ( str(row['entity_id']) if pd.notna(row['entity_id']) else '', str(row['item_id']) if pd.notna(row['item_id']) else '', str(row['config_id']) if pd.notna(row['config_id']) else '' ) if key not in lookup_dict: lookup_dict[key] = row return lookup_dict # Usage source_dict = prepare_dict_index(source_df) # O(1) lookup row = source_dict.get(('123', '456', '789')) ``` ## Performance Results I added timing instrumentation to measure the real-world impact: ```python import time class ExecutionTimer: def __init__(self, logger): self.logger = logger self.timings = {} self.start_time = None def start(self, name: str): self.timings[name] = {'start': time.perf_counter()} def stop(self, name: str): if name in self.timings: self.timings[name]['elapsed'] = ( time.perf_counter() - self.timings[name]['start'] ) def print_summary(self): total = time.perf_counter() - self.start_time print("=" * 60) print("EXECUTION TIME SUMMARY") print("=" * 60) for name, data in self.timings.items(): elapsed = data['elapsed'] pct = (elapsed / total * 100) if total > 0 else 0 print(f" {name}: {elapsed:.3f}s ({pct:.1f}%)") print(f" TOTAL: {total:.3f}s") ``` ### Before Optimization ``` ============================================================ EXECUTION TIME SUMMARY (Original) ============================================================ Load data: 0.045s (0.2%) Generate combinations: 12.340s (52.1%) Replicate to entities: 11.250s (47.5%) Write output: 0.052s (0.2%) ------------------------------------------------------------ TOTAL: 23.687s ============================================================ ``` ### After Optimization ``` ============================================================ EXECUTION TIME SUMMARY (Optimized) ============================================================ Load data: 0.045s (0.4%) Generate combinations: 12.340s (94.8%) Index preparation: 0.089s (0.7%) Replicate to entities: 0.480s (3.7%) Write output: 0.052s (0.4%) ------------------------------------------------------------ TOTAL: 13.006s ============================================================ ``` **The replication step went from 11.25s to 0.48s — a 23x speedup!** The overall script execution improved by **45%**, and the bottleneck shifted from DataFrame lookups to the actual data generation logic. ## Key Takeaways ### 1. Profile Before Optimizing I added very basic timing instrumentation to identify where time was actually being spent. Without measurement, I might have optimized the wrong thing. For more sophisticated profiling, tools like Python's [`cProfile`](https://docs.python.org/3/library/profile.html) or [`line_profiler`](https://github.com/pyutils/line_profiler) can provide detailed insights. ### 2. Understand Your Data Access Patterns - **Single lookup**: `df.query()` is fine - **Many lookups on same DataFrame**: Pre-build an index - **Lookups by simple keys**: Use a dictionary ### 3. Trade Memory for Speed The indexed DataFrame and dictionary use more memory, but the time savings are massive. In my case: - Extra memory: ~2x the DataFrame size - Time savings: 23x faster ### 4. Index Preparation Is an Investment Building the index takes time (0.089s in my case), but it pays off after just a few lookups. The break-even point is around 20 lookups for my data size. ### 5. Test Your Optimization I kept the original implementation to validate that the optimized version produced identical outputs: ```bash # Run both versions python process_data_old.py --output old_output.json python process_data.py --output new_output.json # Compare outputs diff <(cat old_output.json | python -m json.tool --sort-keys) \ <(cat new_output.json | python -m json.tool --sort-keys) ``` ## When NOT to Use This This optimization isn't always the right choice: - **Single or few lookups**: The index preparation overhead isn't worth it - **Frequently changing data**: Rebuilding the index negates the benefits - **Memory-constrained environments**: The extra memory usage might be prohibitive - **Simple column filtering**: `df[df['col'] == value]` is already optimized ## Conclusion Sometimes the biggest performance wins come from understanding the [algorithmic complexity](https://en.wikipedia.org/wiki/Time_complexity) of your code. The original `df.query()` approach was O(n) per lookup. With 15,000 lookups on a 5,000-row DataFrame, that's 75 million operations. By switching to indexed lookups, each operation became O(1), reducing those 75 million operations to just 15,000 hash lookups plus the one-time cost of building the index. This is similar to how [database indexes](https://en.wikipedia.org/wiki/Database_index) work — they're an upfront investment that pays dividends on repeated queries. The lesson? **When you're doing repeated lookups on the same data, build an index first.** --- ## Code Reference For convenience, here's the full implementation as a reusable module. ```python # df_indexer.py - Reusable DataFrame indexing utilities import pandas as pd from typing import Dict, List, Tuple, Any def prepare_lookup_index( df: pd.DataFrame, index_columns: List[str], sentinel: str = '__NAN_SENTINEL__' ) -> Tuple[pd.DataFrame, str]: """ Create a multi-index DataFrame for fast O(1) lookups. Args: df: Source DataFrame index_columns: Columns to use for the multi-index sentinel: Value to replace NaN (can't use NaN in index) Returns: Tuple of (indexed_df, sentinel_value) """ df_indexed = df.copy() for col in index_columns: if col in df_indexed.columns: df_indexed[col] = df_indexed[col].fillna(sentinel) df_indexed = df_indexed.set_index(index_columns, drop=False) df_indexed.index = pd.MultiIndex.from_tuples( df_indexed.index, names=index_columns ) return df_indexed, sentinel def prepare_dict_index( df: pd.DataFrame, key_columns: List[str] ) -> Dict[Tuple, pd.Series]: """ Create a dictionary for O(1) lookups by composite key. Args: df: Source DataFrame key_columns: Columns to use for the lookup key Returns: Dictionary mapping tuple keys to DataFrame rows """ lookup_dict = {} for idx, row in df.iterrows(): key = tuple( str(row[col]) if pd.notna(row[col]) else '' for col in key_columns ) if key not in lookup_dict: lookup_dict[key] = row return lookup_dict def build_lookup_key( row: pd.Series, key_columns: List[str], sentinel: str = '__NAN_SENTINEL__' ) -> Tuple: """Build a tuple key for multi-index lookup.""" return tuple( sentinel if pd.isna(row[col]) else row[col] for col in key_columns ) ``` --- ## ✏️ Personal Notes - This optimization came from a real-world scenario where a script that "worked fine" during development became painfully slow at production scale. The fix was straightforward once I knew where the issue was — but only after I took the time to actually measure where the time was going. - If there's one thing to keep in mind, it's this: **add timing instrumentation before optimizing**. Without data, you're just guessing. And guessing usually means optimizing the wrong thing. Python's [`cProfile`](https://docs.python.org/3/library/profile.html) or [`line_profiler`](https://github.com/pyutils/line_profiler) can help identify these bottlenecks without manual instrumentation. - The pandas `query()` method is convenient and readable — and for single lookups, it's perfectly fine. The problem isn't the tool; it's using the wrong tool for the job. Context matters. - I kept the old implementation around to validate that the new one produced identical outputs for real-world data. Testing optimizations is just as important as making them. A faster wrong answer isn't a win. - This pattern isn't unique to pandas. The same principle applies to [database indexes](https://en.wikipedia.org/wiki/Database_index), [hash maps](https://en.wikipedia.org/wiki/Hash_table), and any repeated lookup scenario. When you're searching the same data multiple times, build an index first. - This optimization isn't always worth it — if you're only doing a handful of lookups, the index preparation overhead isn't justified. Measure first, then decide. - At the end of the day, performance optimization is about trade-offs: memory vs. speed, complexity vs. maintainability, development time vs. execution time. There's no universal answer — just the right answer for your situation. This aligns with the [value quadrant framework](/posts/valueinsoftware/) I've written about before. - Because sometimes, the biggest wins are hiding in the code you wrote six months ago and never looked at again.