ifc_lite_server/services/
parquet_data_model.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//! Parquet serialization for IFC data model (entities, properties, relationships, spatial hierarchy).
6
7use 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/// Errors during data model Parquet serialization.
29#[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
41/// Serialize data model to Parquet format.
42///
43/// Creates 5 Parquet tables:
44/// 1. Entities (entity_id, type_name, global_id, name, has_geometry)
45/// 2. Properties (pset_id, pset_name, property_name, property_value, property_type)
46/// 3. Quantities (qset_id, qset_name, method_of_measurement, quantity_name, quantity_value, quantity_type)
47/// 4. Relationships (rel_type, relating_id, related_id)
48/// 5. Spatial (entity_id, parent_id, level, path, type_name, name, elevation, children_ids, element_ids)
49///    Plus lookup tables: element_to_storey, element_to_building, element_to_site, element_to_space
50pub fn serialize_data_model_to_parquet(
51    data_model: &DataModel,
52) -> Result<Vec<u8>, DataModelParquetError> {
53    // Serialize all tables in parallel using rayon
54    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    // Classification / material / document tables (issue #900). Appended after
82    // the original five so older decoders that stop at `spatial` simply ignore
83    // the trailing bytes — the format stays backward compatible.
84    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    // Write format: [entities_len][entities_data][properties_len][properties_data][quantities_len][quantities_data][relationships_len][relationships_data][spatial_len][spatial_data][classifications_len][...][materials_len][...][documents_len][...]
98    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
119/// Serialize classification associations table.
120fn 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
160/// Serialize document associations table.
161fn 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
201/// Serialize entities table.
202fn serialize_entities_table(entities: &[EntityMetadata]) -> Result<Vec<u8>, DataModelParquetError> {
203    let count = entities.len();
204
205    // Build arrays in parallel using rayon
206    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    // Split into separate vectors
235    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        // Additive nullable columns (issue #1765; data-model cache bumped to
264        // v4). `description`/`object_type` were already read by the decoder
265        // by name; `tag`/`predefined_type` are new end-to-end.
266        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
290/// Serialize properties table.
291fn serialize_properties_table(
292    property_sets: &[PropertySet],
293) -> Result<Vec<u8>, DataModelParquetError> {
294    // Flatten property sets into rows using parallel iteration
295    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                    // Candidate arrays ride as a JSON string in a nullable Utf8
312                    // column (issue #1766): the decoder keeps its bulk column
313                    // reads and JSON-parses only the sparse multi-value rows. A
314                    // serialization failure is propagated, not conflated with
315                    // "no candidates" (which would silently drop them from the
316                    // cached payload) — though a Vec<String> never fails to encode.
317                    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    // Split into separate vectors
336    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        // Additive nullable column (data-model cache bumped to v3). Old decoders
363        // that don't request it simply ignore the extra column.
364        Field::new("data_type", DataType::Utf8, true),
365        // JSON-encoded candidate array (data-model cache bumped to v5).
366        Field::new("values_json", DataType::Utf8, true),
367        // IfcPropertyTableValue: no single data_type by design (#5224,
368        // data-model cache bumped to v7).
369        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
389/// Serialize quantities table.
390fn serialize_quantities_table(
391    quantity_sets: &[QuantitySet],
392) -> Result<Vec<u8>, DataModelParquetError> {
393    use arrow::array::Float64Array;
394
395    // Flatten quantity sets into rows using parallel iteration
396    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    // Split into separate vectors
413    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
453/// Serialize relationships table.
454fn serialize_relationships_table(
455    relationships: &[Relationship],
456) -> Result<Vec<u8>, DataModelParquetError> {
457    let count = relationships.len();
458
459    // Build arrays in parallel
460    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    // `rel_id` is appended last (data-model v6 payload, issue #3860): an older
485    // client reads the three columns it knows by name and ignores this one,
486    // and a newer client treats it as absent when an older server omits it.
487    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
507/// Serialize spatial hierarchy with nodes and lookup tables.
508/// Returns combined binary: [nodes_len][nodes_data][element_to_storey_len][element_to_storey_data]...
509fn serialize_spatial_hierarchy(
510    hierarchy: &SpatialHierarchyData,
511) -> Result<Vec<u8>, DataModelParquetError> {
512    let mut result = Vec::new();
513
514    // Serialize nodes table
515    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    // Serialize lookup tables
520    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    // Add project_id as final u32
541    result.extend_from_slice(&hierarchy.project_id.to_le_bytes());
542
543    Ok(result)
544}
545
546/// Serialize spatial nodes table with all fields.
547fn 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    // Build list arrays for children_ids and element_ids
580    // Flatten all values and build offset array
581    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    // Build ListArray using builder pattern
597    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    // Schema must match what ListBuilder produces - inner items are nullable by default
613    let schema = Schema::new(vec![
614        Field::new("entity_id", DataType::UInt32, false),
615        Field::new("parent_id", DataType::UInt32, true), // Nullable
616        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), // Nullable
620        Field::new("elevation", DataType::Float64, true), // Nullable
621        Field::new(
622            "children_ids",
623            DataType::new_list(DataType::UInt32, true),
624            false,
625        ), // Inner items nullable
626        Field::new(
627            "element_ids",
628            DataType::new_list(DataType::UInt32, true),
629            false,
630        ), // Inner items nullable
631    ]);
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
651/// Serialize a lookup table (element_id -> spatial_id pairs).
652fn 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
681/// Write a RecordBatch to a Parquet buffer with Zstd compression.
682fn 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;