How to do batch prediction | Structured Data and Graph Models

There are three ways to do inference with KumoRFM.

Method	Scale	Best Use Case
Single entity	1	Testing / debugging / Agents
Multi-entity	≤ 1,000 per request	Small batches
Batch prediction	Unlimited	Full scoring pipelines

Let’s take a look at how this works by using an e-commerce dataset. First, import the data and setup the graph.

import data & setup graph

1 import kumoai.rfm as rfm
2 import pandas as pd
3 
4 root = "s3://kumo-sdk-public/rfm-datasets/online-shopping"
5 
6 df_users = pd.read_parquet(f"{root}/users.parquet")
7 df_items = pd.read_parquet(f"{root}/items.parquet")
8 df_orders = pd.read_parquet(f"{root}/orders.parquet")
9 
10 graph = rfm.Graph.from_data({
11     'users': df_users,
12     'items': df_items,
13     'orders': df_orders,
14 })
15 
16 model = rfm.KumoRFM(graph)

Predict how many orders a user will place in the next 30 days.

Predictive Query expression:

1 PREDICT COUNT(orders.*, 0, 30, days)

Single-Entity Inference

Run inference for a single entity:

single-entity

1 query = "PREDICT COUNT(orders.*, 0, 30, days) FOR users.user_id = 0"
2 result = model.predict(query)

Multi-Entity Inference (≤ 1,000 Entities)

Score multiple specific entities in a single call:

multi-entity

1 query = (
2     "PREDICT COUNT(orders.*, 0, 30, days) "
3     "FOR users.user_id IN (0, 1, 2, 3)"
4 )
5 
6 result = model.predict(query)

Batch Prediction (Full Dataset or Large Batches)

Batch prediction allows scoring all entities. The SDK automatically handles batching and retries.

batch-prediction

1 # collect all user IDs
2 indices = df_users["user_id"].tolist()
3 
4 with model.batch_mode(batch_size="max", num_retries=1):
5     result = model.predict(
6         "PREDICT COUNT(orders.*, 0, 30, days) FOR EACH users.user_id",
7         indices=indices
8     )

Parameters

batch_size=“max”: automatically uses the largest valid batch size
num_retries: retry count for transient failures