Python for Data Engineers Cheat Sheet

Core Libraries Every Data Engineer Needs

Production-Safe Snowflake Connection Pattern

import os
import snowflake.connector
from contextlib import contextmanager

@contextmanager
def snowflake_conn():
    """Always close connection, even on error."""
    conn = snowflake.connector.connect(
        account=os.environ['SF_ACCOUNT'],
        user=os.environ['SF_USER'],
        password=os.environ['SF_PASSWORD'],
        warehouse=os.environ['SF_WAREHOUSE'],
        database=os.environ['SF_DATABASE'],
        schema=os.environ['SF_SCHEMA'],
        role=os.environ['SF_ROLE'],
        client_session_keep_alive=False,
    )
    try:
        yield conn
    finally:
        conn.close()

def run_query(sql: str, params: dict | None = None):
    with snowflake_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(sql, params or {})
            return cur.fetchall()

# Parameterized query (SQL-injection safe)
rows = run_query(
    "SELECT * FROM orders WHERE customer_id = %(cid)s AND order_date >= %(d)s",
    {"cid": 12345, "d": "2026-01-01"},
)

Efficient Pandas Patterns

import pandas as pd

# Read Parquet (10-100x faster than CSV)
df = pd.read_parquet('s3://bucket/orders.parquet')

# Specify dtypes to cut memory 50%+
df = pd.read_csv('data.csv', dtype={
    'order_id': 'int32',
    'amount': 'float32',
    'status': 'category',
})

# Vectorize — NEVER iterate row-by-row
# BAD:   for idx, row in df.iterrows(): ...
# GOOD:  df['total'] = df['qty'] * df['price']

# Use assign() for readable chained transforms
clean = (
    df
    .query('amount > 0')
    .assign(
        order_month=lambda d: d['order_date'].dt.to_period('M'),
        is_big=lambda d: d['amount'] > 1000,
    )
    .groupby('order_month', as_index=False)
    .agg(revenue=('amount', 'sum'), orders=('order_id', 'count'))
)

# Process giant files in chunks
for chunk in pd.read_csv('huge.csv', chunksize=100_000):
    processed = chunk.pipe(transform).pipe(validate)
    processed.to_parquet(f'out/{chunk.index[0]}.parquet')

Generators & Memory-Efficient Iteration

# Stream rows from Snowflake without loading all into memory
def stream_query(sql: str, batch_size: int = 10_000):
    with snowflake_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(sql)
            while True:
                rows = cur.fetchmany(batch_size)
                if not rows:
                    break
                yield from rows

# Consumer uses constant memory no matter the table size
for row in stream_query("SELECT * FROM huge_table"):
    process(row)

# Batch generator for bulk inserts
from itertools import islice
def batched(iterable, n):
    it = iter(iterable)
    while (batch := list(islice(it, n))):
        yield batch

# Bulk write 1000 rows at a time
for batch in batched(records, 1000):
    cursor.executemany(INSERT_SQL, batch)

Testing Data Pipelines

# test_transforms.py
import pandas as pd
import pytest
from my_pipeline import clean_orders

def test_clean_orders_removes_negative_amounts():
    raw = pd.DataFrame({
        'order_id': [1, 2, 3],
        'amount': [100.0, -5.0, 250.0],
    })
    result = clean_orders(raw)
    assert len(result) == 2
    assert (result['amount'] > 0).all()

def test_clean_orders_preserves_schema():
    raw = pd.DataFrame({'order_id': [1], 'amount': [100.0]})
    result = clean_orders(raw)
    assert set(result.columns) == {'order_id', 'amount', 'order_month'}

@pytest.mark.parametrize("amount,expected", [
    (0.01, True), (999.99, True), (1000.01, False), (-1, False),
])
def test_valid_amount(amount, expected):
    assert is_valid_amount(amount) is expected

# Run: pytest -v --cov=my_pipeline --cov-report=term-missing

Production Pipeline Best Practices

• ✓Never hardcode credentials — use environment variables, AWS Secrets Manager, or Snowflake key-pair auth
• ✓Always use parameterized queries (%(name)s or %s) — never f-string SQL (SQL injection risk)
• ✓Wrap connections in context managers (with statement) so they close even on exception
• ✓Use logging, not print() — configure root logger with JSON formatter for production
• ✓Validate input schemas with pydantic BaseModel — fail fast on bad data
• ✓Prefer Parquet over CSV for intermediate files: 10-100x faster I/O, preserves dtypes
• ✓Use Snowpark DataFrames when possible — pushes compute to Snowflake, avoids data egress
• ✓Pin dependencies in requirements.txt or pyproject.toml — unpinned versions break builds silently
• ✓Add retries with exponential backoff for network calls (tenacity library)
• ✓Write tests before productionizing — pytest + moto (mock AWS) + testcontainers (real Postgres)

Frequently Asked Questions

Related Cheat Sheets