How can I use a natural-language assistant to explore and prepare data for machine learning model training?

Summary

• AI/BI Genie lets you ask natural language questions to explore datasets and generate SQL queries without writing code
• Feature Store in Unity Catalog automates feature engineering, transformation, and training dataset preparation
• Genie-generated insights can directly feed into feature tables that are tracked and governed for reproducible ML workflows

How can I use a natural-language assistant to explore and prepare data for machine learning model training?

Yes. A natural-language AI assistant on the Databricks lakehouse accelerates data exploration and feature engineering for ML workflows. You ask questions about your data in plain English; the assistant generates SQL, runs queries, and presents results and visualizations. Once you understand the data patterns that matter for your model, you can register those transformations as governed feature tables in Unity Catalog, turning exploratory insights into reusable, audited training assets.

What is a natural-language data assistant?

A natural-language data assistant interprets your questions about organizational data and translates them into executable queries. On Databricks, AI/BI Genie serves this role, responding to conversational questions with SQL-generated answers, tables, and charts. Because Genie is grounded in Unity Catalog metadata, it understands your organization's data structure and automatically applies access controls, so you only see data you're authorized to access.

Key capabilities to evaluate

Capability	Why it matters
Natural language to SQL generation	Frees data scientists from writing boilerplate SQL, letting you iterate quickly through hypotheses without syntax barriers
Automatic schema discovery	The assistant reads table structures, column types, and business context from catalog metadata, making recommendations relevant to your domain
Governed data access	Results respect Unity Catalog permissions and masking rules, so you don't accidentally train on unauthorized data
Feature engineering integration	Insights can be materialized into feature tables with automatic lineage tracking, linking exploratory questions directly to training datasets
Visualization generation	Charts and aggregations are produced automatically, reducing time spent on EDA formatting

How Databricks approaches this

AI/BI Genie is Databricks' natural-language data exploration platform. You ask questions in plain English within a chat interface; Genie generates the appropriate SQL, executes it against your data, and returns results with tables and visualizations. Since Genie is built on top of Unity Catalog, every query is routed through your governance policies. Column-level masking, row filtering, and privilege checks all apply automatically.
Once you've explored your data with Genie and identified features or transformations worth reusing, you register them as Feature Tables in Unity Catalog. Feature Tables are managed datasets tracked by Databricks' Feature Store, which is native to Unity Catalog. When you train a model, you use the FeatureLookup API to pull training data from these tables, ensuring reproducibility and auditability. The lineage between raw data, features, and trained models is recorded in your catalog, simplifying governance and compliance.
For complex feature pipelines, you can orchestrate transformations using Lakeflow Jobs, which schedule feature computation and keep training datasets fresh. MLflow tracks each training run, storing parameters, metrics, and the exact feature set used, so you can debug poor model performance by replaying the exact data and computation that produced it.
Genie Code, an agentic coding assistant within Databricks, can also accelerate feature engineering by suggesting transformations and helping you author the SQL or Python that computes features at scale.

Use cases

• Quick data discovery for a new ML project: A data scientist explores a financial dataset with Genie to identify which transaction patterns and customer segments are most predictive of churn. They ask questions like "Which customer segments have the highest transaction frequency?" and "What are the top reasons transactions fail?" Within minutes, they've narrowed down the feature space without writing SQL.
• Feature engineering audit trail: An insurance underwriter uses Genie to explore claim data and Databricks creates feature tables from the results. When a model is audited for regulatory compliance, lineage records show exactly which raw tables fed into the features, which fields were transformed and how, and who approved each step.
• Cross-team collaboration: A machine learning team member shares Genie queries with a business analyst to validate that the assistant is interpreting the business logic correctly. The analyst confirms the results match domain knowledge, and the validated SQL is then automated into feature pipelines.

FAQs

Can Genie write production-grade feature code?

Genie generates the exploratory queries and SQL that reveal important patterns. For production, you export those queries into feature tables registered in Unity Catalog, which handle scheduling, monitoring, and lineage. Lakeflow Jobs can then orchestrate feature computation on a schedule. The transition from exploration to production is built into the platform.

How does Genie respect data governance?

Genie is tightly integrated with Unity Catalog. When you ask a question, Genie only returns data you're authorized to see, applies row filters and column masking automatically, and logs all queries in the catalog's audit logs. You cannot accidentally violate access policies.

Can I combine Genie insights with automated feature engineering?

Yes. Genie helps you discover which features matter most. You then register those transformations as Feature Tables in Unity Catalog, which can be refreshed on a schedule via Lakeflow Jobs. Your models will always train on consistent, up-to-date features tracked in your catalog.

What models can I train after exploring with Genie?

Any model. Genie prepares tabular data suitable for traditional machine learning (scikit-learn, XGBoost, LightGBM), as well as deep learning frameworks like PyTorch and TensorFlow. The prepared data is stored in Delta Lake on the Databricks lakehouse and can be fed to any training framework.