Python
Intermediate 13 min read Python

Data Structures for AI Work

The right container for each pipeline stage

AI pipelines move data through lists, dicts, generators and sets every stage. Choosing the right structure decides memory use, lookup speed and how cleanly your prompt/retrieval code composes.

listsdictssetsgeneratorsdequeCounterdataclasses
Prerequisites: Python for AI Engineering
Before we start

What you will learn

  • Pick containers by access pattern: index, key, membership or stream
  • Use generators to avoid materializing large document sets in memory
  • Compose dicts, dataclasses and counters into clean retrieval code
Introduction

What is this?

Python's core data structures — lists, dicts, sets, tuples, collections.deque and Counter — each answer a specific access pattern. Lists are ordered and indexable, dicts map keys to values with O(1) lookup, sets test membership fast, and generators stream items without loading them all. Choosing deliberately is what keeps document-processing and retrieval code fast and readable.

Motivation

Why do you need it?

  • A documents dataset with 200k chunks is a list size we can stream, but a dict of chunk_id → content is what retrieval needs.
  • Membership tests with `in` on a list are O(n); on a set they are O(1) — vital for dedup and filtering.
  • Generators cut peak memory when processing large corpora: you never hold the whole thing at once.
  • `Counter` turns 'top frequent tokens' from ten lines of code into one.
Think about it

Mental model

Imagine

Different drawers for different things

Think of

Data structures in Python

You store keys in a magnetic tray (dict by key), spare batteries in a box you check one by one (list), and a 'have I seen this?' checklist you mark with one glance (set). Choosing the wrong drawer makes every use slower.

Mechanics

How it works

  • A dict hashes keys into buckets, giving O(1) membership and lookup with a small memory overhead.
  • A set is a dict with only keys — optimal for deduplication and intersection logic.
  • A list keeps order and random access by index; appends are amortized O(1), inserts in the middle are O(n).
  • A generator yields one element at a time, resuming work only as consumers ask for more.
  • `collections.deque` gives O(1) append/pop at both ends — perfect for sliding-window token contexts.

Matching structures to pipeline stages.

Ingest docs
Chunk IDs
Lookup
Frequency
Retrieval output
Hands-on

Examples

Dedup + lookup with sets and dicts

Dedup + lookup with sets and dicts
from collections import Counter
 
chunks = ["alpha", "beta", "alpha", "gamma", "beta"]
 
unique = set(chunks)          # O(1) membership, dedup
by_id = {i: c for i, c in enumerate(chunks)}
counts = Counter(chunks)      # frequency map
 
print(sorted(unique))
print("alpha" in by_id.values())
print(counts.most_common(2))

One pass builds a dedup set, an id→content map for retrieval and a frequency counter — three access patterns served by three structures instead of three loops.

Real world

In production

  • A chunker streams a 5GB corpus as a generator, dedups chunk hashes against a set, and only materializes the small set of chunks that will be embedded.
  • A retrieval hot-path stores chunk_id → embedding_id and text in dicts, so a rerank pipeline looks up by key instead of scanning.
  • A token-budget governor uses `deque(maxlen=N)` to hold the sliding context window without reallocating.
Guidelines

Best practices

  • Default to tuples for fixed records, dataclasses for records with behavior, dicts for maps.
  • Stream large inputs; materialize only what a consumer needs.
  • Use sets for any uniqueness guarantee you mutate over time.
  • Prefer `Counter` and `defaultdict` over handwritten counting and grouping loops.
Watch out

Common mistakes

  • Testing `item in list` inside a hot loop, turning O(n) scans into a pipeline bottleneck.
  • Building a list of a million documents where a generator would do.
  • Using a dict as a namespace and discovering typos surface as KeyError at 3am.
  • Mutating a list while iterating it and watching indices shift.
Root cause analysis

Common issues

Error
MemoryError while processing a large corpus
Root cause

A generator chain collapsed into a full materialized list

Fix

Stream with generators end-to-end and hold only the aggregates you need.

Error
Unhashable type: 'list' when building a set
Root cause

Putting a mutable list where a hashable key (tuple/frozenset) is required

Fix

Convert to tuples or frozen subsets before insertion.

Error
Duplicate chunks get embedded twice
Root cause

Dedup used a list membership scan with a subtle identity mismatch

Fix

Hash the chunk content into a set and normalize (strip, casefold) before hashing.

When things break

Debugging tips

  • If 'unexpected duplicates' appear, the bug is usually mutating a set keys or a non-hashable member.
  • Memory spikes in a processing pipeline usually mean somewhere you dropped the generator to a list.
  • Slow lookups in hot paths — convert the list you scan to a set or dict and re-measure.
  • For sliding-window corruption, check for modifying a deque's referents rather than the deque itself.
Interviews

Interview checkpoint

Practice

Hands-on lab

Build a Streaming Chunk Deduplicator
25 min · intermediate
In progress
Objectives
  • Stream a synthetic corpus with a generator
  • Deduplicate chunks using a set of fingerprints
  • Build a chunk_id → content dict without materializing everything

Build a Streaming Chunk Deduplicator

25 min · intermediate

0/4 steps
0%
sandbox · Linux sandbox
$
Try these commands
Test yourself

Quick quiz

0/2
1

Which structure gives O(1) membership testing?

2

Generators can reduce peak memory in document-processing pipelines.

Go deeper

Official documentation

Python data structures
docs
collections module
docs