1use crate::services::data_model::{
8 ClassificationAssociation, DataModel, DocumentAssociation, EntityMetadata,
9 PropertySet, QuantitySet, Relationship, SpatialHierarchyData, SpatialNode,
10};
11use arrow::array::builder::ListBuilder;
12use arrow::array::UInt32Builder;
13use arrow::array::{BooleanArray, StringArray, UInt16Array, UInt32Array};
14use arrow::datatypes::{DataType, Field, Schema};
15use arrow::record_batch::RecordBatch;
16use parquet::arrow::ArrowWriter;
17use parquet::basic::Compression;
18use parquet::file::properties::WriterProperties;
19use rayon::prelude::*;
20use std::io::Cursor;
21use std::sync::Arc;
22use thiserror::Error;
23
24#[path = "parquet_data_model_materials.rs"]
25mod materials;
26use materials::serialize_materials_table;
27
28#[derive(Debug, Error)]
30pub enum DataModelParquetError {
31 #[error("Arrow error: {0}")]
32 Arrow(#[from] arrow::error::ArrowError),
33 #[error("Parquet error: {0}")]
34 Parquet(#[from] parquet::errors::ParquetError),
35 #[error("IO error: {0}")]
36 Io(#[from] std::io::Error),
37 #[error("JSON error: {0}")]
38 Json(#[from] serde_json::Error),
39}
40
41pub fn serialize_data_model_to_parquet(
51 data_model: &DataModel,
52) -> Result<Vec<u8>, DataModelParquetError> {
53 let (entities_data, ((properties_data, quantities_data), (relationships_data, spatial_data))) =
55 rayon::join(
56 || serialize_entities_table(&data_model.entities),
57 || {
58 rayon::join(
59 || {
60 rayon::join(
61 || serialize_properties_table(&data_model.property_sets),
62 || serialize_quantities_table(&data_model.quantity_sets),
63 )
64 },
65 || {
66 rayon::join(
67 || serialize_relationships_table(&data_model.relationships),
68 || serialize_spatial_hierarchy(&data_model.spatial_hierarchy),
69 )
70 },
71 )
72 },
73 );
74
75 let entities_data = entities_data?;
76 let properties_data = properties_data?;
77 let quantities_data = quantities_data?;
78 let relationships_data = relationships_data?;
79 let spatial_data = spatial_data?;
80
81 let ((classifications_data, materials_data), documents_data) = rayon::join(
85 || {
86 rayon::join(
87 || serialize_classifications_table(&data_model.classifications),
88 || serialize_materials_table(&data_model.materials),
89 )
90 },
91 || serialize_documents_table(&data_model.documents),
92 );
93 let classifications_data = classifications_data?;
94 let materials_data = materials_data?;
95 let documents_data = documents_data?;
96
97 let mut result = Vec::new();
99 result.extend_from_slice(&(entities_data.len() as u32).to_le_bytes());
100 result.extend_from_slice(&entities_data);
101 result.extend_from_slice(&(properties_data.len() as u32).to_le_bytes());
102 result.extend_from_slice(&properties_data);
103 result.extend_from_slice(&(quantities_data.len() as u32).to_le_bytes());
104 result.extend_from_slice(&quantities_data);
105 result.extend_from_slice(&(relationships_data.len() as u32).to_le_bytes());
106 result.extend_from_slice(&relationships_data);
107 result.extend_from_slice(&(spatial_data.len() as u32).to_le_bytes());
108 result.extend_from_slice(&spatial_data);
109 result.extend_from_slice(&(classifications_data.len() as u32).to_le_bytes());
110 result.extend_from_slice(&classifications_data);
111 result.extend_from_slice(&(materials_data.len() as u32).to_le_bytes());
112 result.extend_from_slice(&materials_data);
113 result.extend_from_slice(&(documents_data.len() as u32).to_le_bytes());
114 result.extend_from_slice(&documents_data);
115
116 Ok(result)
117}
118
119fn serialize_classifications_table(
121 rows: &[ClassificationAssociation],
122) -> Result<Vec<u8>, DataModelParquetError> {
123 let count = rows.len();
124 let mut element_ids = Vec::with_capacity(count);
125 let mut system_names: Vec<Option<String>> = Vec::with_capacity(count);
126 let mut identifications: Vec<Option<String>> = Vec::with_capacity(count);
127 let mut names: Vec<Option<String>> = Vec::with_capacity(count);
128 let mut locations: Vec<Option<String>> = Vec::with_capacity(count);
129
130 for row in rows {
131 element_ids.push(row.element_id);
132 system_names.push(row.system_name.clone());
133 identifications.push(row.identification.clone());
134 names.push(row.name.clone());
135 locations.push(row.location.clone());
136 }
137
138 let schema = Schema::new(vec![
139 Field::new("element_id", DataType::UInt32, false),
140 Field::new("system_name", DataType::Utf8, true),
141 Field::new("identification", DataType::Utf8, true),
142 Field::new("name", DataType::Utf8, true),
143 Field::new("location", DataType::Utf8, true),
144 ]);
145
146 let batch = RecordBatch::try_new(
147 Arc::new(schema),
148 vec![
149 Arc::new(UInt32Array::from(element_ids)),
150 Arc::new(StringArray::from(system_names)),
151 Arc::new(StringArray::from(identifications)),
152 Arc::new(StringArray::from(names)),
153 Arc::new(StringArray::from(locations)),
154 ],
155 )?;
156
157 write_parquet_batch(batch)
158}
159
160fn serialize_documents_table(
162 rows: &[DocumentAssociation],
163) -> Result<Vec<u8>, DataModelParquetError> {
164 let count = rows.len();
165 let mut element_ids = Vec::with_capacity(count);
166 let mut identifications: Vec<Option<String>> = Vec::with_capacity(count);
167 let mut names: Vec<Option<String>> = Vec::with_capacity(count);
168 let mut locations: Vec<Option<String>> = Vec::with_capacity(count);
169 let mut descriptions: Vec<Option<String>> = Vec::with_capacity(count);
170
171 for row in rows {
172 element_ids.push(row.element_id);
173 identifications.push(row.identification.clone());
174 names.push(row.name.clone());
175 locations.push(row.location.clone());
176 descriptions.push(row.description.clone());
177 }
178
179 let schema = Schema::new(vec![
180 Field::new("element_id", DataType::UInt32, false),
181 Field::new("identification", DataType::Utf8, true),
182 Field::new("name", DataType::Utf8, true),
183 Field::new("location", DataType::Utf8, true),
184 Field::new("description", DataType::Utf8, true),
185 ]);
186
187 let batch = RecordBatch::try_new(
188 Arc::new(schema),
189 vec![
190 Arc::new(UInt32Array::from(element_ids)),
191 Arc::new(StringArray::from(identifications)),
192 Arc::new(StringArray::from(names)),
193 Arc::new(StringArray::from(locations)),
194 Arc::new(StringArray::from(descriptions)),
195 ],
196 )?;
197
198 write_parquet_batch(batch)
199}
200
201fn serialize_entities_table(entities: &[EntityMetadata]) -> Result<Vec<u8>, DataModelParquetError> {
203 let count = entities.len();
204
205 type Row = (
207 u32,
208 String,
209 String,
210 String,
211 bool,
212 Option<String>,
213 Option<String>,
214 Option<String>,
215 Option<String>,
216 );
217 let results: Vec<Row> = entities
218 .par_iter()
219 .map(|entity| {
220 (
221 entity.entity_id,
222 entity.type_name.clone(),
223 entity.global_id.clone().unwrap_or_default(),
224 entity.name.clone().unwrap_or_default(),
225 entity.has_geometry,
226 entity.description.clone(),
227 entity.object_type.clone(),
228 entity.tag.clone(),
229 entity.predefined_type.clone(),
230 )
231 })
232 .collect();
233
234 let mut entity_ids = Vec::with_capacity(count);
236 let mut type_names = Vec::with_capacity(count);
237 let mut global_ids = Vec::with_capacity(count);
238 let mut names = Vec::with_capacity(count);
239 let mut has_geometry = Vec::with_capacity(count);
240 let mut descriptions = Vec::with_capacity(count);
241 let mut object_types = Vec::with_capacity(count);
242 let mut tags = Vec::with_capacity(count);
243 let mut predefined_types = Vec::with_capacity(count);
244
245 for (id, type_name, global_id, name, has_geom, desc, obj, tag, predef) in results {
246 entity_ids.push(id);
247 type_names.push(type_name);
248 global_ids.push(global_id);
249 names.push(name);
250 has_geometry.push(has_geom);
251 descriptions.push(desc);
252 object_types.push(obj);
253 tags.push(tag);
254 predefined_types.push(predef);
255 }
256
257 let schema = Schema::new(vec![
258 Field::new("entity_id", DataType::UInt32, false),
259 Field::new("type_name", DataType::Utf8, false),
260 Field::new("global_id", DataType::Utf8, true),
261 Field::new("name", DataType::Utf8, true),
262 Field::new("has_geometry", DataType::Boolean, false),
263 Field::new("description", DataType::Utf8, true),
267 Field::new("object_type", DataType::Utf8, true),
268 Field::new("tag", DataType::Utf8, true),
269 Field::new("predefined_type", DataType::Utf8, true),
270 ]);
271
272 let batch = RecordBatch::try_new(
273 Arc::new(schema),
274 vec![
275 Arc::new(UInt32Array::from(entity_ids)),
276 Arc::new(StringArray::from(type_names)),
277 Arc::new(StringArray::from(global_ids)),
278 Arc::new(StringArray::from(names)),
279 Arc::new(BooleanArray::from(has_geometry)),
280 Arc::new(StringArray::from(descriptions)),
281 Arc::new(StringArray::from(object_types)),
282 Arc::new(StringArray::from(tags)),
283 Arc::new(StringArray::from(predefined_types)),
284 ],
285 )?;
286
287 write_parquet_batch(batch)
288}
289
290fn serialize_properties_table(
292 property_sets: &[PropertySet],
293) -> Result<Vec<u8>, DataModelParquetError> {
294 type Row = (
296 u32,
297 String,
298 String,
299 String,
300 String,
301 Option<String>,
302 bool,
303 Option<String>,
304 );
305 let rows: Vec<Row> = property_sets
306 .par_iter()
307 .flat_map_iter(|pset| {
308 pset.properties
309 .iter()
310 .map(move |prop| -> Result<Row, DataModelParquetError> {
311 let values_json = match &prop.values {
318 Some(v) => Some(serde_json::to_string(v)?),
319 None => None,
320 };
321 Ok((
322 pset.pset_id,
323 pset.pset_name.clone(),
324 prop.property_name.clone(),
325 prop.property_value.clone(),
326 prop.property_type.clone(),
327 prop.data_type.clone(),
328 prop.data_type_mixed,
329 values_json,
330 ))
331 })
332 })
333 .collect::<Result<Vec<Row>, _>>()?;
334
335 let mut pset_ids = Vec::with_capacity(rows.len());
337 let mut pset_names = Vec::with_capacity(rows.len());
338 let mut property_names = Vec::with_capacity(rows.len());
339 let mut property_values = Vec::with_capacity(rows.len());
340 let mut property_types = Vec::with_capacity(rows.len());
341 let mut data_types = Vec::with_capacity(rows.len());
342 let mut data_types_mixed = Vec::with_capacity(rows.len());
343 let mut values_jsons = Vec::with_capacity(rows.len());
344
345 for (pset_id, pset_name, prop_name, prop_value, prop_type, data_type, data_type_mixed, values_json) in rows {
346 pset_ids.push(pset_id);
347 pset_names.push(pset_name);
348 property_names.push(prop_name);
349 property_values.push(prop_value);
350 property_types.push(prop_type);
351 data_types.push(data_type);
352 data_types_mixed.push(data_type_mixed);
353 values_jsons.push(values_json);
354 }
355
356 let schema = Schema::new(vec![
357 Field::new("pset_id", DataType::UInt32, false),
358 Field::new("pset_name", DataType::Utf8, false),
359 Field::new("property_name", DataType::Utf8, false),
360 Field::new("property_value", DataType::Utf8, false),
361 Field::new("property_type", DataType::Utf8, false),
362 Field::new("data_type", DataType::Utf8, true),
365 Field::new("values_json", DataType::Utf8, true),
367 Field::new("data_type_mixed", DataType::Boolean, false),
370 ]);
371
372 let batch = RecordBatch::try_new(
373 Arc::new(schema),
374 vec![
375 Arc::new(UInt32Array::from(pset_ids)),
376 Arc::new(StringArray::from(pset_names)),
377 Arc::new(StringArray::from(property_names)),
378 Arc::new(StringArray::from(property_values)),
379 Arc::new(StringArray::from(property_types)),
380 Arc::new(StringArray::from(data_types)),
381 Arc::new(StringArray::from(values_jsons)),
382 Arc::new(BooleanArray::from(data_types_mixed)),
383 ],
384 )?;
385
386 write_parquet_batch(batch)
387}
388
389fn serialize_quantities_table(
391 quantity_sets: &[QuantitySet],
392) -> Result<Vec<u8>, DataModelParquetError> {
393 use arrow::array::Float64Array;
394
395 let rows: Vec<(u32, String, String, String, f64, String)> = quantity_sets
397 .par_iter()
398 .flat_map_iter(|qset| {
399 qset.quantities.iter().map(move |quant| {
400 (
401 qset.qset_id,
402 qset.qset_name.clone(),
403 qset.method_of_measurement.clone().unwrap_or_default(),
404 quant.quantity_name.clone(),
405 quant.quantity_value,
406 quant.quantity_type.clone(),
407 )
408 })
409 })
410 .collect();
411
412 let mut qset_ids = Vec::with_capacity(rows.len());
414 let mut qset_names = Vec::with_capacity(rows.len());
415 let mut methods = Vec::with_capacity(rows.len());
416 let mut quantity_names = Vec::with_capacity(rows.len());
417 let mut quantity_values = Vec::with_capacity(rows.len());
418 let mut quantity_types = Vec::with_capacity(rows.len());
419
420 for (qset_id, qset_name, method, quant_name, quant_value, quant_type) in rows {
421 qset_ids.push(qset_id);
422 qset_names.push(qset_name);
423 methods.push(method);
424 quantity_names.push(quant_name);
425 quantity_values.push(quant_value);
426 quantity_types.push(quant_type);
427 }
428
429 let schema = Schema::new(vec![
430 Field::new("qset_id", DataType::UInt32, false),
431 Field::new("qset_name", DataType::Utf8, false),
432 Field::new("method_of_measurement", DataType::Utf8, true),
433 Field::new("quantity_name", DataType::Utf8, false),
434 Field::new("quantity_value", DataType::Float64, false),
435 Field::new("quantity_type", DataType::Utf8, false),
436 ]);
437
438 let batch = RecordBatch::try_new(
439 Arc::new(schema),
440 vec![
441 Arc::new(UInt32Array::from(qset_ids)),
442 Arc::new(StringArray::from(qset_names)),
443 Arc::new(StringArray::from(methods)),
444 Arc::new(StringArray::from(quantity_names)),
445 Arc::new(Float64Array::from(quantity_values)),
446 Arc::new(StringArray::from(quantity_types)),
447 ],
448 )?;
449
450 write_parquet_batch(batch)
451}
452
453fn serialize_relationships_table(
455 relationships: &[Relationship],
456) -> Result<Vec<u8>, DataModelParquetError> {
457 let count = relationships.len();
458
459 let results: Vec<(String, u32, u32, u32)> = relationships
461 .par_iter()
462 .map(|rel| {
463 (
464 rel.rel_type.clone(),
465 rel.relating_id,
466 rel.related_id,
467 rel.rel_id,
468 )
469 })
470 .collect();
471
472 let mut rel_types = Vec::with_capacity(count);
473 let mut relating_ids = Vec::with_capacity(count);
474 let mut related_ids = Vec::with_capacity(count);
475 let mut rel_ids = Vec::with_capacity(count);
476
477 for (rel_type, relating_id, related_id, rel_id) in results {
478 rel_types.push(rel_type);
479 relating_ids.push(relating_id);
480 related_ids.push(related_id);
481 rel_ids.push(rel_id);
482 }
483
484 let schema = Schema::new(vec![
488 Field::new("rel_type", DataType::Utf8, false),
489 Field::new("relating_id", DataType::UInt32, false),
490 Field::new("related_id", DataType::UInt32, false),
491 Field::new("rel_id", DataType::UInt32, false),
492 ]);
493
494 let batch = RecordBatch::try_new(
495 Arc::new(schema),
496 vec![
497 Arc::new(StringArray::from(rel_types)),
498 Arc::new(UInt32Array::from(relating_ids)),
499 Arc::new(UInt32Array::from(related_ids)),
500 Arc::new(UInt32Array::from(rel_ids)),
501 ],
502 )?;
503
504 write_parquet_batch(batch)
505}
506
507fn serialize_spatial_hierarchy(
510 hierarchy: &SpatialHierarchyData,
511) -> Result<Vec<u8>, DataModelParquetError> {
512 let mut result = Vec::new();
513
514 let nodes_data = serialize_spatial_nodes_table(&hierarchy.nodes)?;
516 result.extend_from_slice(&(nodes_data.len() as u32).to_le_bytes());
517 result.extend_from_slice(&nodes_data);
518
519 let element_to_storey_data =
521 serialize_lookup_table(&hierarchy.element_to_storey, "element_to_storey")?;
522 result.extend_from_slice(&(element_to_storey_data.len() as u32).to_le_bytes());
523 result.extend_from_slice(&element_to_storey_data);
524
525 let element_to_building_data =
526 serialize_lookup_table(&hierarchy.element_to_building, "element_to_building")?;
527 result.extend_from_slice(&(element_to_building_data.len() as u32).to_le_bytes());
528 result.extend_from_slice(&element_to_building_data);
529
530 let element_to_site_data =
531 serialize_lookup_table(&hierarchy.element_to_site, "element_to_site")?;
532 result.extend_from_slice(&(element_to_site_data.len() as u32).to_le_bytes());
533 result.extend_from_slice(&element_to_site_data);
534
535 let element_to_space_data =
536 serialize_lookup_table(&hierarchy.element_to_space, "element_to_space")?;
537 result.extend_from_slice(&(element_to_space_data.len() as u32).to_le_bytes());
538 result.extend_from_slice(&element_to_space_data);
539
540 result.extend_from_slice(&hierarchy.project_id.to_le_bytes());
542
543 Ok(result)
544}
545
546fn serialize_spatial_nodes_table(
548 spatial_nodes: &[SpatialNode],
549) -> Result<Vec<u8>, DataModelParquetError> {
550 use arrow::array::Float64Array;
551
552 let count = spatial_nodes.len();
553 let mut entity_ids = Vec::with_capacity(count);
554 let mut parent_ids: Vec<Option<u32>> = Vec::with_capacity(count);
555 let mut levels = Vec::with_capacity(count);
556 let mut paths = Vec::with_capacity(count);
557 let mut type_names = Vec::with_capacity(count);
558 let mut names: Vec<Option<&str>> = Vec::with_capacity(count);
559 let mut elevations: Vec<Option<f64>> = Vec::with_capacity(count);
560 let mut children_ids_list = Vec::with_capacity(count);
561 let mut element_ids_list = Vec::with_capacity(count);
562
563 for node in spatial_nodes {
564 entity_ids.push(node.entity_id);
565 parent_ids.push(if node.parent_id == 0 {
566 None
567 } else {
568 Some(node.parent_id)
569 });
570 levels.push(node.level);
571 paths.push(node.path.as_str());
572 type_names.push(node.type_name.as_str());
573 names.push(node.name.as_deref());
574 elevations.push(node.elevation);
575 children_ids_list.push(node.children_ids.clone());
576 element_ids_list.push(node.element_ids.clone());
577 }
578
579 let mut children_values = Vec::new();
582 let mut children_offsets = vec![0i32];
583 let mut element_values = Vec::new();
584 let mut element_offsets = vec![0i32];
585
586 for children_ids in &children_ids_list {
587 children_values.extend_from_slice(children_ids);
588 children_offsets.push(children_values.len() as i32);
589 }
590
591 for element_ids in &element_ids_list {
592 element_values.extend_from_slice(element_ids);
593 element_offsets.push(element_values.len() as i32);
594 }
595
596 let mut children_builder =
598 ListBuilder::new(UInt32Builder::with_capacity(children_values.len()));
599 for children_ids in &children_ids_list {
600 children_builder.values().append_slice(children_ids);
601 children_builder.append(true);
602 }
603 let children_list_array = children_builder.finish();
604
605 let mut element_builder = ListBuilder::new(UInt32Builder::with_capacity(element_values.len()));
606 for element_ids in &element_ids_list {
607 element_builder.values().append_slice(element_ids);
608 element_builder.append(true);
609 }
610 let element_list_array = element_builder.finish();
611
612 let schema = Schema::new(vec![
614 Field::new("entity_id", DataType::UInt32, false),
615 Field::new("parent_id", DataType::UInt32, true), Field::new("level", DataType::UInt16, false),
617 Field::new("path", DataType::Utf8, false),
618 Field::new("type_name", DataType::Utf8, false),
619 Field::new("name", DataType::Utf8, true), Field::new("elevation", DataType::Float64, true), Field::new(
622 "children_ids",
623 DataType::new_list(DataType::UInt32, true),
624 false,
625 ), Field::new(
627 "element_ids",
628 DataType::new_list(DataType::UInt32, true),
629 false,
630 ), ]);
632
633 let batch = RecordBatch::try_new(
634 Arc::new(schema),
635 vec![
636 Arc::new(UInt32Array::from(entity_ids)),
637 Arc::new(UInt32Array::from(parent_ids)),
638 Arc::new(UInt16Array::from(levels)),
639 Arc::new(StringArray::from(paths)),
640 Arc::new(StringArray::from(type_names)),
641 Arc::new(StringArray::from(names)),
642 Arc::new(Float64Array::from(elevations)),
643 Arc::new(children_list_array),
644 Arc::new(element_list_array),
645 ],
646 )?;
647
648 write_parquet_batch(batch)
649}
650
651fn serialize_lookup_table(
653 pairs: &[(u32, u32)],
654 _table_name: &str,
655) -> Result<Vec<u8>, DataModelParquetError> {
656 let count = pairs.len();
657 let mut element_ids = Vec::with_capacity(count);
658 let mut spatial_ids = Vec::with_capacity(count);
659
660 for (element_id, spatial_id) in pairs {
661 element_ids.push(*element_id);
662 spatial_ids.push(*spatial_id);
663 }
664
665 let schema = Schema::new(vec![
666 Field::new("element_id", DataType::UInt32, false),
667 Field::new("spatial_id", DataType::UInt32, false),
668 ]);
669
670 let batch = RecordBatch::try_new(
671 Arc::new(schema),
672 vec![
673 Arc::new(UInt32Array::from(element_ids)),
674 Arc::new(UInt32Array::from(spatial_ids)),
675 ],
676 )?;
677
678 write_parquet_batch(batch)
679}
680
681fn write_parquet_batch(batch: RecordBatch) -> Result<Vec<u8>, DataModelParquetError> {
683 let mut buffer = Vec::new();
684 let cursor = Cursor::new(&mut buffer);
685
686 let props = WriterProperties::builder()
687 .set_compression(Compression::LZ4_RAW)
688 .build();
689
690 let mut writer = ArrowWriter::try_new(cursor, batch.schema(), Some(props))?;
691 writer.write(&batch)?;
692 writer.close()?;
693
694 Ok(buffer)
695}
696
697#[cfg(test)]
698#[path = "parquet_data_model_tests.rs"]
699mod tests;