​
Solution Background and Business Value

Traditional search engines in e-commerce rely on ranking functions like Okapi BM25 or large language models (LLMs) to measure the semantic similarity between queries and product descriptions. While effective, these methods only consider textual relevance and do not leverage historical user behavior or relational data to optimize for purchases. A data-driven approach goes further by integrating historical business data, user-item relationships, and high-intent behavioral signals. The next frontier in search optimization is task-specific graph learning: models that use the structure of user, query, and item interactions to generate embeddings tuned for conversion, relevance, or revenue. Graph-based approaches excel at:
  • Generalizing to sparse queries or new items by using multi-hop connections.
  • Learning rich user and item representations that capture both preferences and context, enabling personalization beyond simple collaborative filtering.
  • Optimizing business KPIs directly, such as conversion rate, revenue, or engagement, rather than proxy measures like text similarity.
Kumo’s distributed graph learning infrastructure allows companies to train these models at web scale, leveraging billions of interactions, and to rapidly experiment and deploy new search features with minimal engineering overhead. Personalized search with graph learning considers factors such as:
  • User preferences and past interactions—captured through learned embeddings.
  • Conversion likelihood, with ranking objectives tailored to business value (conversions, revenue, etc).
  • Business objectives, ensuring high-value or trending products gain appropriate visibility.

​
Data Requirements and Schema

To build a search recommendation model powered by task-specific graph learning, structured behavioral data is required. Kumo enables customers to start with standard e-commerce tables and progressively add new data types as needed. Core Tables
  1. Stemmed Queries Table
    • Stores unique, normalized search queries.
    • Key attributes:
      • stemmed_query_id: Unique identifier for the canonical/processed query.
      • stemmed_text: Normalized form of the query.
  2. Search Events Table
    • Captures individual user search actions.
    • Key attributes:
      • search_event_id: Unique identifier for the search event.
      • user_id: Who performed the search.
      • stemmed_query_id: Canonical query text.
      • timestamp: When the search occurred.
      • device: Optional device info.
      • raw_query_text: Original search string.
  3. Items Table
    • Stores product information.
    • Key attributes:
      • item_id: Unique identifier for each item.
      • item_name, category, price: Core product details.
      • Optional attributes: description, embeddings from LLMs/vision models.
  4. Users Table
    • Stores customer information.
    • Key attributes:
      • user_id: Unique identifier for each user.
      • Optional features: age, location, join_timestamp.
  5. Conversion Table
    • Captures purchases or high-intent conversions.
    • Key attributes:
      • conversion_id: Unique identifier.
      • user_id, item_id, search_event_id: Who, what, and from which search.
      • timestamp, revenue: Details about the conversion.
Entity Relationship Diagram (ERD)

​
Predictive Modeling with Graph Learning

Modern search and recommendation systems can be substantially improved by learning graph-based embeddings—vector representations for users, queries, and items—that are directly optimized for business goals such as conversions or revenue. Kumo enables this by supporting distributed training of graph neural networks (GNNs) over large, heterogeneous interaction graphs. These embeddings serve multiple roles:
  • Retriever: Rapid candidate generation using vector similarity (e.g., dot product or approximate nearest neighbor search) between query/user and item embeddings.
  • Reranker: Enhanced ranking models (e.g., learning-to-rank architectures) that use these embeddings or their similarity scores as additional input features, capturing complex behavioral patterns and higher-order relationships.

​
Predictive Query Examples

Below are sample predictive queries for common search and recommendation tasks: 1. Query-Item Retrieval Find the top N items most likely to be converted for each canonical (stemmed) query over the past 30 days: 
PREDICT
    LIST_DISTINCT(conversions.item_id, 0, 30 days)
    RANK TOP 50
FOR EACH
    stemmed_queries.stemmed_query_id
2. User-Item Personalization Retrieve the top N items for each user, capturing the user’s preferences and purchase intent: 
PREDICT
    LIST_DISTINCT(conversions.item_id, 0, 30 days)
    RANK TOP 50
FOR EACH
    users.user_id
These predictive queries instruct the Kumo platform to train GNN models that generate embeddings and ranking scores tailored to specific business objectives, such as increasing conversion rate or driving revenue.

