Implement in-memory DB querying
Company: OpenAI
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: medium
Interview Round: Technical Screen
##### Question
Implement an in-memory database that supports: 1. Querying the whole table and returning only selected columns (projection). 2. Adding WHERE clause filtering with simple conditions like (column, operator, value). 3. Adding ORDER BY on one or more columns with ascending/descending control. 4. Explaining how you would design and build an index to accelerate such queries (no code required). Example public API: db = DB() db.insert("users", {"id": "1", "name": "Ada", "birthday": "1815-12-10"}) … db.query("users", ["id"], conditions=[("name", "=", "Charles")], order_by=(["birthday"], False)) # returns sorted projection
Overview: This question evaluates a candidate's competency in data structures, algorithms, and systems-level thinking for implementing query processing in an in-memory database, including projection, predicate filtering, ordering, and index design.
Part 1: Project Selected Columns from an In-Memory Table
Implement **relational projection** over an in-memory table.
You are given a table as a **list of rows**, where each row is a dictionary mapping string column names to values. Given a list of column names to keep, return a **new table** that contains only those columns for every row.
## Function
```python
def solution(rows, selected_columns):
...
```
- **`rows`** — a list of dictionaries. Each dictionary represents one row, with string keys (column names).
- **`selected_columns`** — a list of strings naming the columns to project, in the order they should appear in the output.
## What to return
A **new list of dictionaries** — one output row per input row — where each output row contains **exactly** the columns named in `selected_columns`:
- Keep the rows in their **original order**.
- In each output row, include **only** the keys in `selected_columns`, in that **same order**. Drop any other columns present in the source row.
- If a requested column is **missing** from a source row, still include that key in the output with value **`None`**.
Do **not** modify the input `rows` (build fresh dictionaries; the source rows are read-only).
## Examples
- `rows = [{"id": "1", "name": "Ada", "birthday": "1815-12-10"}, {"id": "2", "name": "Charles", "birthday": "1791-12-26"}]`, `selected_columns = ["id", "name"]`
→ `[{"id": "1", "name": "Ada"}, {"id": "2", "name": "Charles"}]` (the `birthday` column is dropped).
- `rows = [{"id": "1"}, {"name": "Ada"}]`, `selected_columns = ["id", "name"]`
→ `[{"id": "1", "name": None}, {"id": None, "name": "Ada"}]` (missing columns become `None`).
## Edge cases
- If `rows` is empty, return `[]`.
- If `selected_columns` is empty, return one **empty dictionary** `{}` for each input row (e.g. two rows → `[{}, {}]`).
## Constraints
- `0 <= len(rows) <= 10000`
- `0 <= len(selected_columns) <= 100`
- Each row is a dictionary with string keys.
- Do not modify the input rows in place.
Constraints
- 0 <= len(rows) <= 10000
- 0 <= len(selected_columns) <= 100
- Each row is a dictionary with string keys
- Do not modify the input rows in place
Examples
Input: ([{'id': '1', 'name': 'Ada', 'birthday': '1815-12-10'}, {'id': '2', 'name': 'Charles', 'birthday': '1791-12-26'}], ['id', 'name'])
Expected Output: [{'id': '1', 'name': 'Ada'}, {'id': '2', 'name': 'Charles'}]
Explanation: Keep only the id and name columns from each row.
Input: ([{'id': '1', 'name': 'Ada'}], ['name'])
Expected Output: [{'name': 'Ada'}]
Explanation: A single-row table should still project correctly.
Hints
- Build a fresh output row for each input row instead of deleting keys from the original.
- A dictionary get lookup is useful when a selected column may be missing.
Part 2: Filter Rows with Simple WHERE Conditions
Filter the rows of a table that satisfy **every** condition in a set of simple SQL-style `WHERE` clauses.
## What to implement
Implement `solution(rows, conditions)`.
- **`rows`** — a list of rows, where each row is a dictionary mapping a column name (string) to its value.
- **`conditions`** — a list of conditions, where each condition is a tuple `(column, operator, value)`.
Return the list of rows that match **all** of the given conditions (the conditions are combined with logical **AND**), in their **original order**.
## Operators
`operator` is one of the following comparison operators, each applied as `row[column] operator value`:
| Operator | Meaning |
|----------|-----------------------|
| `=` | equal to |
| `!=` | not equal to |
| `<` | less than |
| `<=` | less than or equal |
| `>` | greater than |
| `>=` | greater than or equal |
## Matching rules
- A row matches only if it satisfies **every** condition. If any condition fails, the row is excluded.
- **Missing column:** if a row does not contain the `column` referenced by a condition, that condition is treated as **not matching** (the row is excluded).
- **No conditions:** if `conditions` is empty, every row matches and all rows are returned (in original order).
- Within a single condition, `row[column]` and `value` are comparable (e.g. numbers with numbers, strings with strings).
## Examples
**Example 1**
```
rows = [{'id': 1, 'age': 36}, {'id': 2, 'age': 28}, {'id': 3, 'age': 36}]
conditions = [('age', '=', 36)]
returns [{'id': 1, 'age': 36}, {'id': 3, 'age': 36}]
```
**Example 2** — multiple conditions are ANDed together:
```
rows = [
{'id': 1, 'age': 36, 'name': 'Ada'},
{'id': 2, 'age': 28, 'name': 'Bob'},
{'id': 3, 'age': 40, 'name': 'Ada'},
]
conditions = [('name', '=', 'Ada'), ('age', '>', 36)]
returns [{'id': 3, 'age': 40, 'name': 'Ada'}]
```
**Example 3** — a row missing the referenced column does not match:
```
rows = [{'id': 1}, {'id': 2, 'name': 'Ada'}]
conditions = [('name', '=', 'Ada')]
returns [{'id': 2, 'name': 'Ada'}]
```
## Constraints
- `0 <= len(rows) <= 10000`
- `0 <= len(conditions) <= 20`
- Each operator is one of `=`, `!=`, `<`, `<=`, `>`, `>=`.
- Values compared within a condition are mutually comparable.
Constraints
- 0 <= len(rows) <= 10000
- 0 <= len(conditions) <= 20
- Each operator is one of '=', '!=', '<', '<=', '>', '>='
- Values compared within a condition are mutually comparable
Examples
Input: ([{'id': 1, 'age': 36}, {'id': 2, 'age': 28}, {'id': 3, 'age': 36}], [('age', '=', 36)])
Expected Output: [{'id': 1, 'age': 36}, {'id': 3, 'age': 36}]
Explanation: Only rows with age equal to 36 should remain.
Input: ([{'id': 1, 'age': 36, 'name': 'Ada'}, {'id': 2, 'age': 28, 'name': 'Bob'}, {'id': 3, 'age': 40, 'name': 'Ada'}], [('name', '=', 'Ada'), ('age', '>', 36)])
Expected Output: [{'id': 3, 'age': 40, 'name': 'Ada'}]
Explanation: A row must satisfy both conditions.
Approach
This is a straightforward predicate-filter over rows, mirroring a SQL WHERE clause where every condition is ANDed together.
Approach. The helper matches(row) decides whether a single row satisfies all conditions. It loops over each (column, operator, value) triple:
- If the row is missing the referenced column (column not in row), it returns False immediately — exactly the "treat missing column as not matching" rule.
- Otherwise it reads current = row[column] and dispatches on the operator string, comparing current against value with the matching Python comparison (==, !=, <, <=, >, >=).
- The instant any condition fails (not ok), it short-circuits with False. Only if every condition passes does it return True.
An unknown operator raises ValueError, guarding against bad input even though the constraints promise a valid operator.
Building the result. The final line is a comprehension: [dict(row) for row in rows if matches(row)]. It walks rows in their original order (so order is preserved automatically) and keeps those that match. Each kept row is wrapped in dict(row) to return a shallow copy, so the caller can't mutate the inputs through the result.
Why it's correct. AND-ing conditions = "row matches iff it satisfies each one," which the all-conditions-must-pass loop encodes directly. Short-circuiting never changes the boolean outcome, only the work done. Comparisons rely on the constraint that values within a condition are mutually comparable. Empty conditions makes matches trivially return True, so every row passes — the expected pass-through behavior.
Time complexity: O(n * m), where n = number of rows and m = number of conditions. Each row is checked against up to m conditions (short-circuiting only helps in practice). Copying matched rows adds O(total kept fields), bounded by O(n * c) for c columns per row.
Space complexity: O(r * c) for the output, where r = matching rows and c = fields per row, since each matched row is shallow-copied via dict(row). The matching logic itself uses O(1) auxiliary space.
Hints
- Write a helper that checks whether one row satisfies all conditions.
- An empty list of conditions should match every row.
Part 3: Sort Rows with ORDER BY on Multiple Columns
Implement a multi-column `ORDER BY`: sort a table of rows by one or more columns, where each column has its own ascending/descending direction.
## Function
```python
def solution(rows, order_columns, ascending_flags):
...
```
## Input
- **`rows`** — the table, as a list of row dictionaries. Each row maps **column name → value**.
- **`order_columns`** — a list of column names to sort by, given in **priority order** (the first column is the primary sort key, the second breaks ties on the first, and so on).
- **`ascending_flags`** — a list of booleans, one per column in `order_columns` and aligned by position:
- `True` → sort that column in **ascending** order
- `False` → sort that column in **descending** order
So `ascending_flags[i]` is the direction for `order_columns[i]`.
## Output
Return a **new** list of rows sorted according to the ORDER BY rules.
- Sort **lexicographically** by the columns in priority order: compare on the first column; for rows that tie, compare on the second column; and so on through the list. Each comparison respects that column's own ascending/descending flag.
- Do **not** modify the input — neither the `rows` list nor any of the original row dictionaries. The returned rows should be fresh copies.
## Rules and edge cases
- **No order columns** — if `order_columns` is empty, return the rows in their **original order** (as fresh copies).
- **Empty table** — if `rows` is empty, return an empty list.
## Constraints
- `0 <= len(rows) <= 10000`
- `0 <= len(order_columns) <= 10`
- `len(order_columns) == len(ascending_flags)`
- Every column in `order_columns` exists in every row and holds **comparable** values.
Constraints
- 0 <= len(rows) <= 10000
- 0 <= len(order_columns) <= 10
- len(order_columns) == len(ascending_flags)
- Every order column exists in every row and contains comparable values
Examples
Input: ([{'id': 1, 'name': 'Charles'}, {'id': 2, 'name': 'Ada'}, {'id': 3, 'name': 'Bob'}], ['name'], [True])
Expected Output: [{'id': 2, 'name': 'Ada'}, {'id': 3, 'name': 'Bob'}, {'id': 1, 'name': 'Charles'}]
Explanation: Sort by name ascending.
Input: ([{'id': 1, 'age': 30}, {'id': 2, 'age': 20}, {'id': 3, 'age': 40}], ['age'], [False])
Expected Output: [{'id': 3, 'age': 40}, {'id': 1, 'age': 30}, {'id': 2, 'age': 20}]
Explanation: Sort by age descending.
Approach
This solves a multi-column ORDER BY by exploiting stable sorting instead of building a composite comparator.
Setup. First it copies every row with dict(row) into result, so the input rows are never mutated and the output is a fresh list of fresh dicts (sorting in place on copies is safe).
Core idea — least-significant key first. It zips order_columns with ascending_flags, then iterates over those pairs in reversed priority order. So it sorts by the lowest-priority column first and the highest-priority column last:
Why it's correct. Python's list.sort() (Timsort) is stable: equal keys keep their current relative order. By sorting on the top-priority column last, ties on that column fall back to the order produced by the previous (next-lower-priority) sort, which in turn already encoded the column below it, and so on. The net effect is a true lexicographic, multi-key ordering. Per-column direction is handled by reverse=not ascending.
Test 3 demonstrates this: with ['age','name','id'] / [True,True,False], rows id=4 and id=2 (both age 30, name 'Ada') end up id=4 before id=2 because id is the last, descending sort and acts as the tiebreaker.
Edge cases. Empty rows returns []; empty order_columns skips the loop entirely and returns the copied rows in original order. A guard also raises ValueError if the two list lengths disagree. The constraint that every order column exists and holds comparable values is what keeps row[column] and the comparisons safe.
Time complexity: O(k * n log n), where n is the number of rows and k is the number of ORDER BY columns — one stable O(n log n) sort pass per column.
Space complexity: O(n) for the copied output list of row dicts (sorting is in place; Timsort's own temporary buffer is also O(n)).
Hints
- Python's sort is stable, so sorting from the last key to the first key is a clean way to handle mixed directions.
- Make a copy of the rows before sorting if you do not want to mutate the input.
Part 4: Build a Simple Equality Index for Fast Lookups
Build a simple **equality (hash) index** over a table, then use it to answer equality lookups quickly.
In many in-memory databases, repeated equality filters on a single column can be sped up by precomputing a map from each column value to the row positions that hold it. Your task is to build such an index and use it to resolve a batch of lookup values.
### Implement
```python
def solution(rows, index_column, lookup_values):
...
```
### Input
- **`rows`** — a list of rows, where each row is a dictionary mapping **column name → value**. Rows may have different sets of columns.
- **`index_column`** — the name (string) of the column to index on.
- **`lookup_values`** — a list of values to look up against the indexed column.
### Output
Return a list **parallel to `lookup_values`**. For each lookup value, in the same order, return the list of **0-based positions** of the rows whose value in `index_column` equals that lookup value.
- Positions within each list must be in **ascending order** (i.e. in the order the rows appear in `rows`).
- If a lookup value matches no row, return an **empty list** `[]` for it.
- The result has exactly one entry per lookup value. If the same lookup value appears more than once in `lookup_values`, it produces a separate matching list each time.
### Rules
- A row at position `i` matches a value `v` when `row[index_column] == v`.
- **Ignore any row that does not contain `index_column`** — such rows are never included in any result and their positions are skipped (but positions of the remaining rows are still their original indices in `rows`).
- Values in the indexed column are **hashable**.
### Examples
- `rows = [{'id': '1', 'name': 'Ada'}, {'id': '2', 'name': 'Charles'}, {'id': '3', 'name': 'Ada'}]`, `index_column = 'name'`, `lookup_values = ['Ada', 'Charles', 'Eve']` → `[[0, 2], [1], []]`
- `rows = [{'id': 1}, {'id': 2, 'city': 'Paris'}, {'id': 3, 'city': 'Paris'}]`, `index_column = 'city'`, `lookup_values = ['Paris', 'London']` → `[[1, 2], []]` (the first row has no `city`, so it is ignored, but rows 1 and 2 keep their original positions)
- `rows = []`, `index_column = 'name'`, `lookup_values = ['Ada']` → `[[]]`
### Constraints
- `0 <= len(rows) <= 100000`
- `0 <= len(lookup_values) <= 100000`
- Indexed values are hashable.
- Ignore rows missing the indexed column.
Constraints
- 0 <= len(rows) <= 100000
- 0 <= len(lookup_values) <= 100000
- Indexed values are hashable
- Ignore rows missing the indexed column
Examples
Input: ([{'id': '1', 'name': 'Ada'}, {'id': '2', 'name': 'Charles'}, {'id': '3', 'name': 'Ada'}], 'name', ['Ada', 'Charles', 'Eve'])
Expected Output: [[0, 2], [1], []]
Explanation: The index groups all row positions for each name.
Input: ([{'id': 1, 'age': 36}, {'id': 2, 'age': 28}, {'id': 3, 'age': 36}, {'id': 4, 'age': 40}], 'age', [40, 36])
Expected Output: [[3], [0, 2]]
Explanation: Multiple rows can share the same indexed value.
Approach
The solution mimics a hash index in two phases: build, then probe.
Build phase. Iterate over rows once with enumerate so each row keeps its original position i. For every row that actually contains index_column, read its value key = row[index_column] and append i to a list keyed by that value:
setdefault creates an empty list the first time a value is seen and reuses it afterward, so index ends up mapping each distinct column value to all row positions that hold it, in ascending order (positions are appended in iteration order). Rows missing the column are simply skipped, satisfying the "ignore rows without the indexed column" rule.
Probe phase. For each value in lookup_values, fetch its position list with index.get(value, []), defaulting to an empty list when the value was never indexed:
The wrapping list(...) makes a fresh copy, so the returned lists are independent of the internal index (callers can't mutate the index by accident, and the same lookup value appearing twice yields two separate lists).
Why it's correct. Equality lookups on a column are exactly "give me every row whose value equals X," which a value→positions map answers directly. Because values are guaranteed hashable, dictionary insert/lookup are O(1) average, turning what would be an O(n) scan per query into a single up-front pass plus constant-time probes. Output ordering of positions matches row order, matching the expected results (e.g. 'Ada' → [0, 2]).
Time complexity: O(n + q + S)
Space complexity: O(n)
Hints
- Build the mapping from value to row positions once, then answer each lookup in O(1) average time.
- If a value appears multiple times, store all matching positions, not just one.