L2 · 12 · AI for Data Engineering (Spark/Scala)

Data work is full of tasks AI does well: writing SQL and Spark code, explaining query plans, generating data-quality checks, documenting tables. It is also full of ways to be confidently wrong, such as a missing filter or a double-counting join. This page shows how to use AI for data engineering with validation built in.

Context is everything: schema + semantics

The model can't see your warehouse. Give it the schema, the meaning of columns, and the business definitions. Most wrong SQL comes from ambiguous definitions, not syntax.

Dialect: Spark SQL 3.5. Tables:
  orders(order_id STRING, customer_id STRING, store_id STRING, status STRING, created_at TIMESTAMP,
         ready_at TIMESTAMP, picked_up_at TIMESTAMP, total_cents BIGINT)
  stores(store_id STRING, name STRING, region STRING)
Semantics:
  - status in ('PLACED','READY_FOR_PICKUP','PICKED_UP','CANCELLED','REFUNDED')
  - A "no-show" is an order that reached READY_FOR_PICKUP and was CANCELLED without picked_up_at.
  - Times are UTC; report by store local date is NOT required.
  - Exclude test stores: store_id starting with 'TEST-'.
Question: no-show rate per region for the last 30 days, highest first.
Return: the SQL, then a list of assumptions you made.

Validating generated SQL

//> using scala 3

/** A lightweight guard, NOT a SQL parser. Always also use a read-only database user. */
object SqlGuard:
  private val Forbidden = Set("INSERT", "UPDATE", "DELETE", "MERGE", "DROP", "ALTER", "TRUNCATE",
                              "CREATE", "GRANT", "REVOKE", "CALL", "COPY", "EXEC")
  private val TableRef  = """(?i)\b(?:FROM|JOIN)\s+([A-Za-z_][\w.]*)""".r
  private val CteName   = """(?i)\b([A-Za-z_]\w*)\s+AS\s*\(""".r

  def check(sql: String, allowedTables: Set[String], maxRows: Int = 1000): Either[List[String], String] =
    val stripped = sql.replaceAll("(?s)/\\*.*?\\*/", " ").replaceAll("--[^\\n]*", " ").trim.stripSuffix(";").trim
    val upper    = stripped.toUpperCase
    val words    = upper.split("\\W+").toSet
    val ctes     = CteName.findAllMatchIn(stripped).map(_.group(1).toLowerCase).toSet
    val tables   = TableRef.findAllMatchIn(stripped).map(_.group(1).toLowerCase).toSet -- ctes
    val errors =
      Option.when(!(upper.startsWith("SELECT") || upper.startsWith("WITH")))("only SELECT/WITH queries are allowed").toList ++
      Option.when(stripped.contains(";"))("multiple statements are not allowed").toList ++
      (words intersect Forbidden).toList.sorted.map(k => s"forbidden keyword: $k") ++
      (tables -- allowedTables.map(_.toLowerCase)).toList.sorted.map(t => s"unknown table: $t")
    if errors.nonEmpty then Left(errors)
    else if """(?i)\bLIMIT\s+\d+\s*$""".r.findFirstIn(stripped).isDefined then Right(stripped)
    else Right(s"$stripped\nLIMIT $maxRows")

@main def sqlGuardDemo(): Unit =
  val allowed = Set("orders", "stores")
  List(
    """SELECT s.region, COUNT(*) AS n FROM orders o JOIN stores s ON o.store_id = s.store_id GROUP BY s.region""",
    """WITH recent AS (SELECT * FROM orders WHERE created_at > current_date - 30) SELECT COUNT(*) FROM recent LIMIT 1""",
    """DELETE FROM orders WHERE status = 'CANCELLED'""",
    """SELECT * FROM customers_pii; DROP TABLE orders"""
  ).foreach { q =>
    println(s"\n$q\n  -> ${SqlGuard.check(q, allowed).fold(e => "REJECTED: " + e.mkString(", "), ok => "OK:\n" + ok)}")
  }

Evaluate text-to-SQL with golden questions

Compare results, not SQL text: two different queries can both be right. This harness uses an in-memory SQLite database so it runs anywhere. In practice, run it against a small, fixed snapshot of your warehouse.

import sqlite3

SCHEMA = """
CREATE TABLE stores(store_id TEXT PRIMARY KEY, name TEXT, region TEXT);
CREATE TABLE orders(order_id TEXT PRIMARY KEY, store_id TEXT, status TEXT, total_cents INTEGER);
INSERT INTO stores VALUES ('RIV','Riverside','North'),('CEN','Central','North'),('AIR','Airport','South'),('TEST-1','Test','North');
INSERT INTO orders VALUES
 ('o1','RIV','PICKED_UP',12000),('o2','RIV','CANCELLED',5000),('o3','CEN','PICKED_UP',8000),
 ('o4','AIR','PICKED_UP',30000),('o5','AIR','REFUNDED',30000),('o6','TEST-1','PICKED_UP',999999);
"""

