Source-linked AI summary
The Case for TPU-Accelerated Analytical Query Processing
Yeounoh Chung, Jaehyun Ha, Helena Caminal, Fatma Ozcan, Aditya K Kamath, Arvind Krishnamurthy, Henry M. Levy
TL;DR
Analytical database workloads remain largely unexplored on TPUs, whose static execution model poses a challenge for data-dependent relational operations. NTA co-designs TPU-native operators and compilation for this setting. It outperformed DuckDB on all 22 TPC-H queries, with a 7.0× geometric-mean speedup at SF100.
Problem
TPUs’ analytical database use remains largely unexplored, while relational workloads have data-dependent shapes and memory accesses that conflict with static execution.
Source: §1 · p. 1Method
NTA co-designs relational operators for TPU hardware and uses a compiler to bound intermediate capacities and select operator implementations before JIT compilation.
Source: §1 · pp. 1–2Results
7.0× geometric-mean speedup over DuckDB at SF100; NTA outperformed DuckDB on all 22 TPC-H queries.
Source: §1 · p. 2Takeaways & Limitations
The results show that TPUs can deliver competitive, massively scalable analytical query processing on TPC-H.
Source: §10 · p. 12Variable-length strings can waste memory and bandwidth because static tensor representations pad shorter strings to worst-case bounds.
Source: §8.1 · p. 11
Abstract
from arXiv · showhide
Large-scale analytical query processing requires computing substrates capable of immense distributed scale-out. Tensor Processing Units (TPUs) perfectly embody this extreme-scale architecture, offering massive computational density and highly scalable optical pod-level interconnects. Despite these advantages, TPUs have so far remained largely unexplored for analytical database workloads. In this paper, we present the first comprehensive case for TPU-accelerated analytical query processing by running all the TPC-H queries entirely on the TPUs. We introduce Native TPU Accelerator (NTA) for analytical query processing, which provides hardware-conscious algorithm design and optimizations in JAX/XLA. By exposing the TPU's structural mechanics, we demonstrate how dynamic, memory-bound database operations can be rewritten into deterministic JAX programs. We engineer a suite of novel execution strategies that leverage both the TPU's TensorCore and SparseCore to trade sparse, unpredictable memory accesses for dense, hardware-aligned ones. At the core of NTA, a compiler uses catalog statistics to assign guarded static capacities and a hardware-calibrated local cost model to select a TPU-native implementation for each relational operator before JIT compilation. We evaluate our system using the TPC-H benchmark. Taking full advantage of TPU slices and JAX distributed computing, we seamlessly scale our execution environment to run TPC-H benchmark up to Scale Factor 1000 (SF1000). Evaluated against a rigorous CPU ground truth (DuckDB), our native TPU accelerator achieves state-of-the-art performance, achieving a 7.0x single-device speedup over DuckDB at SF100 and a 68.6x speedup with 16 TPUs at SF1000 (1TB). The results establish TPUs not merely as dense neural compute hardware, but as a highly competitive, massively scalable substrate for the next generation of analytical query processing.
1 INTRODUCTION
TPUs offer a highly scalable substrate for analytics, but their static execution model conflicts with relational workloads’ data-dependent shapes and accesses. NTA addresses this mismatch with TPU-native execution and compilation, and outperforms DuckDB across TPC-H at large scale.
1 INTRODUCTION
- TPUs remain largely unexplored for analytical databases, despite their computational density and scalable pod-level interconnects.Relational workloads’ data-dependent shapes and memory accesses conflict with TPUs’ static, compile-time execution requirements.Source: §1 · p. 1
- NTA is a JAX/XLA tensor-based engine that co-designs relational operators for the TPU’s memory and compute hierarchy, including TensorCore and SparseCore.Its strategies trade sparse, padded hash-table accesses for sorting hashed composite keys and validating retrieved candidates against packed keys.Source: §1 · pp. 1–2
- NTA uses catalog statistics to bound intermediate capacities and select operator implementations before lowering a physical plan into a statically shaped JAX program.XLA compiles and further optimizes the program; NTA’s contributions also include TPU-native kernels and distributed execution scaling to SF1000.Source: §1 · p. 2
- 7.0× geometric-mean speedup at SF100 and 68.6× at SF1000 over DuckDB; NTA outperformed DuckDB on all 22 TPC-H queries.The corresponding total-runtime speedups were 5.7× and 50.1×.Source: §1 · p. 2
2 BACKGROUND
TPU query processing must reconcile relational operators’ data-dependent behavior with statically scheduled execution and explicit memory management. The background describes the hardware and software constraints, the gather bottleneck, and the resulting design requirements.
2.1 TPU Architecture
- TensorCores use MXUs and VPUs with software-managed VMEM, while memory transfers require statically specified DMA instructions and compile-time-known sizes and addresses.SparseCores handle irregular gather/scatter streams, but transfers between SparseCore and TensorCore require explicit DMA staging and CMEM synchronization.Source: §2.1 · p. 2
2.2 The Software Stack and Tensor Query Processing
- JAX/XLA favors statically scheduled, fused tensor programs, but a data-dependent gather breaks fusion and forces intermediate materialization in HBM.Directly porting tensor-query mappings designed for GPUs can therefore perform poorly on TPUs.Source: §2.2 · p. 3
2.3 The Gather/Scatter Bottleneck
- At 1.75% selectivity, HBM bank contention overtakes pipeline idling as the primary gather bottleneck.The microbenchmark maps a 60-million-row probe relation to a 15-million-row build relation; clustered, highly selective accesses can funnel queued DMA requests through the most heavily loaded SparseCore lane.Source: Figure 2 · p. 3, §2.3 · p. 3
2.4 Data-Dependent Dynamism on TPUs
- Shape dynamism incurs worst-case padding, while address dynamism breaks fusion and forces HBM materialization; their costs compound across query plans.The resulting design requirements are tightly bounded dynamic shapes, predictable tensor arithmetic for data-dependent accesses where possible, and SparseCore execution for unavoidable dynamic accesses.Source: §2.4 · p. 4
3 NTA SYSTEM OVERVIEW
NTA turns logical query plans into statically shaped TPU programs, using catalog statistics, hardware properties, and calibrated costs to choose implementations before compilation.
3 NTA SYSTEM OVERVIEW
- NTA computes compile-time cardinality bounds for intermediate relations, then selects TPU-native operators using those bounds, key distributions, and hardware cost profiles.Filters and projections are fused into adjacent kernels; the lowered program is compiled with XLA SPMD and distributed across TPU devices.Source: §3 · p. 4
- NTA co-partitions relations across TPU devices, processes local partitions with 32-bit local tuple identifiers, and merges partial results into the final relation.This distributed execution workflow follows static shape planning, physical operator selection, and JAX/Pallas lowering.Source: §3 · p. 4
4 TPU-NATIVE RELATIONAL OPERATORS
NTA represents relations as statically shaped column tensors with validity masks and provides multiple TPU-native implementations for materializing operators. Implementations exploit input properties such as layout, domain size, ordering, and working-set size.
4 TPU-NATIVE RELATIONAL OPERATORS
- NTA represents intermediate relations as statically shaped column tensors with validity masks; selections and projections fuse into the next materializing operator.Joins, aggregations, duplicate elimination, and sorting can each use multiple implementations selected according to input properties and calibrated crossover rules.Source: §4 · p. 4
4.1 Joins
- Join implementations exploit row alignment, bounded PK–FK fanout, or unique build keys; indexed joins use VPU or SparseCore gathers, while sorted-hash lookup handles unique-build joins without a prebuilt index.All implementations preserve the logical join contract, including inner, left-outer, semi, and anti semantics.Source: §4.1 · pp. 4–5
4.2 Aggregations
- Aggregation maps grouping keys to state using dense-domain MXU methods, VMEM histograms, segment reduction, or existing order and clustered layout.The applicable strategy depends on domain and state footprint, input ordering, or bounded fanout; segment reduction is the fallback for large or sparse domains.Source: §4.2 · pp. 5–6
4.3 Distinct
- Distinct uses MXU one-hot membership for bounded integer domains and sort-based methods otherwise, with compacting or VMEM sorting when their conditions apply.Grouped COUNT(DISTINCT x) deduplicates (g, x) pairs before counting survivors for each group.Source: §4.3 · p. 6
4.4 Sort and Top-K
- ORDER BY ... LIMIT K uses in-kernel Top-K; full ordering uses Global Sort, while small statically bounded candidate sets may use Host-Hybrid Sort.Host-Hybrid Sort transfers keys and row identifiers for final ordering, then fetches selected payload columns lazily.Source: §4.4 · p. 6
5 TPU HW-CONSCICOUS COMPILATION & OPTIMIZATION
NTA plans bounded static shapes, selects physical operators using input properties and hardware-calibrated costs, then fuses and lowers the resulting plan for XLA compilation.
5.1 Static Shape Planning
- NTA plans each logical operator's compiled cardinality bound from catalog statistics and semantics, using it to size buffers and account for memory and padded work.Bounds are guarded at runtime: overflow invalidates the result and triggers bound updates, renewed operator selection, and recompilation.Source: §5.1 · pp. 6–7
5.2 Physical Operator Selection
- Physical selection applies input-property and working-set constraints before calibrated hardware crossover rules choose among eligible implementations.The v5p policy switches indexed joins to SparseCore at 50K probe rows; selection thresholds are target-specific, not universal constants.Source: §5.2 · p. 7
5.3 Plan Finalization and Lowering
- Plan finalization fuses non-materializing operations, propagates shapes and validity masks, adds overflow checks, and lowers tensor computation to JAX and custom data movement to Pallas.XLA compiles the statically shaped program, allocates buffers, and schedules execution.Source: §5.3 · p. 7
6 ORCHESTRATION AT EXTREME SCALE
NTA distributes physical plans across a TPU slice and preserves compatible partitioning so many operations can execute locally. Its current FIFO instruction queue limits concurrent-query latency scaling.
6 ORCHESTRATION AT EXTREME SCALE
- NTA uses JAX distributed and XLA GSPMD over a logical device mesh, with each device processing its assigned relation partitions.Large relations are partitioned across devices, while small relations are replicated when they fit in device memory.Source: §6 · p. 8
- Scans, selections, co-partitioned indexed joins, and local aggregations execute without communication; exchanges are needed for repartitioning and combining partial results.Compatible high-volume PK–FK relations are co-partitioned, and partition-local identifiers allow indexed joins to use local partitions.Source: §6 · p. 8
- The serialized FIFO instruction queue makes average query latency scale linearly with concurrent tenants, although aggregate cluster throughput remains constant.Unlike CPU databases, NTA does not time-slice concurrent analytical queries in this deployment.Source: §6 · p. 8
7 EVALUATION
NTA is evaluated on TPC-H against DuckDB and Direct JAX, with warm, precompiled query-execution measurements. It outperforms DuckDB across the reported scales, while operator-level analyses examine hardware mapping and distributed scaling.
7.1 Experimental Setup
- Measurements use TPC-H at SF1, SF100, and SF1000, comparing DuckDB, Direct JAX, and NTA; query execution is timed after warm-up and compilation.The report uses median latency across five runs and reports geometric-mean and total-runtime speedups.Source: §7.1 · pp. 8–9
7.3 Extreme Scale (SF1000)
- At SF1000, row-wise sharding keeps scans, gathers, and local reductions partition-local; the slowest query is still 5.5× faster than DuckDB.State-heavy queries receive less benefit, making the 50.1× total-runtime gain lower than the 68.6× geometric mean.Source: §7.3 · p. 9, Figure 6 · p. 10
7.2 Macrobenchmarks (SF1 & SF100)
- At SF100, NTA accelerates every query by 1.9–30.5×, achieving 7.0× geometric-mean and 5.7× total-runtime speedups over DuckDB.At SF1, it achieves 10.3× geometric-mean and 8.5× total-runtime speedups; Direct JAX is slower than DuckDB on most queries at both scales.Source: §7.2 · p. 9, Figure 5 · p. 10
7.3 Extreme Scale (SF1000)
- 68.6× geometric-mean speedup over DuckDB at SF1000; the 16-chip query set finishes in 1.613 seconds versus 80.86 seconds, a 50.1× total-runtime speedup.NTA accelerates every query by 5.5–250.7×.Source: §7.3 · p. 9, Figure 6 · p. 10
7.4 Ablation Study
- The HBM roofline places most scans, joins, and aggregations in a 0.1–1.0 FLOPS/byte memory-bound region, while many move toward compute-bound operation relative to VMEM.Heavy sorting operators such as Q18 approach the compute-bound region.Source: §7.4 · p. 9, Figure 7 · p. 10
- For PK-FK joins, NTA selects SparseCore Gather, reaching up to 41.4× speedup over DuckDB at SF10; for FK-FK joins, Sorted Hash Lookup Join beats naïve TC Sort-Merge by over 15×.The FK-FK method retains an 18.4× speedup over the CPU baseline at SF10.Source: §7.4 · p. 9, Figure 8 · p. 10
- For 600M-row aggregation, One-Hot Matmul exceeds 1000× speedup at K=4 but runs out of HBM memory at K≥1024; Segment Reduction avoids that matrix and achieves 4.2× speedup.NTA selects between the methods based on group cardinality.Source: §7.4 · pp. 9–10, Figure 9 · p. 10
- The selected join and aggregation kernels exhibit near-linear distributed speedup up to 128 TPU devices.JAX/XLA SPMD shards data row-wise across the devices.Source: Figure 10 · p. 10, §7.4 · p. 10
8 DISCUSSION & FUTURE WORK
NTA’s ClickBench results extend beyond TPC-H, with performance strongest when queries produce compact, bounded, fusable state. The discussion identifies workload and hardware boundaries and proposes co-design directions, including unified execution of relational and semantic AI queries.
8.1 Generalization & Limitations
- 1.79×, 5.17×, and 14.60× geometric-mean speedups over DuckDB at 1, 4, and 16 TPUs, respectively; 16 TPUs win 40 of 43 queries.Four TPUs win 39 queries; a single TPU wins 26 of 43, with a 4.36× geomean on those wins.Source: Figure 12 · p. 11, §8.1 · p. 11
- For ClickBench, NTA takes fixed logical plans from DuckDB, leaving logical rewrites upstream while focusing on physical operator selection and tile scheduling.The evaluation uses a 100M-row flat-table benchmark.Source: §8.1 · p. 11
- ClickBench performance tracks how early a query reduces to compact, statically bounded, fusable state more than its query complexity.Metadata lookups, scalar reductions, and small histograms scale especially well; Q29 reaches up to 417× speedup.Source: §8.1 · p. 11
- Exact distinct counts in Q8 and Q10 require table-scale composite-identity resolution, while Q27’s unbounded URL domain exceeds on-chip VMEM capacity.Q27 scans 100M rows, forcing histogram updates into latency-heavy HBM round-trips.Source: §8.1 · p. 11
- Variable-length strings and selective matching such as LIKE or regex remain challenging because static worst-case padding wastes memory, bandwidth, and vector compute.When average strings are much shorter than the maximum, padding consumes scarce HBM bandwidth and VMEM.Source: §8.1 · p. 11
8.2 HW & SW Co-Design Opportunities
- Proposed hardware directions include skipping invalid padded tiles and compacting sparse vectors near VPU units to reduce redundant work.The discussion identifies shape dynamism as a source of padded rows and masked SIMD lanes.Source: §8.2 · pp. 11–12
- Kernel-managed scratchpad placement could keep hot data fast, while hardware partitioning and isolation could support concurrent query execution.The current lack of hardware partitioning enforces strict sequential execution of query kernels.Source: §8.2 · p. 12
8.3 End-to-End Semantic Queries
- NTA could compile relational operators and LLM invocations into one JAX program running natively on a TPU.The discussion presents this unification as a promising application for AI-native database engines.Source: §8.3 · p. 12
9 RELATED WORK
Prior accelerator query work has centered on GPUs, while TPU studies often treated them as GPU substitutes. NTA instead designs operators around TPU-specific execution constraints, including a join strategy that avoids sparse hash materialization.
Hardware-Accelerated Query Processing
- GPU query frameworks and earlier TPU studies leave a gap: NTA targets TPU-specific HBM stalls and fusion breaks with custom Pallas kernels for TensorCore and SparseCore.The cited TPU studies [14] [16] largely treated TPUs as black-box GPU substitutes.Source: §9 · p. 12
Join Algorithms on Accelerators
- Sorted-Hash Lookup Join replaces sparse hash materialization and erratic probing with packed-key hashing, parallel sorting, and vectorized search.The passage contrasts this with accelerator hash-join methods relying on dynamic allocation and fine-grained scatter/gather.Source: §9 · p. 12
Other TPU Applications
- Unlike predictable tensor workloads such as nearest-neighbor search [4], analytical queries have data-dependent output sizes and access patterns; NTA reports execution up to SF1000.The passage describes NTA as the first engine to execute strict relational algebra effectively across multi-terabyte datasets.Source: §9 · p. 12
10 CONCLUSION
The conclusion argues that TPUs remain underexplored for analytical databases and presents NTA as a hardware-conscious JAX/XLA system with competitive TPC-H performance. It also identifies shared execution of relational and semantic AI operations as an opportunity for accelerator-based analytics.
10 CONCLUSION
- NTA demonstrates competitive TPU-accelerated analytical query processing, including state-of-the-art performance on TPC-H benchmarks.The system uses hardware-conscious algorithm design and JAX/XLA optimizations that leverage SparseCore and TPU slices.Source: §10 · p. 12
- Running relational queries and semantic AI operations on a shared computing substrate is presented as an opportunity to unify enterprise data management with ML/AI accelerator stacks.The conclusion frames this as an opportunity for richer structured and unstructured data analytics.Source: §10 · p. 12