AI is becoming part of almost every data engineering workflow.
We are asking AI to generate SQL, PySpark pipelines, data-quality rules, unit tests, technical specifications and even complete project structures. The results can be impressive, but there is another side that we do not discuss enough: cost and unnecessary complexity.
Recently, I asked an AI coding tool to generate validation scripts for a large data pipeline. It produced thousands of lines of code and consumed a significant number of tokens. But when I reviewed the output carefully, only a small portion of the required columns and business rules were actually covered.
The code volume looked impressive. The actual coverage was not.
That experience made me question a common assumption:
Does every data engineering task really need the most powerful LLM?
In most cases, the answer is no.
Not Every Task Requires AI
Consider some typical data engineering activities:
Generate null checks for 100 columns.
Create default-value checks for dimension keys.
Compare source and target row counts.
Generate a standard Delta table definition.
Validate naming conventions.
Create repetitive unit-test queries.
These are predictable and rule-based tasks. A Python or SQL generator can often perform them faster, more consistently and at almost zero AI cost.
Now consider a different set of problems:
Diagnose why YTD values reset incorrectly.
Explain why a join is unexpectedly multiplying records.
Interpret unclear transformation rules from an STTM.
Recommend a design for handling late-arriving dimensions.
Analyze whether a pipeline meets a complex business requirement.
These tasks require reasoning and context. This is where an LLM provides real value.
The problem is that many implementations send both types of work to the same expensive model.
Cost Control Alone Is Not Enough
Databricks recently introduced AI spend controls in Unity AI Gateway. Organizations can define budgets and monitor AI expenditure across users, workspaces and use cases. It also provides request-level cost information through Unity Catalog system tables.
This is important because token usage alone does not tell us whether an AI task delivered value.
A request may consume 100,000 tokens and still generate incomplete or unusable code. Another request may consume 5,000 tokens and identify a production issue that saves several days of investigation.
The real metric should not be:
How many tokens did we use?
It should be:
What useful and validated outcome did we receive for the cost?
Databricks has also introduced Smart Routing in Unity AI Gateway, which selects models based on task complexity. According to Databricks, its internal coding workloads achieved meaningful cost reduction by routing simpler tasks to less expensive models and escalating complex tasks when necessary.
This is a strong direction. But for data engineering, I believe routing should go one step further.
We should not only select the appropriate model. We should first determine whether an LLM is required at all.
A Data Engineering-Specific Routing Pattern
I call this pattern the AI Pipeline Engineering Governor.
The governor sits between the user and the execution engine. It analyzes the requested task and chooses one of three execution paths.
Path 1: Deterministic Engine
The task is handled using Python, SQL templates or metadata-driven rules.
Examples include:
DDL generation
Standard data-quality checks
Null and default-value validation
Naming-standard validation
Schema comparison
Reconciliation query generation
Repetitive test generation
This path is inexpensive, predictable and easy to validate.
Path 2: Low-Cost Model
A smaller model handles tasks requiring limited interpretation.
Examples include:
Explaining a generated SQL statement
Classifying a transformation rule
Summarizing a pipeline failure
Converting a clearly defined rule into SQL
Generating documentation from validated metadata
Path 3: Advanced Reasoning Model
A more capable model is used only when the problem requires deeper reasoning.
Examples include:
Investigating reconciliation mismatches
Diagnosing record-count explosions
Understanding complex temporal logic
Resolving conflicting business rules
Recommending architectural trade-offs
Analyzing unfamiliar pipelines with incomplete documentation
The router should default to the deterministic or lower-cost path. It should escalate only when the task complexity, ambiguity or validation result justifies the additional expense.
Proposed Databricks Architecture
The solution can be implemented using native Databricks capabilities.
1. Databricks App
A simple Databricks App provides the user interface.
The engineer submits:
Task description
Source and target schemas
STTM or transformation rules
Existing SQL or PySpark code
Data-quality expectations
The interface displays the selected execution path before running the request.
2. Task Classifier
The classifier evaluates factors such as:
The first version does not need another expensive AI model for classification. A rule-based scoring mechanism is enough.
For example:
if task_type in ["null_check", "schema_compare", "ddl_generation"]:
route = "DETERMINISTIC"
elif ambiguity_score <= 2 and transformation_count < 10:
route = "LOW_COST_MODEL"
else:
route = "ADVANCED_MODEL"Over time, the router can learn from execution results and human feedback.
3. Execution Layer
The execution layer invokes the selected option:
The execution decision and its reason should be stored with every request.
4. Validation Layer
Generating code is not the final step.
Every output should pass relevant validations before it is presented for approval.
The validation layer can check:
If a lower-cost route fails validation, the task can be escalated automatically with the failure details included as context.
This is better than immediately sending the original request to the most expensive model.
5. MLflow Tracing and Evaluation
MLflow 3 can capture the complete execution trace, including inputs, outputs, intermediate steps and evaluation results.
For every task, I would record:
MLflow scorers can evaluate whether the generated output is complete, grounded in the supplied metadata and aligned with engineering standards.
6. Unity Catalog Governance
The generated artifacts should remain traceable to their original inputs.
Unity Catalog can help govern:
Sensitive STTM documents, schemas and production code should not become uncontrolled prompt content.
Governance must apply to both the data and the AI interaction.
7. Human Approval
The governor should never promote AI-generated code directly into production.
The final workflow should be:
Request
ā
Task Classification
ā
Deterministic Engine or AI Model
ā
Automated Validation
ā
Cost and Quality Evaluation
ā
Human Approval
ā
CI/CD Promotion
AI can accelerate engineering work, but the engineer should remain accountable for the production outcome.
Measuring Whether the Router Is Working
A cost-aware router should not be evaluated only by total savings.
The following metrics provide a more complete picture:
Cost per accepted artifact
Percentage of tasks completed without an LLM
Percentage completed by a low-cost model
Escalation rate
First-pass validation rate
Transformation-rule coverage
Unit-test coverage
Human rejection rate
Average delivery time
Production defect rate
Suppose a cheaper model reduces the cost by 60% but doubles the engineerās review time. That may not be a real saving.
Similarly, an advanced model may be more expensive but justified when it correctly diagnoses a complex reconciliation issue.
The objective is not simply to minimize AI cost.
The objective is to maximize validated engineering output per dollar.
How This Differs from Generic Model Routing
Generic model routing asks:
Which model should handle this task?
A data engineering governor asks several questions:
Can this task be completed deterministically?
If AI is required, what is the least expensive capable model?
Is the output aligned with the supplied metadata and business rules?
Did the generated artifact pass technical validation?
Is escalation justified?
Has a human approved the final result?
This distinction matters because many data engineering tasks are metadata-driven and repeatable. Using an LLM for every step may introduce inconsistency where deterministic logic would be more reliable.
Start Small
The first version does not need a complex multi-agent architecture.
A useful MVP could support three tasks:
Generate data-quality checks from table metadata.
Translate transformation rules into SQL or PySpark.
Diagnose source-to-target reconciliation mismatches.
The first task can primarily use deterministic generation. The second can use a lower-cost model with validation. The third can use an advanced reasoning model.
This is enough to demonstrate routing, tracing, validation, cost measurement and human approval end to end.
Final Thought
LLMs are powerful, but using the most capable model for every task is neither an architecture nor a governance strategy.
Some work requires reasoning. Some work requires deterministic automation. The important part is knowing the difference.
Databricks now provides many of the building blocks through Unity AI Gateway, Unity Catalog and MLflow 3. The next step is to apply them to measurable engineering workflows where cost, quality and accountability can be evaluated together.
The future of agentic data engineering is not about generating the maximum amount of code.
It is about producing the right engineering outcome, using the right execution path, at a cost we can explain.
References
#Databricks #DataEngineering #AgenticAI #GenerativeAI #MLflow #UnityCatalog #UnityAIGateway #AIGovernance #FinOps #LLMOps #PySpark #Lakehouse #DataQuality #AIAgents #CostOptimization
Amit Kumar Singh
Lead Data Engineer | AI-Assisted Data Engineering