GOLDEN = [  # (question, reference SQL written and reviewed by a human)
    ("Revenue from picked-up orders per region, excluding test stores",
     "SELECT s.region, SUM(o.total_cents) FROM orders o JOIN stores s USING(store_id) "
     "WHERE o.status='PICKED_UP' AND s.store_id NOT LIKE 'TEST-%' GROUP BY s.region"),
    ("Number of cancelled orders",
     "SELECT COUNT(*) FROM orders WHERE status='CANCELLED'"),
]

# What a model produced (paste real model output here). The first one forgot to exclude test stores.
CANDIDATES = [
    "SELECT s.region, SUM(o.total_cents) FROM orders o JOIN stores s ON o.store_id=s.store_id "
    "WHERE o.status='PICKED_UP' GROUP BY s.region",
    "SELECT COUNT(order_id) FROM orders WHERE status = 'CANCELLED'",
]


def rows(db: sqlite3.Connection, sql: str) -> list[tuple]:
    return sorted(db.execute(sql).fetchall())


def main() -> None:
    db = sqlite3.connect(":memory:")
    db.executescript(SCHEMA)
    correct = 0
    for (question, ref), cand in zip(GOLDEN, CANDIDATES):
        try:
            ok = rows(db, cand) == rows(db, ref)
        except sqlite3.Error as e:
            ok, question = False, f"{question} (error: {e})"
        correct += ok
        print(f"{'PASS' if ok else 'FAIL'}  {question}")
        if not ok:
            print(f"      expected {rows(db, ref)}\n      got      {rows(db, cand)}")
    print(f"\nexecution accuracy: {correct}/{len(GOLDEN)}")


if __name__ == "__main__":
    main()

Spark with AI

Writing Spark code

• State Spark version and language. Spark 3.5 and 4.x publish Scala 2.13 artefacts. Scala 3 projects can use them through the 2.13 compatibility mode with the DataFrame API. Typed Dataset[T] encoders rely on Scala 2 reflection, so check community libraries if you need them in Scala 3.
• Ask for built-in functions only (org.apache.spark.sql.functions) and no UDFs unless necessary. Built-ins are optimised and visible to the planner.
• Ask for a small local test with a tiny DataFrame for every transformation.

import org.apache.spark.sql.{DataFrame, SparkSession}
import org.apache.spark.sql.functions.*

object NoShowRate:
  /** No-show = reached READY_FOR_PICKUP, then CANCELLED without pickup. Test stores excluded. */
  def compute(orders: DataFrame, stores: DataFrame): DataFrame =
    orders
      .filter(col("ready_at").isNotNull && !col("store_id").startsWith("TEST-"))
      .withColumn("no_show", (col("status") === "CANCELLED" && col("picked_up_at").isNull).cast("int"))
      .join(broadcast(stores.select("store_id", "region")), Seq("store_id"))
      .groupBy("region")
      .agg(count(lit(1)).as("ready_orders"), sum("no_show").as("no_shows"))
      .withColumn("no_show_rate", round(col("no_shows") / col("ready_orders"), 4))
      .orderBy(desc("no_show_rate"))

Let AI read the plan

Physical plans are long and cryptic, and models are good at reading them. Paste the output of df.explain("formatted") (plus data sizes) and ask:

Here is a Spark physical plan and approximate table sizes (orders 2B rows, stores 900 rows).
1) Explain in plain words what Spark does, step by step.
2) Identify expensive operations (shuffles, sort-merge joins, skew, full scans, missing pruning).
3) Suggest concrete changes (broadcast hints, partition filters, column pruning, AQE settings), with expected impact.

Data-quality checks

Give the AI a schema and a sample (synthetic or anonymised), and ask it to propose checks. Review them, because the AI doesn't know which anomalies are normal for your business.

For the orders table (schema + 20 sample rows below), propose data-quality checks grouped as:
completeness (nulls), validity (allowed values, ranges), uniqueness (keys), consistency
(picked_up_at >= ready_at; PICKED_UP implies picked_up_at not null), freshness (max(created_at) within 2h).
Write each as a SQL query that returns violating rows, plus a severity (block / warn).

LLMs inside pipelines

Sometimes the pipeline itself calls a model, for example to classify support tickets or extract attributes from product descriptions.

• Small model + structured output (page 02) for classification and extraction at volume.
• Batch and cache: use provider batch APIs or local models, and cache by input hash so re-runs don't re-pay.
• Idempotent and versioned: store model + prompt version with each result, so you can re-process when they change.
• Quality gates: sample outputs for human review. Track the label distribution and alert on drift.
• Cost estimate first: rows × average tokens × price. A 10-million-row backfill can be expensive.
• Privacy: minimise and redact personal data before sending it to any model.