Batch Inference Made Simple: Using Databricks Serverless Model Serving and AI Functions

Once you have deployed a machine learning model to production, it typically falls into one of two categories. First, there are real-time models that need to be always available, ready to serve predictions with minimal latency. These handle requests one at a time(or in small batches), but the requests come in unpredictably throughout the day. Think fraud detection systems that need to evaluate transactions as they happen-speed is critical. Then there are batch inference models that run on a schedule, usually processing large volumes of predictions all at once. Maybe you are generating demand forecasts for thousands of products every week, or creating monthly revenue projections across all your sales regions. These jobs don’t need instant responses, but they need to handle significant scale. Traditionally, serverless model serving has been reserved for real-time use cases. For batch predictions, teams would typically pull a saved model from a registry and incorporate it into a data pipeline to process requests in bulk. This means a data engineer has to write expensive Python code to loop through batches, run predictions, and integrate all of that into their ETL pipelines. On top of that, they need to manage the infrastructure-spinning up clusters when needed, ensuring they are sized correctly for the workloads, and remembering to shut them down afterward to avoid unnecessary costs. This approach also creates a barrier for SQL analysts who want to use these models directly in their queries and scripts. They end up dependent on engineering teams to run predictions for them. To mitigate these issues, I will walk you through how to deploy models on Databricks model serving and use AI Functions to run these batch predictions. That results in zero maintenance in terms of infrastructure and enables sql analysts to incorporate model inference as part of their workflow. Retail Sales Forecasting with Databricks Model Serving Retail organizations need accurate sales forecasts to optimize inventory management, staffing levels and promotional planning, and supply chain operations. Traditional forecasting often requires significant manual effort from data science teams to generate predictions, creating bottlenecks when business analysts need quick insights across multiple stores, products or time horizons. Generating these forecasts required data engineers to manually run batch prediction jobs, manage compute infrastructure and create custom pipelines for each request-making the process slow, expensive and inaccessible to business users. By transforming batch inference into a simple SQL function, we transform what was once a complex engineering tasks requiring Python expertise and infrastructure management into a simple query. In this example we will show you how to train a forecasting model on Databricks, serve it in model serving and run batch inference using AI Functions. Step 1 : Create a sample fake data set that contains retail sales data with realistic patterns. The generated dataset is a synthetic retail sales transactions dataset designed for time series forecasting. Each row represents a single sales transaction and includes the following fields: date: Transaction date store_id, store_name, region: Store identifiers and region product_id, product_name, category: Product identifiers and category quantity: Number of units sold in the transaction unit_price: Price per unit (with random variation) total_sales: Total sales amount before discount discount_pct, discount_amount: Discount applied and its value final_sales: Sales amount after discount day_of_week, month, year, quarter: Temporal features for analysis date: Transaction date store_id, store_name, region: Store identifiers and region product_id, product_name, category: Product identifiers and category quantity: Number of units sold in the transaction unit_price: Price per unit (with random variation) total_sales: Total sales amount before discount discount_pct, discount_amount: Discount applied and its value final_sales: Sales amount after discount day_of_week, month, year, quarter: Temporal features for analysis The dataset covers multiple years, stores, and products, and incorporates realistic retail patterns such as seasonality, holidays, promotions, and growth trends. import pandas as pd import numpy as np def generate_retail_sales_data(start_date='2021-01-01', end_date='2024-12-31', num_stores=10, num_products=50, seed=42): """Generate synthetic retail sales data with realistic patterns.""" np.random.seed(seed) # Generate date range date_range = pd.date_range(start=start_date, end=end_date, freq='D') # Product categories and their base prices categories = ['Electronics', 'Clothing', 'Food', 'Home & Garden', 'Sports'] store_regions = ['North', 'South', 'East', 'West', 'Central'] data = [] # Generate product catalog products = [] for i in range(num_products): category = np.random.choice(categories) base_price = np.random.uniform(10, 500) if category == 'Electronics': base_price = np.random.uniform(100, 2000) elif category == 'Food': base_price = np.random.uniform(5, 50) products.append({ 'product_id': f'P{i+1:04d}', 'product_name': f'{category}_Product_{i+1}', 'category': category, 'base_price': base_price }) # Generate store information stores = [] for i in range(num_stores): stores.append({ 'store_id': f'S{i+1:03d}', 'store_name': f'Store_{i+1}', 'region': store_regions[i % len(store_regions)] }) # Generate sales transactions for date in date_range: day_of_week = date.dayofweek month = date.month # Apply seasonal factors weekend_factor = 1.2 if day_of_week >= 5 else 1.0 # 20% boost on weekends holiday_factor = 1.5 if month in [11, 12] else 1.0 # 50% boost in holiday season summer_factor = 0.85 if month in [6, 7, 8] else 1.0 # 15% slowdown in summer # Apply growth trend (5% annual growth) days_since_start = (date - pd.Timestamp(start_date)).days trend_factor = 1 + (days_since_start / 365) * 0.05 # Generate transactions for each store for store in stores: store_factor = np.random.uniform(0.7, 1.3) # Calculate number of transactions for this store on this day base_transactions = int(np.random.poisson(15) * weekend_factor * holiday_factor * summer_factor * trend_factor * store_factor) for _ in range(base_transactions): product = np.random.choice(products) # Quantity sold per transaction quantity = np.random.choice([1, 2, 3, 4, 5], p=[0.5, 0.25, 0.15, 0.07, 0.03]) # Apply price variation (promotions, dynamic pricing) price_variation = np.random.uniform(0.85, 1.15) unit_price = product['base_price'] * price_variation # Calculate sales amounts total_sales = unit_price * quantity # Apply discounts discount_pct = np.random.choice([0, 5, 10, 15, 20], p=[0.6, 0.2, 0.1, 0.07, 0.03]) discount_amount = total_sales * (discount_pct / 100) final_sales = total_sales - discount_amount # Append transaction record data.append({ 'date': date, 'store_id': store['store_id'], 'store_name': store['store_name'], 'region': store['region'], 'product_id': product['product_id'], 'product_name': product['product_name'], 'category': product['category'], 'quantity': quantity, 'unit_price': round(unit_price, 2), 'total_sales': round(total_sales, 2), 'discount_pct': discount_pct, 'discount_amount': round(discount_amount, 2), 'final_sales': round(final_sales, 2), 'day_of_week': date.day_name(), 'month': date.month, 'year': date.year, 'quarter': f'Q{(date.month-1)//3 + 1}' }) df = pd.DataFrame(data) return df def save_to_unity_catalog(df, catalog_name, schema_name, table_name): """Save DataFrame to Unity Catalog as a Delta table.""" spark_df = spark.createDataFrame(df) # Handle schema names with hyphens if "-" in schema_name: schema_name = f"`{schema_name}`" full_table_name = f"{catalog_name}.{schema_name}.{table_name}" # Write as Delta table spark_df.write \ .format("delta") \ .mode("overwrite") \ .option("overwriteSchema", "true") \ .saveAsTable(full_table_name) return full_table_name # Generate retail sales data sales_df = generate_retail_sales_data( start_date='2021-01-01', end_date='2024-12-31', num_stores=10, num_products=50, seed=42 ) # Configure Unity Catalog location catalog = "samaya_madhavan" schema = "serverless-ml" table = "retail_sales_forecast" # Save to Unity Catalog table_path = save_to_unity_catalog(sales_df, catalog, schema, table) # Display summary print(f"Generated {len(sales_df):,} transaction records") print(f"Date range: {sales_df['date'].min()} to {sales_df['date'].max()}") print(f"Total sales: ${sales_df['final_sales'].sum():,.2f}") print(f"Data saved to: {table_path}") import pandas as pd import numpy as np def generate_retail_sales_data(start_date='2021-01-01', end_date='2024-12-31', num_stores=10, num_products=50, seed=42): """Generate synthetic retail sales data with realistic patterns.""" np.random.seed(seed) # Generate date range date_range = pd.date_range(start=start_date, end=end_date, freq='D') # Product categories and their base prices categories = ['Electronics', 'Clothing', 'Food', 'Home & Garden', 'Sports'] store_regions = ['North', 'South', 'East', 'West', 'Central'] data = [] # Generate product catalog products = [] for i in range(num_products): category = np.random.choice(categories) base_price = np.random.uniform(10, 500) if category == 'Electronics': base_price = np.random.uniform(100, 2000) elif category == 'Food': base_price = np.random.uniform(5, 50) products.append({ 'product_id': f'P{i+1:04d}', 'product_name': f'{category}_Product_{i+1}', 'category': category, 'base_price': base_price }) # Generate store information stores = [] for i in range(num_stores): stores.append({ 'store_id': f'S{i+1:03d}', 'store_name': f'Store_{i+1}', 'region': store_regions[i % len(store_regions)] }) # Generate sales transactions for date in date_range: day_of_week = date.dayofweek month = date.month # Apply seasonal factors weekend_factor = 1.2 if day_of_week >= 5 else 1.0 # 20% boost on weekends holiday_factor = 1.5 if month in [11, 12] else 1.0 # 50% boost in holiday season summer_factor = 0.85 if month in [6, 7, 8] else 1.0 # 15% slowdown in summer # Apply growth trend (5% annual growth) days_since_start = (date - pd.Timestamp(start_date)).days trend_factor = 1 + (days_since_start / 365) * 0.05 # Generate transactions for each store for store in stores: store_factor = np.random.uniform(0.7, 1.3) # Calculate number of transactions for this store on this day base_transactions = int(np.random.poisson(15) * weekend_factor * holiday_factor * summer_factor * trend_factor * store_factor) for _ in range(base_transactions): product = np.random.choice(products) # Quantity sold per transaction quantity = np.random.choice([1, 2, 3, 4, 5], p=[0.5, 0.25, 0.15, 0.07, 0.03]) # Apply price variation (promotions, dynamic pricing) price_variation = np.random.uniform(0.85, 1.15) unit_price = product['base_price'] * price_variation # Calculate sales amounts total_sales = unit_price * quantity # Apply discounts discount_pct = np.random.choice([0, 5, 10, 15, 20], p=[0.6, 0.2, 0.1, 0.07, 0.03]) discount_amount = total_sales * (discount_pct / 100) final_sales = total_sales - discount_amount # Append transaction record data.append({ 'date': date, 'store_id': store['store_id'], 'store_name': store['store_name'], 'region': store['region'], 'product_id': product['product_id'], 'product_name': product['product_name'], 'category': product['category'], 'quantity': quantity, 'unit_price': round(unit_price, 2), 'total_sales': round(total_sales, 2), 'discount_pct': discount_pct, 'discount_amount': round(discount_amount, 2), 'final_sales': round(final_sales, 2), 'day_of_week': date.day_name(), 'month': date.month, 'year': date.year, 'quarter': f'Q{(date.month-1)//3 + 1}' }) df = pd.DataFrame(data) return df def save_to_unity_catalog(df, catalog_name, schema_name, table_name): """Save DataFrame to Unity Catalog as a Delta table.""" spark_df = spark.createDataFrame(df) # Handle schema names with hyphens if "-" in schema_name: schema_name = f"`{schema_name}`" full_table_name = f"{catalog_name}.{schema_name}.{table_name}" # Write as Delta table spark_df.write \ .format("delta") \ .mode("overwrite") \ .option("overwriteSchema", "true") \ .saveAsTable(full_table_name) return full_table_name # Generate retail sales data sales_df = generate_retail_sales_data( start_date='2021-01-01', end_date='2024-12-31', num_stores=10, num_products=50, seed=42 ) # Configure Unity Catalog location catalog = "samaya_madhavan" schema = "serverless-ml" table = "retail_sales_forecast" # Save to Unity Catalog table_path = save_to_unity_catalog(sales_df, catalog, schema, table) # Display summary print(f"Generated {len(sales_df):,} transaction records") print(f"Date range: {sales_df['date'].min()} to {sales_df['date'].max()}") print(f"Total sales: ${sales_df['final_sales'].sum():,.2f}") print(f"Data saved to: {table_path}") The output of this is saved to unity catalog : Step 2 : Build, register and deploy a time series forecasting model. The next step in the process is to build, register and deploy a time series forecasting model using scikit-learn’s LinearRegression, with MLFlow for model management and Databricks Model Serving for deployment. We start by loading and preparing retail sales data from the unity catalog table, aggregating daily sales for a specific store. The data is split into training and test sets, and features such as day index, day of week, and month are engineered for the model. A LinearRegression model is trained, evaluated, and logged to MLflow along with a StandardScaler for feature normalization. A custom MLflow pyfunc wrapper is then created to enable multi-period forecasting, and the model is registered in Unity Catalog. The notebook proceeds to deploy the model to a Databricks Model Serving endpoint. import pandas as pd import numpy as np import mlflow import mlflow.sklearn from sklearn.linear_model import LinearRegression from sklearn.preprocessing import StandardScaler from mlflow.models import infer_signature # Configuration catalog = "samaya_madhavan" schema = "serverless-ml" source_table = "retail_sales_forecast" model_name = f"{catalog}.{schema}.retail_forecast_model" endpoint_name = "retail-forecast-endpoint" # Step 1: Load and prepare data print("Loading data...") df = spark.sql(f""" SELECT date, SUM(final_sales) AS sales FROM {catalog}.`{schema}`.{source_table} WHERE store_id = 'S001' GROUP BY date ORDER BY date """).toPandas() df['date'] = pd.to_datetime(df['date']) df['day_index'] = (df['date'] - df['date'].min()).dt.days df['day_of_week'] = df['date'].dt.dayofweek df['month'] = df['date'].dt.month # Create train/test split train_size = int(len(df) * 0.8) train_df = df[:train_size] test_df = df[train_size:] X_train = train_df[['day_index', 'day_of_week', 'month']] y_train = train_df['sales'] X_test = test_df[['day_index', 'day_of_week', 'month']] y_test = test_df['sales'] print(f"Training samples: {len(X_train)}, Test samples: {len(X_test)}") # Step 2: Train sklearn model print("\nTraining LinearRegression model...") mlflow.set_registry_uri("databricks-uc") with mlflow.start_run(run_name="sklearn_forecast") as run: # Train scaler and model scaler = StandardScaler() X_train_scaled = scaler.fit_transform(X_train) X_test_scaled = scaler.transform(X_test) model = LinearRegression() model.fit(X_train_scaled, y_train) # Evaluate train_score = model.score(X_train_scaled, y_train) test_score = model.score(X_test_scaled, y_test) print(f"Train R²: {train_score:.4f}") print(f"Test R²: {test_score:.4f}") # Log metrics mlflow.log_metric("train_r2", train_score) mlflow.log_metric("test_r2", test_score) # Save scaler separately import joblib joblib.dump(scaler, "/tmp/scaler.pkl") mlflow.log_artifact("/tmp/scaler.pkl") # Create input/output examples for signature input_example = pd.DataFrame({'periods': [30]}) # Generate sample forecast last_day = df['day_index'].max() last_date = df['date'].max() forecast_rows = [] for i in range(1, 6): # Sample 5 days future_date = last_date + pd.Timedelta(days=i) future_day = last_day + i X_future = pd.DataFrame([{ 'day_index': future_day, 'day_of_week': future_date.dayofweek, 'month': future_date.month }]) X_future_scaled = scaler.transform(X_future) pred = model.predict(X_future_scaled)[0] forecast_rows.append({ 'date': future_date.strftime('%Y-%m-%d'), 'forecast': float(pred), 'forecast_lower': float(pred * 0.9), 'forecast_upper': float(pred * 1.1) }) output_example = pd.DataFrame(forecast_rows) signature = infer_signature(input_example, output_example) # Log model with sklearn flavor mlflow.sklearn.log_model( sk_model=model, artifact_path="model", signature=signature, input_example=input_example, registered_model_name=model_name, serialization_format="cloudpickle" ) print(f"\nModel registered: {model_name}") # Step 3: Create wrapper model for forecasting print("\nCreating forecast wrapper...") with mlflow.start_run(run_name="sklearn_forecast_wrapper") as run: class SklearnForecastWrapper(mlflow.pyfunc.PythonModel): def load_context(self, context): import joblib self.model = joblib.load(context.artifacts["model"]) self.scaler = joblib.load(context.artifacts["scaler"]) self.last_day_index = int(open(context.artifacts["last_day"]).read()) self.last_date = pd.to_datetime(open(context.artifacts["last_date"]).read()) def predict(self, context, model_input): if isinstance(model_input, pd.DataFrame) and 'periods' in model_input.columns: periods = int(model_input['periods'].iloc[0]) else: periods = 30 forecasts = [] for i in range(1, periods + 1): future_date = self.last_date + pd.Timedelta(days=i) future_day = self.last_day_index + i X_future = pd.DataFrame([{ 'day_index': future_day, 'day_of_week': future_date.dayofweek, 'month': future_date.month }]) X_scaled = self.scaler.transform(X_future) pred = self.model.predict(X_scaled)[0] forecasts.append({ 'date': future_date.strftime('%Y-%m-%d'), 'forecast': float(pred), 'forecast_lower': float(pred * 0.9), 'forecast_upper': float(pred * 1.1) }) # Return in ai_query compatible format - single element array return pd.DataFrame([{"predictions": forecasts}]) # Save artifacts import joblib joblib.dump(model, "/tmp/sklearn_model.pkl") joblib.dump(scaler, "/tmp/sklearn_scaler.pkl") with open("/tmp/last_day.txt", "w") as f: f.write(str(int(last_day))) with open("/tmp/last_date.txt", "w") as f: f.write(str(last_date)) # Create signature input_example = pd.DataFrame({'periods': [30]}) wrapper = SklearnForecastWrapper() wrapper.model = model wrapper.scaler = scaler wrapper.last_day_index = int(last_day) wrapper.last_date = last_date output_example = wrapper.predict(None, input_example) signature = infer_signature(input_example, output_example) # Log wrapper model mlflow.pyfunc.log_model( artifact_path="model", python_model=SklearnForecastWrapper(), artifacts={ "model": "/tmp/sklearn_model.pkl", "scaler": "/tmp/sklearn_scaler.pkl", "last_day": "/tmp/last_day.txt", "last_date": "/tmp/last_date.txt" }, signature=signature, input_example=input_example, registered_model_name=model_name ) print(f"Wrapper model registered: {model_name}") # Step 4: Get latest version and set alias from mlflow.tracking import MlflowClient client = MlflowClient() model_versions = client.search_model_versions(f"name='{model_name}'") latest_version = max([int(mv.version) for mv in model_versions]) client.set_registered_model_alias(model_name, "Champion", str(latest_version)) print(f"\nLatest version: {latest_version} (alias: Champion)") # Step 5: Deploy to Model Serving print("\nDeploying to Model Serving...") from mlflow.deployments import get_deploy_client deploy_client = get_deploy_client("databricks") endpoint_config = { "served_entities": [{ "name": "retail-forecast-model", "entity_name": model_name, "entity_version": str(latest_version), "workload_size": "Small", "scale_to_zero_enabled": True }], "traffic_config": { "routes": [{"served_model_name": "retail-forecast-model", "traffic_percentage": 100}] } } try: deploy_client.get_endpoint(endpoint=endpoint_name) deploy_client.update_endpoint(endpoint=endpoint_name, config=endpoint_config) print(f"✓ Updated endpoint: {endpoint_name}") except: deploy_client.create_endpoint(name=endpoint_name, config=endpoint_config) print(f"✓ Created endpoint: {endpoint_name}") # Step 6: Test the endpoint print("\nTesting endpoint...") import time time.sleep(10) try: result = deploy_client.predict( endpoint=endpoint_name, inputs={"dataframe_records": [{"periods": 7}]} ) print("\nSample forecast (7 days):") if isinstance(result, dict) and 'predictions' in result: print(pd.DataFrame(result['predictions']).head()) else: print(result) except Exception as e: print(f"Endpoint may still be starting: {e}") print(f"\n{'='*60}") print(f"Deployment Complete!") print(f"{'='*60}") print(f"Model: {model_name}") print(f"Version: {latest_version}") print(f"Endpoint: {endpoint_name}") import pandas as pd import numpy as np import mlflow import mlflow.sklearn from sklearn.linear_model import LinearRegression from sklearn.preprocessing import StandardScaler from mlflow.models import infer_signature # Configuration catalog = "samaya_madhavan" schema = "serverless-ml" source_table = "retail_sales_forecast" model_name = f"{catalog}.{schema}.retail_forecast_model" endpoint_name = "retail-forecast-endpoint" # Step 1: Load and prepare data print("Loading data...") df = spark.sql(f""" SELECT date, SUM(final_sales) AS sales FROM {catalog}.`{schema}`.{source_table} WHERE store_id = 'S001' GROUP BY date ORDER BY date """).toPandas() df['date'] = pd.to_datetime(df['date']) df['day_index'] = (df['date'] - df['date'].min()).dt.days df['day_of_week'] = df['date'].dt.dayofweek df['month'] = df['date'].dt.month # Create train/test split train_size = int(len(df) * 0.8) train_df = df[:train_size] test_df = df[train_size:] X_train = train_df[['day_index', 'day_of_week', 'month']] y_train = train_df['sales'] X_test = test_df[['day_index', 'day_of_week', 'month']] y_test = test_df['sales'] print(f"Training samples: {len(X_train)}, Test samples: {len(X_test)}") # Step 2: Train sklearn model print("\nTraining LinearRegression model...") mlflow.set_registry_uri("databricks-uc") with mlflow.start_run(run_name="sklearn_forecast") as run: # Train scaler and model scaler = StandardScaler() X_train_scaled = scaler.fit_transform(X_train) X_test_scaled = scaler.transform(X_test) model = LinearRegression() model.fit(X_train_scaled, y_train) # Evaluate train_score = model.score(X_train_scaled, y_train) test_score = model.score(X_test_scaled, y_test) print(f"Train R²: {train_score:.4f}") print(f"Test R²: {test_score:.4f}") # Log metrics mlflow.log_metric("train_r2", train_score) mlflow.log_metric("test_r2", test_score) # Save scaler separately import joblib joblib.dump(scaler, "/tmp/scaler.pkl") mlflow.log_artifact("/tmp/scaler.pkl") # Create input/output examples for signature input_example = pd.DataFrame({'periods': [30]}) # Generate sample forecast last_day = df['day_index'].max() last_date = df['date'].max() forecast_rows = [] for i in range(1, 6): # Sample 5 days future_date = last_date + pd.Timedelta(days=i) future_day = last_day + i X_future = pd.DataFrame([{ 'day_index': future_day, 'day_of_week': future_date.dayofweek, 'month': future_date.month }]) X_future_scaled = scaler.transform(X_future) pred = model.predict(X_future_scaled)[0] forecast_rows.append({ 'date': future_date.strftime('%Y-%m-%d'), 'forecast': float(pred), 'forecast_lower': float(pred * 0.9), 'forecast_upper': float(pred * 1.1) }) output_example = pd.DataFrame(forecast_rows) signature = infer_signature(input_example, output_example) # Log model with sklearn flavor mlflow.sklearn.log_model( sk_model=model, artifact_path="model", signature=signature, input_example=input_example, registered_model_name=model_name, serialization_format="cloudpickle" ) print(f"\nModel registered: {model_name}") # Step 3: Create wrapper model for forecasting print("\nCreating forecast wrapper...") with mlflow.start_run(run_name="sklearn_forecast_wrapper") as run: class SklearnForecastWrapper(mlflow.pyfunc.PythonModel): def load_context(self, context): import joblib self.model = joblib.load(context.artifacts["model"]) self.scaler = joblib.load(context.artifacts["scaler"]) self.last_day_index = int(open(context.artifacts["last_day"]).read()) self.last_date = pd.to_datetime(open(context.artifacts["last_date"]).read()) def predict(self, context, model_input): if isinstance(model_input, pd.DataFrame) and 'periods' in model_input.columns: periods = int(model_input['periods'].iloc[0]) else: periods = 30 forecasts = [] for i in range(1, periods + 1): future_date = self.last_date + pd.Timedelta(days=i) future_day = self.last_day_index + i X_future = pd.DataFrame([{ 'day_index': future_day, 'day_of_week': future_date.dayofweek, 'month': future_date.month }]) X_scaled = self.scaler.transform(X_future) pred = self.model.predict(X_scaled)[0] forecasts.append({ 'date': future_date.strftime('%Y-%m-%d'), 'forecast': float(pred), 'forecast_lower': float(pred * 0.9), 'forecast_upper': float(pred * 1.1) }) # Return in ai_query compatible format - single element array return pd.DataFrame([{"predictions": forecasts}]) # Save artifacts import joblib joblib.dump(model, "/tmp/sklearn_model.pkl") joblib.dump(scaler, "/tmp/sklearn_scaler.pkl") with open("/tmp/last_day.txt", "w") as f: f.write(str(int(last_day))) with open("/tmp/last_date.txt", "w") as f: f.write(str(last_date)) # Create signature input_example = pd.DataFrame({'periods': [30]}) wrapper = SklearnForecastWrapper() wrapper.model = model wrapper.scaler = scaler wrapper.last_day_index = int(last_day) wrapper.last_date = last_date output_example = wrapper.predict(None, input_example) signature = infer_signature(input_example, output_example) # Log wrapper model mlflow.pyfunc.log_model( artifact_path="model", python_model=SklearnForecastWrapper(), artifacts={ "model": "/tmp/sklearn_model.pkl", "scaler": "/tmp/sklearn_scaler.pkl", "last_day": "/tmp/last_day.txt", "last_date": "/tmp/last_date.txt" }, signature=signature, input_example=input_example, registered_model_name=model_name ) print(f"Wrapper model registered: {model_name}") # Step 4: Get latest version and set alias from mlflow.tracking import MlflowClient client = MlflowClient() model_versions = client.search_model_versions(f"name='{model_name}'") latest_version = max([int(mv.version) for mv in model_versions]) client.set_registered_model_alias(model_name, "Champion", str(latest_version)) print(f"\nLatest version: {latest_version} (alias: Champion)") # Step 5: Deploy to Model Serving print("\nDeploying to Model Serving...") from mlflow.deployments import get_deploy_client deploy_client = get_deploy_client("databricks") endpoint_config = { "served_entities": [{ "name": "retail-forecast-model", "entity_name": model_name, "entity_version": str(latest_version), "workload_size": "Small", "scale_to_zero_enabled": True }], "traffic_config": { "routes": [{"served_model_name": "retail-forecast-model", "traffic_percentage": 100}] } } try: deploy_client.get_endpoint(endpoint=endpoint_name) deploy_client.update_endpoint(endpoint=endpoint_name, config=endpoint_config) print(f"✓ Updated endpoint: {endpoint_name}") except: deploy_client.create_endpoint(name=endpoint_name, config=endpoint_config) print(f"✓ Created endpoint: {endpoint_name}") # Step 6: Test the endpoint print("\nTesting endpoint...") import time time.sleep(10) try: result = deploy_client.predict( endpoint=endpoint_name, inputs={"dataframe_records": [{"periods": 7}]} ) print("\nSample forecast (7 days):") if isinstance(result, dict) and 'predictions' in result: print(pd.DataFrame(result['predictions']).head()) else: print(result) except Exception as e: print(f"Endpoint may still be starting: {e}") print(f"\n{'='*60}") print(f"Deployment Complete!") print(f"{'='*60}") print(f"Model: {model_name}") print(f"Version: {latest_version}") print(f"Endpoint: {endpoint_name}") The model is now served on Databricks Model Serving Step 3 : Use AI Functions to obtain forecasts from the model served on Model Serving. Finally, we demonstrate how to use AI Functions to run batch predictions. AI Functions : Databricks offers AI Functions as native, ready-to-use capabilities that enable you to perform AI operations-such as translating text or analyzing sentiment-directly on data stored within the Databricks platform. These functions are accessible across the entire Databricks ecosystem, including SQL environments, notebooks, Lakeflow Declarative Pipelines, and Workflows. AI Functions AI Functions The AI Functions catalog includes two categories : Specialized functions deliver targeted AI capabilities for specific use cases, such as text summarization and language translation. These functions utilize cutting-edge generative AI models that Databricks hosts and maintains. Specialized functions General purpose functions:ai_query() serves as a flexible, multi-purpose function that enables you to apply virtually any AI model to your datasets. General purpose functions:ai_query() In this blog we will show you examples of how to use ai_query(). The following code shows how to use the SQL ai_query() function to obtain forecasts from the deployed Model Serving endpoint on Databricks. Example 1: 30 day forecast. Example 1: 30 day forecast. forecast_30 = spark.sql(f""" SELECT ai_query( '{endpoint_name}', named_struct('periods', 30) ) AS forecast_json """) display(forecast_30) forecast_30 = spark.sql(f""" SELECT ai_query( '{endpoint_name}', named_struct('periods', 30) ) AS forecast_json """) display(forecast_30) Example 2: Multiple forecast horizons Example 2: Multiple forecast horizons multi_forecast = spark.sql(f""" SELECT horizon, days, ai_query( '{endpoint_name}', named_struct('periods', days) ) AS forecast_json FROM ( SELECT 'Short-term' AS horizon, 7 AS days UNION ALL SELECT 'Medium-term' AS horizon, 30 AS days UNION ALL SELECT 'Long-term' AS horizon, 90 AS days ) """) multi_forecast = spark.sql(f""" SELECT horizon, days, ai_query( '{endpoint_name}', named_struct('periods', days) ) AS forecast_json FROM ( SELECT 'Short-term' AS horizon, 7 AS days UNION ALL SELECT 'Medium-term' AS horizon, 30 AS days UNION ALL SELECT 'Long-term' AS horizon, 90 AS days ) """) Example 3: Generate forecasts for all stores Example 3: Generate forecasts for all stores source_table = "retail_sales_forecast" # Get list of all stores stores = spark.sql(f""" SELECT DISTINCT store_id, region FROM {catalog}.`{schema}`.{source_table} ORDER BY store_id """) display(stores) # Create batch forecast requests batch_requests = spark.sql(f""" CREATE OR REPLACE TABLE {catalog}.`{schema}`.batch_forecast_requests AS SELECT store_id, region, 30 AS forecast_periods, CURRENT_TIMESTAMP() AS request_timestamp FROM ( SELECT DISTINCT store_id, region FROM {catalog}.`{schema}`.{source_table} ) """) print("✓ Batch requests table created") # Execute batch forecasts using ai_query batch_forecasts = spark.sql(f""" CREATE OR REPLACE TABLE {catalog}.`{schema}`.batch_forecasts AS SELECT r.store_id, r.region, r.forecast_periods, r.request_timestamp, pred.date AS forecast_date, pred.forecast AS predicted_sales, pred.forecast_lower, pred.forecast_upper FROM {catalog}.`{schema}`.batch_forecast_requests r LATERAL VIEW explode( ai_query( '{endpoint_name}', named_struct('periods', r.forecast_periods) ).predictions ) AS pred """) print("✓ Batch forecasts completed and saved") # View results results = spark.sql(f""" SELECT * FROM {catalog}.`{schema}`.batch_forecasts ORDER BY store_id, forecast_date LIMIT 50 """) display(results) source_table = "retail_sales_forecast" # Get list of all stores stores = spark.sql(f""" SELECT DISTINCT store_id, region FROM {catalog}.`{schema}`.{source_table} ORDER BY store_id """) display(stores) # Create batch forecast requests batch_requests = spark.sql(f""" CREATE OR REPLACE TABLE {catalog}.`{schema}`.batch_forecast_requests AS SELECT store_id, region, 30 AS forecast_periods, CURRENT_TIMESTAMP() AS request_timestamp FROM ( SELECT DISTINCT store_id, region FROM {catalog}.`{schema}`.{source_table} ) """) print("✓ Batch requests table created") # Execute batch forecasts using ai_query batch_forecasts = spark.sql(f""" CREATE OR REPLACE TABLE {catalog}.`{schema}`.batch_forecasts AS SELECT r.store_id, r.region, r.forecast_periods, r.request_timestamp, pred.date AS forecast_date, pred.forecast AS predicted_sales, pred.forecast_lower, pred.forecast_upper FROM {catalog}.`{schema}`.batch_forecast_requests r LATERAL VIEW explode( ai_query( '{endpoint_name}', named_struct('periods', r.forecast_periods) ).predictions ) AS pred """) print("✓ Batch forecasts completed and saved") # View results results = spark.sql(f""" SELECT * FROM {catalog}.`{schema}`.batch_forecasts ORDER BY store_id, forecast_date LIMIT 50 """) display(results) For more detailed examples and complete code examples refer to github repository. For more detailed examples and complete code examples refer to github repository. github repository By deploying our retail sales forecasting model to Databricks Model Serving and leveraging AI Functions, we eliminated the need for data engineers to write custom Python loops, manage cluster infrastructure, or worry about scaling batch prediction workloads. SQL analysts can now generate 30-day forecasts for a single store, compare short-term versus long-term predictions, or produce forecasts across all ten stores in our retail chain—all with simple SQL queries using ai_query(). By bridging the gap between model deployment and data analysis, AI Functions don't just simplify batch inference—they democratize it. This is the paradigm shift that makes AI truly accessible: when running predictions becomes as simple as writing a WHERE clause, organizations can finally unlock intelligence from their data at the speed of thought rather than the speed of engineering sprints.