ifc_lite_server/services/parquet_stream_shapes.rs
1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4
5//! Whole-stream offsets for the streaming Parquet route, and the cross-batch
6//! shape sharing that rides on them (#5407).
7//!
8//! A streamed model is written one batch at a time, so anything that has to
9//! hold across batches lives here: the running vertex/index offsets every
10//! batch's mesh table is based on, and, when the client opted in, where each
11//! distinct shape already landed. The cache writer and the per-batch client
12//! blob both consume a [`PlannedBatch`], so the two can never disagree about a
13//! range.
14//!
15//! BOUNDED MEMORY. No occurrence outlives the batch that produced it. What
16//! grows across batches is, per DISTINCT shape emitted so far, one
17//! [`EmittedShape`] (16 bytes) plus its `MeshGeometryKey` (24 bytes), whatever
18//! the size of that shape's geometry; and, per instanced representation, the
19//! ONE template mesh its later occurrences are verified against, capped in
20//! total at [`RETAINED_TEMPLATE_BUDGET_BYTES`]. Past the cap a representation
21//! is simply not carried: its later occurrences are shared within their own
22//! batch, or by content hash, instead. Sharing degrades; correctness and the
23//! bound do not.
24
25use crate::services::parquet::{frame_sections, write_parquet_buffer, ParquetError};
26use crate::services::parquet_layout::{ParquetLayout, StreamShapes};
27use crate::services::parquet_mesh_tables::build_mesh_tables;
28use crate::services::parquet_shape_plan::{EmittedShape, PriorShapes, RetainedTemplate, ShapePlan};
29use crate::types::MeshData;
30use arrow::record_batch::RecordBatch;
31use bytes::Bytes;
32
33/// Ceiling on the template geometry a cross-batch stream retains (#5407).
34///
35/// Sized against what it buys, not what it costs a small model: the
36/// representations that repeat across batches are the model's distinct
37/// instanced shapes, whose geometry is what the shared stream sends ONCE
38/// anyway. 256 MiB of raw template buffers is several times that on every
39/// fixture measured for #5407, while staying well under one parse's working
40/// set on the models a streaming deployment exists for.
41pub(super) const RETAINED_TEMPLATE_BUDGET_BYTES: usize = 256 * 1024 * 1024;
42
43/// One batch's mesh / vertex / index tables, planned against the whole stream.
44pub struct PlannedBatch {
45 pub(super) mesh: RecordBatch,
46 pub(super) vertex: RecordBatch,
47 pub(super) index: RecordBatch,
48 /// Whole-stream offset of this batch's first vertex row.
49 pub vertex_base: u32,
50 /// Whole-stream offset of this batch's first index, in indices (three per
51 /// index-table row), the unit of the mesh table's `index_start`.
52 pub index_base: u32,
53}
54
55impl PlannedBatch {
56 /// The batch as a client blob, in the same `[len][mesh][len][vert][len]
57 /// [idx]` framing every batch event carries. Its offsets are whole-stream,
58 /// so it is only meaningful to a client told `vertex_base` / `index_base`.
59 pub fn encode(&self) -> Result<Bytes, ParquetError> {
60 frame_sections(
61 &write_parquet_buffer(&self.mesh)?,
62 &write_parquet_buffer(&self.vertex)?,
63 &write_parquet_buffer(&self.index)?,
64 )
65 }
66
67 /// Mesh rows in this batch.
68 pub fn mesh_count(&self) -> usize {
69 self.mesh.num_rows()
70 }
71}
72
73/// The running state behind a streamed model's whole-stream offsets.
74pub struct StreamShapePlanner {
75 layout: ParquetLayout,
76 /// `None`: every mesh emits its own block, as the stream always did.
77 /// `Some`: cross-batch sharing, against every shape emitted so far.
78 prior: Option<PriorShapes>,
79 /// Raw buffer bytes held by `prior`'s retained templates, and their cap.
80 retained_bytes: usize,
81 template_budget: usize,
82 vertex_offset: u32,
83 index_offset: u32,
84}
85
86impl StreamShapePlanner {
87 /// A planner for `layout`. Sharing needs BOTH the shared-shapes layout and
88 /// the client's cross-batch opt-in; the route rejects the opt-in on the
89 /// flat layout, so ignoring it here is never what a caller sees.
90 pub fn new(layout: ParquetLayout, stream_shapes: StreamShapes) -> Self {
91 let sharing = layout.has_rotation() && stream_shapes == StreamShapes::CrossBatch;
92 Self {
93 layout,
94 prior: sharing.then(PriorShapes::default),
95 retained_bytes: 0,
96 template_budget: RETAINED_TEMPLATE_BUDGET_BYTES,
97 vertex_offset: 0,
98 index_offset: 0,
99 }
100 }
101
102 /// Plan the next batch and advance the whole-stream offsets past the
103 /// shapes it emits. Fails loud, advancing nothing, when an offset would
104 /// overflow the mesh table's `u32` columns: wrapping would decode as
105 /// garbage, not as an error.
106 ///
107 /// `baked_basis` is the row-major `native_to_baked` of the frame the
108 /// batch's vertices are in (`StreamEvent::Batch::baked_basis`); ignored
109 /// unless the planner shares.
110 pub fn plan(
111 &mut self,
112 meshes: &[MeshData],
113 baked_basis: Option<&[f64; 16]>,
114 ) -> Result<PlannedBatch, ParquetError> {
115 let basis = baked_basis.map(|m| ifc_lite_geometry::Matrix4::from_row_slice(m));
116 let (plan, new_shapes) = match &self.prior {
117 Some(prior) => ShapePlan::shared_shapes_after(meshes, basis.as_ref(), prior),
118 None => (ShapePlan::Identity, Vec::new()),
119 };
120 let (vertex_base, index_base) = (self.vertex_offset, self.index_offset);
121 let mut emitted = Vec::with_capacity(new_shapes.len());
122 let (mut vertex_end, mut index_end) = (vertex_base, index_base);
123 for shape in plan.shape_meshes(meshes) {
124 let vertex_count = u32::try_from(shape.positions.len() / 3).ok();
125 let index_count = u32::try_from(shape.indices.len()).ok();
126 let (Some(vertex_count), Some(index_count)) = (vertex_count, index_count) else {
127 return Err(overflow());
128 };
129 if self.prior.is_some() {
130 emitted.push(EmittedShape {
131 vertex_start: vertex_end,
132 vertex_count,
133 index_start: index_end,
134 index_count,
135 });
136 }
137 vertex_end = vertex_end.checked_add(vertex_count).ok_or_else(overflow)?;
138 index_end = index_end.checked_add(index_count).ok_or_else(overflow)?;
139 }
140 let (mesh, vertex, index) =
141 build_mesh_tables(meshes, &plan, self.layout, vertex_base, index_base)?;
142 if let Some(prior) = self.prior.as_mut() {
143 // `new_shapes` runs parallel to the emitted shapes (see
144 // `ShapePlan::shared_shapes_after`), so each lands where it was
145 // just planned.
146 let shape_meshes = plan.shape_meshes(meshes);
147 for ((shape, at), mesh) in new_shapes.into_iter().zip(emitted).zip(shape_meshes) {
148 prior.by_hash.insert(shape.key, at);
149 let Some(rep) = shape.template_of else { continue };
150 let bytes = template_bytes(mesh);
151 if prior.templates.contains_key(&rep)
152 || self.retained_bytes + bytes > self.template_budget
153 {
154 continue;
155 }
156 self.retained_bytes += bytes;
157 prior.templates.insert(rep, RetainedTemplate { mesh: mesh.clone(), at });
158 }
159 }
160 self.vertex_offset = vertex_end;
161 self.index_offset = index_end;
162 Ok(PlannedBatch {
163 mesh,
164 vertex,
165 index,
166 vertex_base,
167 index_base,
168 })
169 }
170}
171
172/// The buffers a retained template holds, in bytes.
173fn template_bytes(mesh: &MeshData) -> usize {
174 (mesh.positions.len() + mesh.normals.len() + mesh.indices.len()) * 4
175}
176
177fn overflow() -> ParquetError {
178 ParquetError::Overflow("global vertex/index offsets exceed u32".to_string())
179}
180
181#[cfg(test)]
182#[path = "parquet_stream_shapes_tests.rs"]
183mod parquet_stream_shapes_tests;