Graph Construction
Accelerating Graph Construction for Production Graph Transformer Workloads
Three optimizations that cut production graph transformer pipeline times by 12-25%: ingestion skip, eager loading, and consistent partitioning for edge reuse.
Josh Joseph, Min Shen, and Jinu Sunil - November 2025
The Bottleneck: Graph Construction in Production
Kumo’s Graph Transformer platform turns relational databases into temporal heterogeneous graphs, then runs transformer-based models over those graphs. Model training and inference get most of the attention, but in production the upstream stages - ingestion, materialization, and graph engine startup - often consume a significant share of the total wall-clock time.
Two usage patterns make this especially painful:
- Iterative development. Data scientists experiment rapidly: modifying which tables to include, adjusting primary keys, tweaking time columns, and changing data types. Each configuration change historically triggered a full pipeline re-run, even when the underlying source data had not changed.
- Repeated production training. Models retrain on fresh data daily or hourly. In most cases only a handful of tables contain new rows, while the majority remain unchanged. Yet the system rebuilt everything from scratch on every run.
The result: teams spent more time waiting for data preparation than for actual model training. Three targeted optimizations, released in Kumo versions 2.11 and 2.12, address this directly. Combined, they deliver 12-25% faster end-to-end training times, averaging around 20% improvement across customer workloads.
In production graph transformer workflows, graph construction, not model training, is often the dominant bottleneck. Optimizing ingestion, materialization, and engine startup yields double-digit end-to-end speedups without touching the model itself.
Anatomy of the Kumo Pipeline
Before examining the optimizations, it helps to understand the stages that every Kumo training run executes. The pipeline has three major phases, each with distinct compute characteristics.
Stage 1: Ingestion
Ingestion copies data from the customer’s source system - Snowflake, Databricks, or S3 - into Kumo’s secure data plane. The goal is to create a stable snapshot that downstream stages can process without contending with live updates. During ingestion, the system also performs data type conversion for each column and data cleaning, including filtering null values in primary keys and timestamps.
Stage 2: Table materialization
Materialization transforms the ingested tabular data into tensor format suitable for graph processing. Each table’s rows are assigned contiguous integer indices, and features are encoded and stored in PyTorch Geometric’s remote backend. This stage converts human-readable database rows into the dense numerical representations the model consumes.
Stage 3: Edge materialization
Foreign key relationships between tables become edges in the graph. Edge materialization performs join operations on primary-foreign key pairs to produce edge lists. These are stored in a remote graph store. For large databases with many foreign key relationships, this stage can be the most expensive part of the pipeline.
Stage 4: Graph engine startup
The graph engine - feature store, graph store, and neighbor sampler - must load all materialized data before training can begin. In the original design, this loading phase waited until all materialization completed, creating idle time between the end of materialization and the start of training.
Kumo training pipeline before optimization
Analogy: Think of the pipeline like preparing ingredients for a restaurant kitchen. Ingestion is sourcing raw ingredients from suppliers. Table materialization is washing, chopping, and portioning each ingredient. Edge materialization is assembling the mise en place, grouping ingredients that go together in each dish. Graph engine startup is arranging everything on the cook’s station. Training is the actual cooking. If the cook cannot start arranging their station until every last ingredient is prepped, service is delayed by the slowest prep task.
Optimization 1: Ingestion Skip
The first optimization targets a specific pain point in iterative development. When a data scientist changes a table’s configuration - adjusting a data type, selecting a different primary key, or modifying the time column - the original system triggered a full re-ingestion of that table from the source system.
This was wasteful. The source data itself had not changed. Only the interpretation of that data - which column is the primary key or what type to assign a column - had changed. Re-copying terabytes from Snowflake or Databricks just because a data scientist tweaked a configuration parameter added minutes or hours of unnecessary wait time.
The fix: decouple configuration from ingestion
The solution separates the concern of “what data do we have?” from “how do we interpret it?” Configuration changes - data type conversions, primary key selection, and time column adjustments - now apply lazily during the materialization stage rather than forcing a data re-copy during ingestion.
Ingestion only re-runs when the source data itself has changed. If a data scientist is experimenting with different graph designs on the same underlying dataset, ingestion is skipped entirely on subsequent runs.
Refresh times are often cut by roughly 50% when experimenting with new graph designs. The optimization is simple in concept - skip unnecessary work - but required careful separation of configuration state from data state throughout the ingestion layer.
Optimization 2: Graph Engine Eager Loading
In the original pipeline, the graph engine - feature store, graph store, and neighbor sampler - only began initializing after all table and edge materialization completed. On a database with 20 tables, the engine sat idle while the last few tables finished materializing, even though the first 18 tables were ready minutes ago.
The fix: progressive warm-up
The feature store and graph store now begin loading as soon as each individual table or edge set finishes materializing. Instead of a synchronization barrier that waits for the slowest table, the engine progressively warms up during the materialization phase.
By the time the last table or edge set completes, the engine has already loaded the vast majority of the graph. The remaining startup time is just the delta: loading the final few artifacts rather than the entire graph from scratch.
Pipeline with eager loading
The practical impact scales with graph size. For small graphs with only a few tables, the engine loads quickly regardless. For medium-to-large graphs with dozens of tables and complex foreign key relationships, eager loading means training starts 30+ minutes earlier than it did before.
Analogy: Instead of waiting until every single ingredient is prepped before arranging the station, the cook starts organizing finished ingredients immediately. By the time the last vegetable is chopped, the station is nearly ready and cooking can begin almost instantly.
Eager loading eliminates the idle gap between materialization and training. On medium-to-large graphs, this alone saves 30+ minutes per training run by overlapping graph engine initialization with the tail end of materialization.
Optimization 3: Consistent Partitioning for Edge Reuse
Edge materialization is one of the most expensive stages in the pipeline. It performs join operations across tables to produce edge lists from foreign key relationships. In repeated training scenarios - retraining on fresh data daily or hourly - much of this work is redundant: if neither the source table nor the target table has changed, the edges between them are identical.
The obvious solution is caching: store previously computed edges and reuse them when the inputs have not changed. But naive caching introduces a subtle and dangerous correctness risk.
The consistency problem
During table materialization, each row is assigned a contiguous integer index.
These indices are what edges reference: edge (3, 7) means “row 3 in the source table connects to row 7 in the target table.”
If the indexing is not perfectly deterministic, the same row could receive index 3 on one run and index 5 on the next.
Cached edges would then point to the wrong rows, producing silently incorrect training data.
In distributed compute environments - Spark, Databricks, and Snowflake - deterministic indexing is not guaranteed by default. Shuffle operations, parallelism settings, and Parquet block sizes can all change row ordering between runs.
The fix: deterministic indexing
Kumo built a deterministic indexing system that enforces consistent partitioning across all supported compute backends. The system locks down:
- Shuffle configurations: fixed partition counts and hash functions so rows land in the same partitions across runs.
- Sort orders: explicit ordering within each partition ensures row sequence is reproducible.
- Parquet block sizes: fixed block boundaries prevent row reordering during serialization and deserialization.
With deterministic indexing in place, the system can safely compare the current run’s inputs against the cached edge data. If the primary keys, time columns, and source data for both tables in a foreign key relationship are unchanged, the cached edges are reused directly.
Consistent partitioning makes edge caching safe across distributed compute environments. By locking down shuffle configurations, sort orders, and Parquet block sizes, the system guarantees that the same input data always produces the same row indices, so cached edges are always correct.
Benchmark Results
Kumo validated the combined optimizations using a synthetic benchmark based on the H&M retail dataset, mirroring common production usage patterns. The benchmark framework automatically logs timing data across pipeline stages via Jenkins integration, with Grafana dashboards planned for continuous monitoring.
Synthetic benchmark (H&M dataset)
The combined effect of all three optimizations on the H&M benchmark:
Production customer results
The optimizations were validated on two production customer workloads with different characteristics:
The variance between 12.5% and 25% is informative. The food delivery company’s workload spent a larger fraction of total runtime on ingestion and materialization, so optimizing those stages had a proportionally bigger impact. The social media platform runs massive-scale graphs where model training itself is the dominant cost. Even there, a 12.5% end-to-end improvement is significant when training runs take hours.
The 12-25% range reflects a real production insight: the speedup depends on the ratio of graph construction time to model training time. Workloads where data preparation dominates see the largest gains. But even on training-dominated workloads, the optimizations deliver meaningful improvements.
Practical Implications and Availability
These three optimizations address the two most common production workflows: iterative experimentation and repeated retraining. The improvements are purely in the data preparation layer. No model architecture changes, accuracy tradeoffs, or changes to the training algorithm are required.
For data scientists iterating on graph design
Ingestion skip is the highest-impact change. Every time you adjust a primary key, change a data type, or modify a time column, the system no longer re-copies data from your warehouse. Experimentation cycles that previously took an hour now complete in roughly half that time. This compounds across a typical development session where a data scientist might try 10-20 configuration variants.
For production teams retraining models
Consistent partitioning and eager loading deliver the bulk of the value. Edge reuse eliminates redundant computation when most tables have not changed between training runs. Eager loading ensures the graph engine is warmed up by the time the last table finishes materializing. Together, these reduce the wall-clock gap between “new data arrives” and “model training begins.”
Availability
- Kumo v2.11: Initial optimizations - ingestion skip and early eager loading.
- Kumo v2.12: Full optimization suite enabled, including consistent partitioning for edge reuse.
All optimizations are enabled by default. No configuration changes are required.
The optimizations are purely in the data preparation layer: no model changes and no accuracy tradeoffs. They are enabled by default in Kumo v2.11 and v2.12, delivering 12-25% faster end-to-end training with zero user-facing configuration changes.