​
Technical Approach

  • Graph Construction: Build a heterogeneous graph with users, stemmed queries, items, search events, and conversions as nodes. Edges represent interactions such as searches, clicks, and purchases.
  • GNN Training: Optimize for the specified predictive queries using negative sampling, multi-hop message passing, and loss functions (e.g., softmax cross-entropy) targeting business KPIs.
  • Embedding Utilization: Use the resulting embeddings in retrieval (ANN search or dot product), and/or as feature inputs for downstream learning-to-rank models in the reranker.
  • Evaluation: Monitor metrics such as NDCG, MRR, and conversion rate during offline validation and online A/B testing to assess real-world impact.

​
Example Workflow

  1. Embedding Generation: Train graph neural network models to produce embeddings for users, queries, and items based on historical search events and conversions.
  2. Batch Export: Export and update embeddings (e.g., daily for items, weekly for users) to downstream systems using the Kumo SDK or API.
  3. Retriever Integration: In the retrieval phase, compute vector similarity (e.g., dot product) between query (or user) and item embeddings to select candidate items for each search.
  4. Reranker Integration: In the reranking phase, use embeddings and their similarity scores as input features alongside other ranking signals to improve relevance and personalization.
  5. Experimentation: Support iterative model development by adjusting the graph schema, training objectives, and refresh schedules to optimize search performance.

​
Building Models with Kumo SDK

import kumoai as kumo

# 1. Initialize Kumo SDK and connect to data source
kumo.init(url="https://<customer_id>.kumoai.cloud/api", api_key=API_KEY)
connector = kumo.S3Connector("s3://your-dataset-location/")

# 2. Define tables
stemmed_queries = kumo.Table.from_source_table(
    source_table=connector.table('stemmed_queries'),
    primary_key='stemmed_query_id',
).infer_metadata()

search_events = kumo.Table.from_source_table(
    source_table=connector.table('search_events'),
    primary_key='search_event_id',
).infer_metadata()

items = kumo.Table.from_source_table(
    source_table=connector.table('items'),
    primary_key='item_id',
).infer_metadata()

users = kumo.Table.from_source_table(
    source_table=connector.table('users'),
    primary_key='user_id',
).infer_metadata()

conversions = kumo.Table.from_source_table(
    source_table=connector.table('conversions'),
    primary_key='conversion_id',
).infer_metadata()

# 3. Construct heterogeneous interaction graph
graph = kumo.Graph(
    tables={
        'stemmed_queries': stemmed_queries,
        'search_events': search_events,
        'items': items,
        'users': users,
        'conversions': conversions,
    },
    edges=[
        dict(src_table='search_events', fkey='stemmed_query_id', dst_table='stemmed_queries'),
        dict(src_table='search_events', fkey='user_id', dst_table='users'),
        dict(src_table='conversions', fkey='user_id', dst_table='users'),
        dict(src_table='conversions', fkey='item_id', dst_table='items'),
        dict(src_table='conversions', fkey='search_event_id', dst_table='search_events'),
    ],
)
graph.validate(verbose=True)

# 4. Define predictive queries (pquery) for training

# Query-Item retrieval: For each stemmed query, retrieve top N items by likelihood of conversion
pquery_query_item = kumo.PredictiveQuery(
    graph=graph,
    query="""
    PREDICT
        LIST_DISTINCT(conversions.item_id, 0, 30 days)
        RANK TOP 50
    FOR EACH
        stemmed_queries.stemmed_query_id
    """
)

# User-Item personalization: For each user, retrieve top N items
pquery_user_item = kumo.PredictiveQuery(
    graph=graph,
    query="""
    PREDICT
        LIST_DISTINCT(conversions.item_id, 0, 30 days)
        RANK TOP 50
    FOR EACH
        users.user_id
    """
)

pquery_query_item.validate(verbose=True)
pquery_user_item.validate(verbose=True)

# 5. Suggest and train a GNN model for the pquery (e.g., for query-item retrieval)
model_plan = pquery_query_item.suggest_model_plan()
trainer = kumo.Trainer(model_plan)
training_job = trainer.fit(
    graph=graph,
    train_table=pquery_query_item.generate_training_table(non_blocking=True),
    non_blocking=False,
)
print(f"Training metrics: {training_job.metrics()}")

# 6. Embedding materialization (downstream usage)
# Embeddings can be exported to a feature store or ANN service for use in both retriever and reranker pipelines.
With Kumo, teams can go from data connection to first working GNN-based search model in under a day, and can optimize performance with dozens of experiments over just a couple of weeks—without the heavy engineering burden of building graph learning infrastructure from scratch. If you want to see more code, deeper math, or production pipeline examples, just let us know!