1use crate::services::axis::{zup_to_yup, zup_to_yup_f64};
15use crate::services::parquet_schema::{instance_schema, ABSENT_SOURCE_ID};
16use crate::types::MeshData;
17use arrow::array::{Float32Array, Float64Array, Int32Array, StringArray, UInt32Array, UInt8Array};
18use arrow::datatypes::{DataType, Field, Schema};
19use arrow::record_batch::RecordBatch;
20use bytes::Bytes;
21use parquet::arrow::ArrowWriter;
22use parquet::basic::Compression;
23use parquet::file::properties::WriterProperties;
24use parquet::schema::types::ColumnPath;
25use rustc_hash::FxHashMap;
26use std::io::Cursor;
27use std::sync::Arc;
28
29use super::parquet::check_u32_len;
30use super::ParquetError;
31
32use crate::services::parquet_instancing::{
33 collate_rotation_aware_placements, mesh_geometry_key, optimized_wire_version,
34 rotation_zup_to_yup, MeshGeometryKey, IDENTITY_ROTATION,
35};
36
37pub const VERTEX_MULTIPLIER: f32 = 10_000.0;
39
40#[inline]
42fn quantize_position(value: f32) -> i32 {
43 (value * VERTEX_MULTIPLIER).round() as i32
44}
45
46#[inline]
48fn color_to_byte(value: f32) -> u8 {
49 (value.clamp(0.0, 1.0) * 255.0).round() as u8
50}
51
52#[derive(Clone, Copy, PartialEq, Eq, Hash)]
54struct MaterialKey {
55 r: u8,
56 g: u8,
57 b: u8,
58 a: u8,
59}
60
61impl MaterialKey {
62 fn from_color(color: &[f32; 4]) -> Self {
63 Self {
64 r: color_to_byte(color[0]),
65 g: color_to_byte(color[1]),
66 b: color_to_byte(color[2]),
67 a: color_to_byte(color[3]),
68 }
69 }
70}
71
72fn serialize_to_parquet_optimized(
87 meshes: &[MeshData],
88 include_normals: bool,
89 baked_basis: Option<&ifc_lite_geometry::Matrix4<f64>>,
90) -> Result<(Bytes, usize, usize), ParquetError> {
91 let rotated_placements = collate_rotation_aware_placements(meshes, baked_basis);
95
96 let mut unique_meshes: Vec<&MeshData> = Vec::new();
98 let mut mesh_lookup: FxHashMap<MeshGeometryKey, u32> = FxHashMap::default();
99 let mut template_lookup: FxHashMap<usize, u32> = FxHashMap::default();
103 let mut unique_materials: Vec<MaterialKey> = Vec::new();
104 let mut material_lookup: FxHashMap<MaterialKey, u32> = FxHashMap::default();
105
106 let mut instance_entity_ids: Vec<u32> = Vec::with_capacity(meshes.len());
108 let mut instance_ifc_types: Vec<&str> = Vec::with_capacity(meshes.len());
109 let mut instance_mesh_indices: Vec<u32> = Vec::with_capacity(meshes.len());
110 let mut instance_material_indices: Vec<u32> = Vec::with_capacity(meshes.len());
111 let mut instance_origin_x: Vec<f64> = Vec::with_capacity(meshes.len());
119 let mut instance_origin_y: Vec<f64> = Vec::with_capacity(meshes.len());
120 let mut instance_origin_z: Vec<f64> = Vec::with_capacity(meshes.len());
121 let mut instance_geometry_class: Vec<u8> = Vec::with_capacity(meshes.len());
122 let mut instance_geometry_item_ids: Vec<u32> = Vec::with_capacity(meshes.len());
125 let mut instance_material_ids: Vec<u32> = Vec::with_capacity(meshes.len());
126 let mut instance_finishes: [Vec<f32>; 2] = Default::default();
129 let mut instance_rotation: [Vec<f32>; 9] = Default::default();
132 let mut emitted_non_identity_rotation = false;
134
135 for (mesh_index, mesh) in meshes.iter().enumerate() {
136 let (mesh_idx, origin_yup, rotation_yup) = match rotated_placements.get(&mesh_index) {
137 Some(placement) => {
138 let template = &meshes[placement.template_mesh_index];
139 let idx = *template_lookup
140 .entry(placement.template_mesh_index)
141 .or_insert_with(|| {
142 let idx = unique_meshes.len() as u32;
143 unique_meshes.push(template);
144 idx
145 });
146 (idx, zup_to_yup_f64(placement.origin_zup), rotation_zup_to_yup(&placement.rotation_zup))
147 }
148 None => {
149 let geo_key = mesh_geometry_key(mesh);
150 let idx = *mesh_lookup.entry(geo_key).or_insert_with(|| {
151 let idx = unique_meshes.len() as u32;
152 unique_meshes.push(mesh);
153 idx
154 });
155 (idx, zup_to_yup_f64(mesh.origin), IDENTITY_ROTATION)
156 }
157 };
158
159 let mat_key = MaterialKey::from_color(&mesh.color);
161 let material_idx = *material_lookup.entry(mat_key).or_insert_with(|| {
162 let idx = unique_materials.len() as u32;
163 unique_materials.push(mat_key);
164 idx
165 });
166
167 instance_entity_ids.push(mesh.express_id);
171 instance_ifc_types.push(&mesh.ifc_type);
172 instance_mesh_indices.push(mesh_idx);
173 instance_material_indices.push(material_idx);
174 instance_origin_x.push(origin_yup[0]);
175 instance_origin_y.push(origin_yup[1]);
176 instance_origin_z.push(origin_yup[2]);
177 instance_geometry_class.push(mesh.geometry_class);
178 instance_geometry_item_ids.push(mesh.geometry_item_id.unwrap_or(ABSENT_SOURCE_ID));
179 instance_material_ids.push(mesh.material_id.unwrap_or(ABSENT_SOURCE_ID));
180 for (col, value) in instance_finishes.iter_mut().zip(mesh.finish_wire()) {
181 col.push(value);
182 }
183 emitted_non_identity_rotation |= rotation_yup != IDENTITY_ROTATION;
184 for (col, value) in instance_rotation.iter_mut().zip(rotation_yup.iter()) {
185 col.push(*value);
186 }
187 }
188
189 let total_vertices: usize = unique_meshes.iter().map(|m| m.positions.len() / 3).sum();
191 let total_indices: usize = unique_meshes.iter().map(|m| m.indices.len()).sum();
192
193 let mut vertex_x: Vec<i32> = Vec::with_capacity(total_vertices);
195 let mut vertex_y: Vec<i32> = Vec::with_capacity(total_vertices);
196 let mut vertex_z: Vec<i32> = Vec::with_capacity(total_vertices);
197
198 let normals_capacity = if include_normals { total_vertices } else { 0 };
200 let mut normal_x: Vec<f32> = Vec::with_capacity(normals_capacity);
201 let mut normal_y: Vec<f32> = Vec::with_capacity(normals_capacity);
202 let mut normal_z: Vec<f32> = Vec::with_capacity(normals_capacity);
203
204 let mut indices: Vec<u32> = Vec::with_capacity(total_indices);
206
207 let mut mesh_vertex_offsets: Vec<u32> = Vec::with_capacity(unique_meshes.len());
209 let mut mesh_vertex_counts: Vec<u32> = Vec::with_capacity(unique_meshes.len());
210 let mut mesh_index_offsets: Vec<u32> = Vec::with_capacity(unique_meshes.len());
211 let mut mesh_index_counts: Vec<u32> = Vec::with_capacity(unique_meshes.len());
212
213 let mut vertex_offset: u32 = 0;
214 let mut index_offset: u32 = 0;
215
216 for mesh in &unique_meshes {
217 let vert_count = mesh.positions.len() / 3;
218
219 mesh_vertex_offsets.push(vertex_offset);
220 mesh_vertex_counts.push(vert_count as u32);
221 mesh_index_offsets.push(index_offset);
222 mesh_index_counts.push(mesh.indices.len() as u32);
223
224 let mesh_has_normals = mesh.normals.len() == mesh.positions.len();
227 if include_normals && !mesh_has_normals && !mesh.normals.is_empty() {
228 tracing::warn!(
229 express_id = mesh.express_id,
230 ifc_type = %mesh.ifc_type,
231 positions = mesh.positions.len(),
232 normals = mesh.normals.len(),
233 "Mesh normals length mismatch; emitting zero normals"
234 );
235 }
236
237 for i in 0..vert_count {
240 let (x, y, z) = zup_to_yup(
241 mesh.positions[i * 3],
242 mesh.positions[i * 3 + 1],
243 mesh.positions[i * 3 + 2],
244 );
245 vertex_x.push(quantize_position(x));
246 vertex_y.push(quantize_position(y));
247 vertex_z.push(quantize_position(z));
248
249 if include_normals {
250 if mesh_has_normals {
251 let (x, y, z) = zup_to_yup(
252 mesh.normals[i * 3],
253 mesh.normals[i * 3 + 1],
254 mesh.normals[i * 3 + 2],
255 );
256 normal_x.push(x);
257 normal_y.push(y);
258 normal_z.push(z);
259 } else {
260 normal_x.push(0.0);
261 normal_y.push(0.0);
262 normal_z.push(0.0);
263 }
264 }
265 }
266
267 indices.extend_from_slice(&mesh.indices);
269
270 vertex_offset += vert_count as u32;
271 index_offset += mesh.indices.len() as u32;
272 }
273
274 let rotation_column_count = if emitted_non_identity_rotation { 9 } else { 0 };
278 let instance_schema = instance_schema(emitted_non_identity_rotation);
279
280 let instance_columns: Vec<Arc<dyn arrow::array::Array>> = vec![
281 Arc::new(UInt32Array::from(instance_entity_ids)) as Arc<dyn arrow::array::Array>,
282 Arc::new(StringArray::from(instance_ifc_types)),
283 Arc::new(UInt32Array::from(instance_mesh_indices)),
284 Arc::new(UInt32Array::from(instance_material_indices)),
285 Arc::new(Float64Array::from(instance_origin_x)),
286 Arc::new(Float64Array::from(instance_origin_y)),
287 Arc::new(Float64Array::from(instance_origin_z)),
288 Arc::new(UInt8Array::from(instance_geometry_class)),
289 Arc::new(UInt32Array::from(instance_geometry_item_ids)),
290 Arc::new(UInt32Array::from(instance_material_ids)),
291 ]
292 .into_iter()
293 .chain(instance_finishes.into_iter().map(|col| Arc::new(Float32Array::from(col)) as Arc<dyn arrow::array::Array>))
294 .chain(instance_rotation.into_iter().take(rotation_column_count).map(|col| Arc::new(Float32Array::from(col)) as Arc<dyn arrow::array::Array>))
295 .collect();
296
297 let instance_batch = RecordBatch::try_new(instance_schema, instance_columns)?;
298
299 let mesh_schema = Arc::new(Schema::new(vec![
301 Field::new("vertex_offset", DataType::UInt32, false),
302 Field::new("vertex_count", DataType::UInt32, false),
303 Field::new("index_offset", DataType::UInt32, false),
304 Field::new("index_count", DataType::UInt32, false),
305 ]));
306
307 let mesh_batch = RecordBatch::try_new(
308 mesh_schema,
309 vec![
310 Arc::new(UInt32Array::from(mesh_vertex_offsets)),
311 Arc::new(UInt32Array::from(mesh_vertex_counts)),
312 Arc::new(UInt32Array::from(mesh_index_offsets)),
313 Arc::new(UInt32Array::from(mesh_index_counts)),
314 ],
315 )?;
316
317 let material_schema = Arc::new(Schema::new(vec![
319 Field::new("r", DataType::UInt8, false),
320 Field::new("g", DataType::UInt8, false),
321 Field::new("b", DataType::UInt8, false),
322 Field::new("a", DataType::UInt8, false),
323 ]));
324
325 let material_batch = RecordBatch::try_new(
326 material_schema,
327 vec![
328 Arc::new(UInt8Array::from(unique_materials.iter().map(|m| m.r).collect::<Vec<_>>())),
329 Arc::new(UInt8Array::from(unique_materials.iter().map(|m| m.g).collect::<Vec<_>>())),
330 Arc::new(UInt8Array::from(unique_materials.iter().map(|m| m.b).collect::<Vec<_>>())),
331 Arc::new(UInt8Array::from(unique_materials.iter().map(|m| m.a).collect::<Vec<_>>())),
332 ],
333 )?;
334
335 let vertex_schema = if include_normals {
337 Arc::new(Schema::new(vec![
338 Field::new("x", DataType::Int32, false),
339 Field::new("y", DataType::Int32, false),
340 Field::new("z", DataType::Int32, false),
341 Field::new("nx", DataType::Float32, false),
342 Field::new("ny", DataType::Float32, false),
343 Field::new("nz", DataType::Float32, false),
344 ]))
345 } else {
346 Arc::new(Schema::new(vec![
347 Field::new("x", DataType::Int32, false),
348 Field::new("y", DataType::Int32, false),
349 Field::new("z", DataType::Int32, false),
350 ]))
351 };
352
353 let vertex_batch = if include_normals {
354 RecordBatch::try_new(
355 vertex_schema,
356 vec![
357 Arc::new(Int32Array::from(vertex_x)),
358 Arc::new(Int32Array::from(vertex_y)),
359 Arc::new(Int32Array::from(vertex_z)),
360 Arc::new(Float32Array::from(normal_x)),
361 Arc::new(Float32Array::from(normal_y)),
362 Arc::new(Float32Array::from(normal_z)),
363 ],
364 )?
365 } else {
366 RecordBatch::try_new(
367 vertex_schema,
368 vec![
369 Arc::new(Int32Array::from(vertex_x)),
370 Arc::new(Int32Array::from(vertex_y)),
371 Arc::new(Int32Array::from(vertex_z)),
372 ],
373 )?
374 };
375
376 let index_schema = Arc::new(Schema::new(vec![Field::new("i", DataType::UInt32, false)]));
378
379 let index_batch =
380 RecordBatch::try_new(index_schema, vec![Arc::new(UInt32Array::from(indices))])?;
381
382 let instance_parquet = write_parquet_buffer(&instance_batch)?;
384 let mesh_parquet = write_parquet_buffer(&mesh_batch)?;
385 let material_parquet = write_parquet_buffer(&material_batch)?;
386 let vertex_parquet = write_parquet_buffer(&vertex_batch)?;
387 let index_parquet = write_parquet_buffer(&index_batch)?;
388
389 let data = assemble_optimized_output(
390 optimized_wire_version(emitted_non_identity_rotation),
391 include_normals,
392 &instance_parquet,
393 &mesh_parquet,
394 &material_parquet,
395 &vertex_parquet,
396 &index_parquet,
397 )?;
398 Ok((data, unique_meshes.len(), unique_materials.len()))
399}
400
401fn check_optimized_section_lengths(
405 instance_len: usize, mesh_len: usize, material_len: usize, vertex_len: usize, index_len: usize,
406) -> Result<(), ParquetError> {
407 let sections = [
408 ("instance", instance_len), ("mesh", mesh_len), ("material", material_len),
409 ("vertex", vertex_len), ("index", index_len),
410 ];
411 for (name, len) in sections {
412 check_u32_len(name, len)?;
413 }
414 Ok(())
415}
416
417fn assemble_optimized_output(
426 version: u8,
427 include_normals: bool,
428 instance_parquet: &[u8],
429 mesh_parquet: &[u8],
430 material_parquet: &[u8],
431 vertex_parquet: &[u8],
432 index_parquet: &[u8],
433) -> Result<Bytes, ParquetError> {
434 check_optimized_section_lengths(
435 instance_parquet.len(), mesh_parquet.len(), material_parquet.len(),
436 vertex_parquet.len(), index_parquet.len(),
437 )?;
438
439 let mut output = Vec::with_capacity(
440 2 + 20 + instance_parquet.len() + mesh_parquet.len() + material_parquet.len()
441 + vertex_parquet.len() + index_parquet.len(),
442 );
443
444 output.push(version);
445 output.push(if include_normals { 1u8 } else { 0u8 });
447
448 output.extend_from_slice(&(instance_parquet.len() as u32).to_le_bytes());
449 output.extend_from_slice(&(mesh_parquet.len() as u32).to_le_bytes());
450 output.extend_from_slice(&(material_parquet.len() as u32).to_le_bytes());
451 output.extend_from_slice(&(vertex_parquet.len() as u32).to_le_bytes());
452 output.extend_from_slice(&(index_parquet.len() as u32).to_le_bytes());
453
454 output.extend_from_slice(instance_parquet);
455 output.extend_from_slice(mesh_parquet);
456 output.extend_from_slice(material_parquet);
457 output.extend_from_slice(vertex_parquet);
458 output.extend_from_slice(index_parquet);
459
460 Ok(Bytes::from(output))
461}
462
463fn write_parquet_buffer(batch: &RecordBatch) -> Result<Vec<u8>, ParquetError> {
467 let mut buffer = Vec::new();
468 let cursor = Cursor::new(&mut buffer);
469
470 let mut props_builder = WriterProperties::builder()
472 .set_compression(Compression::LZ4_RAW)
473 .set_dictionary_enabled(true); for field in batch.schema().fields() {
478 let is_numeric = matches!(
479 field.data_type(),
480 DataType::Float32
481 | DataType::Float64
482 | DataType::UInt32
483 | DataType::UInt64
484 | DataType::Int32
485 | DataType::Int64
486 | DataType::UInt8
487 | DataType::Int8
488 );
489
490 if is_numeric {
491 props_builder = props_builder
492 .set_column_dictionary_enabled(ColumnPath::from(field.name().as_str()), false);
493 }
494 }
495
496 let props = props_builder.build();
497
498 let mut writer = ArrowWriter::try_new(cursor, batch.schema(), Some(props))?;
499 writer.write(batch)?;
500 writer.close()?;
501
502 Ok(buffer)
503}
504
505#[derive(Debug, Clone, serde::Serialize)]
507pub struct OptimizedStats {
508 pub input_meshes: usize,
510 pub unique_meshes: usize,
512 pub unique_materials: usize,
514 pub mesh_reuse_ratio: f32,
516 pub has_normals: bool,
518}
519
520pub fn serialize_to_parquet_optimized_with_stats(
522 meshes: &[MeshData],
523 include_normals: bool,
524 baked_basis: Option<&ifc_lite_geometry::Matrix4<f64>>,
525) -> Result<(Bytes, OptimizedStats), ParquetError> {
526 let (data, unique_mesh_count, unique_material_count) =
527 serialize_to_parquet_optimized(meshes, include_normals, baked_basis)?;
528
529 let stats = OptimizedStats {
530 input_meshes: meshes.len(),
531 unique_meshes: unique_mesh_count,
532 unique_materials: unique_material_count,
533 mesh_reuse_ratio: if unique_mesh_count > 0 {
534 meshes.len() as f32 / unique_mesh_count as f32
535 } else {
536 1.0
537 },
538 has_normals: include_normals,
539 };
540
541 Ok((data, stats))
542}
543
544#[cfg(test)]
545#[path = "parquet_optimized_tests.rs"]
546mod optimized_tests;