Skip to content

Working with the data — best practices

Download notebook

Using this notebook in Databricks

Download the notebook using the button above, then import it into your Databricks workspace:

  1. In the Databricks UI, go to Workspace → Shared
  2. Click the ⋮ menu and choose Import
  3. Select File and upload the downloaded .py file
  4. The notebook will appear in the Shared space and be accessible to everyone in your workspace

Adding it to Shared means all workspace members can access and run it. If you want a personal copy to edit, clone it into your own folder first (File → Clone).

The data mart contains tables that can have over 100 million rows. This page covers how to work with that efficiently and avoid common pitfalls.

Key principle: filter early, collect late.

Let Spark and SQL do the heavy lifting on the full dataset. Only pull data into pandas (or CSV) once you have filtered it down to what you actually need.

1. Understand the size of your data first

Before doing anything, check how large the table is. This avoids accidentally trying to load 100M rows into memory.

-- Row count (fast — doesn't scan the data)
SELECT COUNT(*) AS total_rows FROM <mart_catalog_name>.default.diagnosis;

-- Size and column info
DESCRIBE DETAIL <mart_catalog_name>.default.diagnosis;

2. Filter in SQL or Spark, not in pandas

SQL and Spark process data in parallel across the cluster. Pandas runs on a single node and requires the entire dataset to fit in memory. For large tables, always filter before collecting.

Do this:

# Good — filter 100M rows down to what you need using Spark
df = (
    spark.table("<mart_catalog_name>.default.diagnosis")
    .filter("diag_beg_effective_dt_tm >= 2020")
    .filter("encntr_type_class_cd = 391")
    .select("salted_warehouse_identifier", "diag_beg_effective_dt_tm", "encounter_type_class", "icd10")
)

# Only now convert to pandas — should be a manageable number of rows
print(f"Filtered rows: {df.count():,}")
pdf = df.toPandas()

Not this:

# Bad — pulls 100M rows into driver memory, then filters
# This will be slow, expensive, and may crash with an OutOfMemory error

# pdf = spark.table("<mart_catalog_name>.<schema_name>.<table_name>").toPandas()
# pdf = pdf[pdf['diagnosis_year'] >= 2020]

3. Use SQL for aggregation

If you need summary statistics or group-level counts, do the aggregation in SQL. The result will be small enough to work with anywhere.

-- Aggregate 100M rows down to a summary — fast and returns a small result
-- SELECT
--   diagnosis_year,
--   stage,
--   COUNT(*) AS patient_count,
--   AVG(age_at_diagnosis) AS avg_age
-- FROM <mart_catalog_name>.<schema_name>.<table_name>
-- GROUP BY diagnosis_year, stage
-- ORDER BY diagnosis_year, stage;

The result of that query might be a few hundred rows — safe to pull into pandas for plotting or further analysis.

# Aggregate in Spark, then convert — small result, no memory issues
# summary = spark.sql("""
#     SELECT diagnosis_year, stage, COUNT(*) AS n
#     FROM <mart_catalog_name>.<schema_name>.<table_name>
#     GROUP BY diagnosis_year, stage
# """)
# pdf_summary = summary.toPandas()

4. Sample when exploring

If you want to get a feel for the data before committing to a full analysis, take a sample rather than loading the whole table.

# Random 1% sample of a large table
# df_sample = spark.table("<mart_catalog_name>.<schema_name>.<table_name>").sample(0.01)
# print(f"Sample size: {df_sample.count():,}")
# pdf_sample = df_sample.toPandas()

# Or just take the first N rows (faster, but not random)
# pdf_preview = spark.table("<mart_catalog_name>.<schema_name>.<table_name>").limit(10000).toPandas()

5. Save intermediate results

If you have spent time filtering and transforming a large table, save the result to your workspace catalog so you don't have to reprocess it every time you open the notebook.

# Filter once, save to workspace catalog
# cohort = (
#     spark.table("<mart_catalog_name>.<schema_name>.<table_name>")
#     .filter("diagnosis_year >= 2018")
#     .filter("stage IS NOT NULL")
# )
#
# cohort.write.mode("overwrite").saveAsTable(
#     "<workspace_catalog_name>.my_project.filtered_cohort"
# )

# Next time, read from your saved table — much faster
# cohort = spark.table("<workspace_catalog_name>.my_project.filtered_cohort")

6. When you need a CSV

Sometimes you need a CSV file — for a bash script, an external tool, or a specific workflow. That's fine, but follow these guidelines.

Files written to /tmp are temporary

Files written to /tmp are stored on the cluster's local ephemeral storage. They are deleted when your session ends or the cluster is terminated. If you need results to persist between sessions, save them as a table in your workspace catalog (see section 5) rather than as a CSV on /tmp.

Always filter first

Never export a full 100M-row table to CSV. Filter down to what you need, then export.

# Good — export a filtered subset
# filtered = (
#     spark.table("<mart_catalog_name>.<schema_name>.<table_name>")
#     .filter("diagnosis_year = 2023")
#     .filter("region = 'Thames Valley'")
#     .select("patient_id", "diagnosis_date", "stage")
# )
#
# print(f"Exporting {filtered.count():,} rows")
# filtered.toPandas().to_csv("/tmp/cohort_2023.csv", index=False)

Understand the costs

Exporting to CSV involves compute (Spark reads and serialises the data), memory (the full result must fit in the driver node), and transfer charges if the source storage is in a different region.

For very large exports that don't fit in driver memory, write to CSV using Spark directly — this avoids the pandas memory bottleneck:

# Large export — write CSV using Spark (no memory limit, but creates multiple files)
# (
#     spark.table("<mart_catalog_name>.<schema_name>.<table_name>")
#     .filter("diagnosis_year >= 2020")
#     .write
#     .mode("overwrite")
#     .option("header", "true")
#     .csv("/tmp/large_export")
# )
#
# # This creates a folder with multiple CSV files (one per partition)
# # To combine into a single file:
# # .coalesce(1).write.mode("overwrite").option("header", "true").csv("/tmp/single_export")
# # Note: coalesce(1) forces all data through one node — only use for moderate-sized results

7. Running existing scripts

If you have existing Python or bash scripts that worked with CSV files, you can still use them.

  1. Filter the data in SQL or Spark
  2. Export the filtered result to a temporary CSV
  3. Run your script against the CSV
# Steps 1 & 2: filter and export
# (
#     spark.table("<mart_catalog_name>.<schema_name>.<table_name>")
#     .filter("diagnosis_year >= 2020")
#     .select("patient_id", "diagnosis_date", "stage", "treatment")
#     .toPandas()
#     .to_csv("/tmp/filtered_data.csv", index=False)
# )

# Step 3: run your bash script against the exported CSV
# bash /Workspace/Shared/scripts/my_analysis.sh /tmp/filtered_data.csv

For Python scripts, you can paste them into a cell and change the data loading line from pd.read_csv(...) to spark.table(...).toPandas(). The rest of the script stays the same.

8. Choosing the right approach

What you need Approach Notes
Counts, averages, summaries SQL aggregation Fastest — result is tiny
Explore a large table .sample(0.01) or .limit(10000) Get a feel for the data without loading it all
Filter and analyse a subset Spark filter → .toPandas() Let Spark do the filtering
Reuse a filtered dataset Save to workspace catalog Avoids re-filtering every session
Run an existing pandas script Spark filter → .toPandas() → your script Change one line (the data load)
Run an existing bash script Spark filter → export CSV → %sh bash script.sh Filter first, export the subset
Plot or visualise SQL/Spark aggregate → pandas → matplotlib Aggregate the data before plotting
Full table export Avoid if possible; use Spark .write.csv() if essential Never use .toPandas() on 100M+ rows

Remember: the data mart is always there. You don't need to download everything upfront. Query what you need, when you need it.