| State | Draft |
| Discussion Thread | https://lists.apache.org/thread/dgpgvoszh52vxxszmg65wmcgxnj9zwby |
| Vote Thread | |
| Vote Result Thread | |
| Progress Tracking (PR/GitHub Project/Issue Label) | |
| Date Created |
|
| Version Released | |
| Authors | Pavan Kumar Unknown User (kaxilnaik) |
In today's evolving data landscape, organizations face significant challenges:
Real-World Pain Points:
This proposal leverages Airflow's strengths (production reliability, 1000+ integrations, governance) while adding AI-native capabilities for intelligent data operations.
Based on Pydantic AI's capabilities, while it provides excellent agent framework, multi-model support, and structured outputs, it lacks:
An example view of how Operators look like below:
from airflow.providers.ai.operators import (
LLMSQLQueryOperator,
LLMSchemaCompareOperator,
LLMDataQualityOperator,
LLMFileAnalysisOperator
)
from airflow.sdk import Asset
# Enhanced Asset with structured metadata
customer_postgres = Asset(
name="customer_data_postgres",
uri="postgres://warehouse/public/customers",
conn_id="postgres_warehouse",
schema={
"customer_id": "integer PRIMARY KEY",
"full_name": "varchar(255) NOT NULL",
"email_address": "varchar(255) UNIQUE",
"created_at": "timestamp DEFAULT now()",
"total_revenue": "decimal(10,2)"
},
sensitivity="pii"
)
# SQL operator with automatic context injection and safety
# Option 1: Traditional operator approach
sql_analysis = LLMSQLQueryOperator(
task_id="analyze_customer_segments",
prompt="Find top 10 customers by revenue, include their email and signup date",
data_sources=[customer_postgres],
# Automatic context injection includes:
# - Database type: PostgreSQL
# - Available tables and schemas from DbApiHook
# - Column types and constraints
# - Sample data (first 5 rows) -> explicit approvals with htil to read data if its PII or PCI
# - Built-in SQL safety (blocks DROP, DELETE without WHERE, etc.)
)
# Option 2: Decorator approach with pre-processing
@task.llm_sql_query(data_sources=[customer_postgres])
def analyze_customer_segments_with_preprocessing():
# Custom pre-processing logic
current_date = datetime.now().strftime('%Y-%m-%d')
business_hours = get_business_hours()
# Dynamic prompt generation based on context
prompt = f"""
Find top 10 customers by revenue as of {current_date}.
Include their email and signup date.
Filter for customers active during business hours: {business_hours}
"""
return {"prompt": prompt, "additional_context": {"analysis_date": current_date}} |
Each LLM operator has a corresponding decorator for more flexible, Pythonic workflows:
from airflow.providers.ai.decorators import task
# Schema comparison with custom logic
@task.llm_schema_compare(data_sources=[s3_asset, postgres_asset])
def intelligent_schema_validation():
# Pre-processing: check business calendar
is_migration_window = check_migration_window()
if is_migration_window:
prompt = "Compare schemas and generate migration plan for scheduled maintenance window"
else:
prompt = "Compare schemas and flag breaking changes - no migrations allowed"
return {
"prompt": prompt,
"migration_allowed": is_migration_window,
"additional_context": {"maintenance_window": is_migration_window}
}
# Data quality with dynamic rules
@task.llm_data_quality(data_sources=[customer_asset])
def adaptive_quality_checks():
# Pre-processing: get current business rules
current_rules = fetch_business_rules()
seasonal_adjustments = get_seasonal_data_patterns()
prompt = f"""
Validate customer data against current business rules:
{current_rules}
Apply seasonal adjustments for data volume expectations:
{seasonal_adjustments}
Generate appropriate validation queries.
"""
return {
"prompt": prompt,
"business_rules": current_rules,
"seasonal_context": seasonal_adjustments
}
# File analysis with preprocessing
@task.llm_file_analysis(data_sources=[log_files_asset])
def analyze_logs_with_context():
# Pre-processing: get system context
recent_deployments = get_recent_deployments()
system_alerts = get_active_alerts()
prompt = f"""
Analyze log files for anomalies, considering:
- Recent deployments: {recent_deployments}
- Active system alerts: {system_alerts}
Focus on correlation between deployment events and error patterns.
"""
return {
"prompt": prompt,
"deployment_context": recent_deployments,
"alert_context": system_alerts
} |
Benefits of Decorator Approach:
SQL Operator Context (via DbApiHook integration):
# Automatically injected context for PostgreSQL:
{
"database_type": "postgresql",
"version": "15.2",
"available_tables": ["customers", "orders", "products"],
"schema_info": {
"customers": {
"customer_id": {"type": "integer", "nullable": False, "primary_key": True},
"email_address": {"type": "varchar(255)", "nullable": False, "unique": True},
"total_revenue": {"type": "decimal(10,2)", "nullable": True}
}
},
"sample_data": {
"customers": [
{"customer_id": 1, "email_address": "john@example.com", "total_revenue": 1250.00},
{"customer_id": 2, "email_address": "jane@example.com", "total_revenue": 890.50}
]
},
"dialect_features": {
"supports_window_functions": True,
"supports_cte": True,
"date_functions": ["DATE_TRUNC", "EXTRACT", "AGE"]
}
}
File Operator Context (via S3Hook/GCSHook integration):
# Automatically injected context for S3 Parquet files:
{
"storage_type": "s3",
"file_format": "parquet",
"file_size_mb": 245,
"estimated_rows": 1000000,
"schema_info": {
"id": "int64",
"name": "string",
"email": "string",
"signup_date": "timestamp[ns]"
},
"sample_data": [
{"id": 1, "name": "John Doe", "email": "john@example.com"},
{"id": 2, "name": "Jane Smith", "email": "jane@example.com"}
],
"partitioning": ["year", "month"],
"compression": "snappy"
} |
SQL Operator Built-in Protection:
class LLMSQLQueryOperator(BaseOperator):
# Built-in dangerous operation blocking
BLOCKED_KEYWORDS = ["DROP", "TRUNCATE", "DELETE FROM", "ALTER TABLE", "GRANT", "REVOKE"]
DEFAULT_SYSTEM_PROMPT = """You are a SQL expert integrated with {database_type}. |
SAFETY RULES:
- NEVER generate DROP, TRUNCATE, DELETE without WHERE, or ALTER statements
- Always use proper JOIN syntax for this database type
- Respect column types and constraints provided in schema
- Use database-specific functions when available
CONTEXT: You have access to schema info, sample data, and dialect features.
Generate optimized, safe queries that work with {database_type} version {version}."""
File Operator Built-in Protection:
class LLMFileAnalysisOperator(BaseOperator):
ALLOWED_OPERATIONS = ["read", "analyze", "summarize", "validate"]
DEFAULT_SYSTEM_PROMPT = """You are a file analysis expert for {storage_type} {file_format} files. |
SAFETY RULES:
- ONLY read and analyze files, NEVER modify or delete
- Respect file size limits and memory constraints
- Generate efficient queries for large datasets
CONTEXT: File has {estimated_rows} rows, {file_size_mb}MB, partitioned by {partitioning}."""
# Cross-system schema drift detection
schema_drift = LLMSchemaCompareOperator(
task_id="detect_schema_drift",
data_sources=[customer_s3, customer_postgres, customer_snowflake],
prompt="Identify schema mismatches that would break data loading between systems",
# Automatically gets context from each system:
# - S3: Parquet schema, partitioning, file stats
# - PostgreSQL: Table schema, constraints, indexes
# - Snowflake: Column types, clustering keys, data sharing info
)
A new AnalyticsOperator that provides unified access to multi-cloud data:
from airflow.providers.ai.operators import AnalyticsOperator
# Execute queries across different storage systems uniformly
analytics_task = AnalyticsOperator(
task_id="cross_cloud_analysis",
query="{{ ti.xcom_pull(task_ids='quality_check') }}", # LLM-generated query
data_sources=[customer_s3, orders_gcs, inventory_azure], # Multi-cloud
engine="datafusion", # High-performance query engine
output_format="parquet",
output_location="s3://results/analysis/"
)
Why Apache DataFusion for AI Workloads:
Based on Wren AI's experience, DataFusion provides significant advantages for AI-driven data analysis:
We propose both embedded and separate HITL patterns:
Option A: Embedded HITL
(Out of scope this scenario moving this to New AIP)
quality_check = LLMDataQualityOperator( task_id="customer_quality_analysis", data_sources=[customer_s3], prompt="Generate data quality validation queries", require_approval=True, # Built-in HITL approval_timeout=timedelta(hours=2) ) |
Option B: Separate HITL Steps
# Generate queries
generate_queries = LLMDataQualityOperator(
task_id="generate_quality_queries",
data_sources=[customer_s3],
prompt="Generate data quality validation queries",
dry_run=True # Don't execute, just generate
)
# Human approval step
approve_queries = ApprovalOperator(
task_id="approve_queries",
body="{{ ti.xcom_pull(task_ids='generate_quality_queries') }}",
allow_modifications=True # Users can edit generated queries
)
# Execute approved queries
execute_analysis = AnalyticsOperator(
task_id="execute_quality_checks",
query="{{ ti.xcom_pull(task_ids='approve_queries') }}",
data_sources=[customer_s3]
)
generate_queries >> approve_queries >> execute_analysis |
Here's a real-world scenario combining all components:
from datetime import datetime, timedelta
from airflow.sdk import DAG, Asset
from airflow.providers.ai.operators import (
LLMSchemaCompareOperator,
LLMDataQualityOperator,
AnalyticsOperator
)
from airflow.operators.approval import ApprovalOperator
# Define multi-cloud assets with structured metadata
customer_s3 = Asset(
name="customer_feed_s3",
uri="s3://data-lake/customer/",
conn_id="aws_default",
schema={"id": "int32", "name": "string", "email": "string"},
sensitivity="pii",
format="parquet"
)
customer_postgres = Asset(
name="customer_master_postgres",
uri="postgres://warehouse/public/customers",
conn_id="postgres_default",
schema={"customer_id": "integer", "full_name": "varchar", "email_address": "varchar"},
sensitivity="pii"
)
with DAG(
"intelligent_data_validation",
start_date=datetime(2024, 1, 1),
schedule=timedelta(hours=6),
) as dag:
# 1. Detect schema drift between S3 feed and PostgreSQL master
schema_drift = LLMSchemaCompareOperator(
task_id="detect_schema_drift",
data_sources=[customer_s3, customer_postgres],
prompt="Identify schema mismatches that would break data loading",
output_format="structured_report"
)
# 2. Generate data quality queries for new S3 data
generate_quality_checks = LLMDataQualityOperator(
task_id="generate_quality_queries",
data_sources=[customer_s3],
prompts=[
"Generate summary statistics queries",
"Check for duplicate email addresses",
],
dry_run=True
)
# 3. Human approval for generated queries (with edit capability)
approve_queries = ApprovalOperator(
task_id="approve_quality_queries",
body="{{ ti.xcom_pull(task_ids='generate_quality_queries') }}",
allow_modifications=True,
timeout=timedelta(hours=2)
)
# 4. Execute approved quality checks using DataFusion
execute_quality_checks = AnalyticsOperator(
task_id="run_quality_analysis",
query="{{ ti.xcom_pull(task_ids='approve_quality_queries') }}",
data_sources=[customer_s3],
engine="datafusion",
output_location="s3://results/quality-reports/"
)
schema_drift >> generate_quality_checks >> approve_queries >> execute_quality_checks |
While we start with SQL generation, the architecture supports broader AI workflows:
# Future: General AI task abstraction
process_customer_churn = AITask(
task_id="analyze_churn_patterns",
objective="Identify customers at risk of churning and recommend actions",
resources={
"customer_data": snowflake_conn,
"email_system": sendgrid_conn,
"ml_platform": sagemaker_conn
},
constraints={"budget": "$50", "privacy": "pii_protected"},
output_model=ChurnAnalysisReport
)
Decorator Benefits:
Enhanced Asset System
# Proposed Asset structure evolution
class Asset:
name: str
uri: str
conn_id: Optional[str] = None # Direct connection reference
schema: Optional[Dict[str, str]] = None # Structured schema
sensitivity: Optional[str] = None # pii, sensitive, public
format: Optional[str] = None # parquet, json, csv, etc.
statistics: Optional[Dict] = None # Row counts, update times
extra: Dict[str, Any] = field(default_factory=dict) # Backward compatibility
AnalyticsOperator with DataFusion
Flexible HITL Integration
Key Benefit: Can iterate rapidly on operator design and context injection without any core Airflow changes. All functionality works through existing connection and hook infrastructure.
For Airflow:
For Users:
For the Industry: Creates the missing bridge between AI capabilities and production data infrastructure while making data assets first-class citizens in AI workflows.
Overall the proposal is to build LLM's based operators by leveraging the built in airflow connections and assets functionality
None
None
By completing all the phases