fcb_core/writer/
attribute.rs

1use crate::fb::ColumnType;
2use byteorder::{ByteOrder, LittleEndian};
3use chrono::{DateTime, Utc};
4use cjseq::CityJSONFeature;
5use serde_json::Value;
6use std::collections::BTreeMap;
7use tracing::{debug, warn};
8
9// Schema for attributes. The key is the attribute name, the value is a tuple of the column index and the column type.
10//
11// This is a `BTreeMap`, not a `HashMap`, on purpose. It is not load-bearing
12// for column order or column index: both `to_columns` (serializer.rs) and
13// `encode_attributes_with_schema` (below) collect the schema into a `Vec` and
14// `sort_by_key` on the stored column index before emitting anything, and the
15// column index assigned at insert time comes from `self.len()`, a count that
16// is independent of iteration order. The `BTreeMap` is defensive: it keeps
17// callers that build a schema by iterating it directly (rather than going
18// through the sorted accessors) from reintroducing `HashMap`'s randomly
19// seeded per-process order and making the writer non-reproducible.
20pub type AttributeSchema = BTreeMap<String, (u16, ColumnType)>;
21
22pub trait AttributeSchemaMethods {
23    fn add_attributes(&mut self, attrs: &Value);
24}
25
26impl AttributeSchemaMethods for AttributeSchema {
27    fn add_attributes(&mut self, attrs: &Value) {
28        if !attrs.is_object() {
29            self.insert("json".to_string(), (self.len() as u16, ColumnType::Json));
30            return;
31        }
32
33        let map = attrs.as_object().unwrap();
34        for (key, val) in map.iter() {
35            if !self.contains_key(key) && !val.is_null() {
36                if let Some(coltype) = guess_type(val) {
37                    self.insert(key.clone(), (self.len() as u16, coltype));
38                }
39            }
40        }
41    }
42}
43
44/// Naive type-guessing. You could use your schema or logic as in your Python code.
45fn guess_type(value: &Value) -> Option<ColumnType> {
46    match value {
47        Value::Bool(_) => Some(ColumnType::Bool),
48        Value::Number(n) => {
49            if n.is_f64() {
50                Some(ColumnType::Double)
51            } else if n.is_u64() {
52                Some(ColumnType::ULong)
53            } else if n.is_i64() {
54                Some(ColumnType::Long)
55            } else {
56                Some(ColumnType::ULong) // Fallback for unknown number type.
57            }
58        }
59        Value::String(s) => {
60            // Attempt to parse the string as an RFC3339 date.
61            if chrono::DateTime::parse_from_rfc3339(s).is_ok() {
62                Some(ColumnType::DateTime)
63            } else {
64                Some(ColumnType::String)
65            }
66        }
67        Value::Array(_) => Some(ColumnType::Json),
68        Value::Object(_) => Some(ColumnType::Json),
69        _ => None,
70    }
71}
72
73pub(crate) fn attr_size(coltype: &ColumnType, colval: &Value) -> usize {
74    match *coltype {
75        ColumnType::Byte => size_of::<i8>(),
76        ColumnType::UByte => size_of::<u8>(),
77        ColumnType::Bool => size_of::<u8>(),
78        ColumnType::Short => size_of::<i16>(),
79        ColumnType::UShort => size_of::<u16>(),
80        ColumnType::Int => size_of::<i32>(),
81        ColumnType::UInt => size_of::<u32>(),
82        ColumnType::Long => size_of::<i64>(),
83        ColumnType::ULong => size_of::<u64>(),
84        ColumnType::Float => size_of::<f32>(),
85        ColumnType::Double => size_of::<f64>(),
86        ColumnType::String | ColumnType::DateTime => {
87            size_of::<u32>() + colval.as_str().unwrap().len()
88        }
89        ColumnType::Json => {
90            let json = serde_json::to_string(colval).unwrap_or_default();
91            size_of::<u32>() + json.len()
92        }
93        ColumnType::Binary => size_of::<u32>() + colval.as_str().unwrap().len(), //TODO: check if this is correct
94        _ => unreachable!(),
95    }
96}
97
98pub(crate) fn encode_attributes_with_schema(attr: &Value, schema: &AttributeSchema) -> Vec<u8> {
99    if !attr.is_object() || attr.as_object().unwrap().is_empty() || attr.is_null() {
100        return Vec::new();
101    }
102
103    let mut out = Vec::new();
104    let mut sorted_schema: Vec<_> = schema.iter().collect();
105    sorted_schema.sort_by_key(|(_, (index, _))| *index);
106
107    for (name, (index, coltype)) in sorted_schema {
108        let (_, val) = {
109            let attr_obj = attr.as_object();
110            if let Some(attr_obj) = attr_obj {
111                let value = attr_obj.iter().find(|(k, _)| *k == name);
112                if let Some(value) = value {
113                    (value.0, value.1)
114                } else {
115                    continue;
116                }
117            } else {
118                return Vec::new();
119            }
120        };
121
122        if val.is_null() {
123            continue;
124        }
125
126        let mut offset = out.len();
127        let attr_size = attr_size(coltype, val);
128
129        // Reserve space for index and value
130        out.resize(offset + size_of::<u16>() + attr_size, 0);
131
132        // Write index
133        LittleEndian::write_u16(&mut out[offset..], *index);
134        offset += size_of::<u16>();
135
136        match *coltype {
137            ColumnType::Bool => {
138                let b = val.as_bool().unwrap_or(false);
139                out[offset] = b as u8;
140            }
141            ColumnType::Int => {
142                let i = val.as_i64().unwrap_or(0);
143                LittleEndian::write_i32(&mut out[offset..], i as i32);
144            }
145            ColumnType::UInt => {
146                let i = val.as_u64().unwrap_or(0);
147                LittleEndian::write_u32(&mut out[offset..], i as u32);
148            }
149            ColumnType::Byte => {
150                let b = val.as_i64().unwrap_or(0);
151                out[offset] = b as u8;
152            }
153            ColumnType::UByte => {
154                let b = val.as_u64().unwrap_or(0);
155                out[offset] = b as u8;
156            }
157
158            ColumnType::Short => {
159                let i = val.as_i64().unwrap_or(0);
160                LittleEndian::write_i16(&mut out[offset..], i as i16);
161            }
162            ColumnType::UShort => {
163                let i = val.as_u64().unwrap_or(0);
164                LittleEndian::write_u16(&mut out[offset..], i as u16);
165            }
166
167            ColumnType::Long => {
168                let i = val.as_i64().unwrap_or(0);
169                LittleEndian::write_i64(&mut out[offset..], i);
170            }
171            ColumnType::ULong => {
172                let i = val.as_u64().unwrap_or(0);
173                LittleEndian::write_u64(&mut out[offset..], i);
174            }
175            ColumnType::Float => {
176                let f = val.as_f64().unwrap_or(0.0);
177                LittleEndian::write_f32(&mut out[offset..], f as f32);
178            }
179            ColumnType::Double => {
180                let f = val.as_f64().unwrap_or(0.0);
181                LittleEndian::write_f64(&mut out[offset..], f);
182            }
183            ColumnType::String | ColumnType::DateTime => {
184                let s = val.as_str().unwrap_or("");
185                LittleEndian::write_u32(&mut out[offset..], s.len() as u32);
186                out[offset + size_of::<u32>()..offset + size_of::<u32>() + s.len()]
187                    .copy_from_slice(s.as_bytes());
188            }
189            ColumnType::Json => {
190                let json = serde_json::to_string(val).unwrap_or_default();
191                LittleEndian::write_u32(&mut out[offset..], json.len() as u32);
192                out[offset + size_of::<u32>()..offset + size_of::<u32>() + json.len()]
193                    .copy_from_slice(json.as_bytes());
194            }
195            ColumnType::Binary => {
196                let s = val.as_str().unwrap_or("");
197                LittleEndian::write_u32(&mut out[offset..], s.len() as u32);
198                out[offset + size_of::<u32>()..offset + size_of::<u32>() + s.len()]
199                    .copy_from_slice(s.as_bytes());
200            }
201            _ => unreachable!(),
202        }
203    }
204    out
205}
206
207#[derive(Clone, PartialEq, Debug)]
208pub enum AttributeIndexEntry {
209    Bool { index: u16, val: bool },
210    Int { index: u16, val: i32 },
211    UInt { index: u16, val: u32 },
212    Long { index: u16, val: i64 },
213    ULong { index: u16, val: u64 },
214    Float { index: u16, val: f32 },
215    Double { index: u16, val: f64 },
216    String { index: u16, val: String },
217    DateTime { index: u16, val: DateTime<Utc> },
218    Short { index: u16, val: i16 },
219    UShort { index: u16, val: u16 },
220    Byte { index: u16, val: u8 },
221    UByte { index: u16, val: u8 },
222    Json { index: u16, val: String },
223    Binary { index: u16, val: String },
224}
225
226pub fn cityfeature_to_index_entries(
227    cityfeature: &CityJSONFeature,
228    schema: &AttributeSchema,
229    indexing_attr: &[String],
230) -> Vec<AttributeIndexEntry> {
231    let mut index_entries = Vec::new();
232    // `city_objects` is a `HashMap` in cjseq, so iterate it in a fixed (id)
233    // order -- the entry order reaches the B+tree payload lists for duplicate
234    // keys, and a random order there is a random output file.
235    let mut object_ids: Vec<_> = cityfeature.city_objects.keys().collect();
236    object_ids.sort_unstable();
237    for object in object_ids
238        .into_iter()
239        .filter_map(|id| cityfeature.city_objects.get(id))
240    {
241        if let Some(attr) = &object.attributes {
242            let attr_index_entries = attribute_to_index_entries(attr, schema, indexing_attr);
243            index_entries.extend(attr_index_entries);
244        }
245    }
246
247    index_entries
248}
249
250// this attr should be a json object with attribute name as key and attribute value as value
251pub fn attribute_to_index_entries(
252    attr: &Value,
253    schema: &AttributeSchema,
254    indexing_attr: &[String],
255) -> Vec<AttributeIndexEntry> {
256    if !attr.is_object() || attr.is_null() || attr.as_object().unwrap().is_empty() {
257        return Vec::new();
258    }
259
260    let mut index_entries = Vec::new();
261
262    let map = attr.as_object().unwrap();
263    for attr in indexing_attr {
264        let val: &Value = match map.get(attr) {
265            Some(val) => val,
266            None => {
267                // Never `println!`: library diagnostics on stdout corrupt
268                // `fcb_cli ser <input> -`, which writes the binary there.
269                debug!("feature is missing indexed attribute {attr}");
270                continue;
271            }
272        };
273
274        let index_coltype = schema.get(attr);
275        if let Some((index, coltype)) = index_coltype {
276            match *coltype {
277                ColumnType::Bool => {
278                    let b = val.as_bool().unwrap_or(false);
279                    index_entries.push(AttributeIndexEntry::Bool {
280                        index: *index,
281                        val: b,
282                    });
283                }
284                ColumnType::Int => {
285                    let i = val.as_i64().unwrap_or(0);
286                    index_entries.push(AttributeIndexEntry::Int {
287                        index: *index,
288                        val: i as i32,
289                    });
290                }
291                ColumnType::UInt => {
292                    let i = val.as_u64().unwrap_or(0);
293                    index_entries.push(AttributeIndexEntry::UInt {
294                        index: *index,
295                        val: i as u32,
296                    });
297                }
298                ColumnType::Long => {
299                    let i = val.as_i64().unwrap_or(0);
300                    index_entries.push(AttributeIndexEntry::Long {
301                        index: *index,
302                        val: i as i64,
303                    });
304                }
305                ColumnType::ULong => {
306                    let i = val.as_u64().unwrap_or(0);
307                    index_entries.push(AttributeIndexEntry::ULong {
308                        index: *index,
309                        val: i as u64,
310                    });
311                }
312                ColumnType::Float => {
313                    let f = val.as_f64().unwrap_or(0.0);
314                    index_entries.push(AttributeIndexEntry::Float {
315                        index: *index,
316                        val: f as f32,
317                    });
318                }
319                ColumnType::Double => {
320                    let f = val.as_f64().unwrap_or(0.0);
321                    index_entries.push(AttributeIndexEntry::Double {
322                        index: *index,
323                        val: f,
324                    });
325                }
326                ColumnType::String => {
327                    index_entries.push(AttributeIndexEntry::String {
328                        index: *index,
329                        val: val.as_str().unwrap_or("").to_string(),
330                    });
331                }
332                ColumnType::DateTime => {
333                    let dt = match chrono::DateTime::parse_from_rfc3339(val.as_str().unwrap_or(""))
334                    {
335                        Ok(dt) => dt.to_utc(),
336                        Err(e) => {
337                            warn!("failed to parse DateTime for {attr}, defaulting to epoch: {e}");
338                            // Choose whether to skip, default, or handle differently
339                            // For example, default to 1970-01-01:
340                            DateTime::<Utc>::from_timestamp(0, 0).unwrap()
341                        }
342                    };
343                    index_entries.push(AttributeIndexEntry::DateTime {
344                        index: *index,
345                        val: dt,
346                    });
347                }
348                _ => {
349                    //Byte, Ubyte,
350                    debug!("attribute {attr} is not supported for indexing");
351                }
352            }
353        }
354    }
355
356    index_entries
357}
358
359#[cfg(test)]
360mod tests {
361    use crate::{
362        deserializer::decode_attributes,
363        root_as_city_feature, root_as_header,
364        serializer::{to_columns, to_fcb_attribute},
365        CityFeature, CityFeatureArgs, CityObject, CityObjectArgs, Header, HeaderArgs,
366    };
367
368    use super::*;
369
370    use anyhow::Result;
371    use flatbuffers::FlatBufferBuilder;
372    use pretty_assertions::assert_eq;
373    use serde_json::json;
374
375    #[test]
376    fn test_add_attributes() -> Result<()> {
377        let json_data = json!({
378            "attributes": {
379                "int": -10,
380                "uint": 5,
381                "bool": true,
382                "float": 1.0,
383                "string": "hoge",
384                "array": [1, 2, 3],
385                "json": {
386                    "hoge": "fuga"
387                },
388                "null": null
389            }
390        });
391
392        let mut attr_schema: AttributeSchema = AttributeSchema::new();
393
394        attr_schema.add_attributes(&json_data["attributes"]);
395
396        // Check if the schema contains the expected keys and types
397        assert_eq!(attr_schema.get("int").unwrap().1, ColumnType::Long);
398        assert_eq!(attr_schema.get("uint").unwrap().1, ColumnType::ULong);
399        assert_eq!(attr_schema.get("bool").unwrap().1, ColumnType::Bool);
400        assert_eq!(attr_schema.get("float").unwrap().1, ColumnType::Double);
401        assert_eq!(attr_schema.get("string").unwrap().1, ColumnType::String);
402        assert_eq!(attr_schema.get("array").unwrap().1, ColumnType::Json); //TODO: check if this is correct
403        assert_eq!(attr_schema.get("json").unwrap().1, ColumnType::Json);
404
405        Ok(())
406    }
407
408    #[test]
409    fn test_attribute_serialization() -> Result<()> {
410        let test_cases = vec![
411            // Case 1: Same schema
412            (
413                json!({
414                        "int": -10,
415                        "uint": 5,
416                        "bool": true,
417                        "float": 1.0,
418                        "string": "hoge",
419                        "array": [1, 2, 3],
420                        "json": {
421                            "hoge": "fuga"
422                        }
423                }),
424                json!({
425                        "int": -10,
426                        "uint": 5,
427                            "bool": true,
428                        "float": 1.0,
429                        "string": "hoge",
430                        "array": [1, 2, 3],
431                        "json": {
432                            "hoge": "fuga"
433                    },
434                }),
435                json!({
436                    "attributes": {
437                        "int": -10,
438                        "uint": 5,
439                        "bool": true,
440                        "float": 1.0,
441                        "string": "hoge",
442                        "array": [1, 2, 3],
443                        "json": {
444                            "hoge": "fuga"
445                        }
446                    }
447                }),
448                "same schema",
449            ),
450            // Case 2: JSON with null value
451            (
452                json!({
453                            "int": -10,
454                        "uint": 5,
455                        "bool": true,
456                        "float": 1.0,
457                        "string": "hoge",
458                        "array": [1, 2, 3],
459                        "json": {
460                            "hoge": "fuga"
461                        },
462                        "exception": null
463                }),
464                json!({
465                            "int": -10,
466                        "uint": 5,
467                        "bool": true,
468                        "float": 1.0,
469                        "string": "hoge",
470                        "array": [1, 2, 3],
471                        "json": {
472                            "hoge": "fuga"
473                        },
474                }),
475                json!({
476                    "attributes": {
477                        "int": -10,
478                        "uint": 5,
479                        "bool": true,
480                        "float": 1.0,
481                        "string": "hoge",
482                        "array": [1, 2, 3],
483                        "json": {
484                            "hoge": "fuga"
485                        },
486                        "exception": 1000
487                    }
488                }),
489                "JSON with null value",
490            ),
491            // Case 3: JSON is empty
492            (
493                json!({}),
494                json!({}),
495                json!({
496                    "attributes": {
497                        "int": -10,
498                        "uint": 5,
499                        "bool": true,
500                        "float": 1.0,
501                        "string": "hoge",
502                        "array": [1, 2, 3],
503                        "json": {
504                            "hoge": "fuga"
505                        },
506                        "exception": 1000
507                    }
508                }),
509                "JSON is empty",
510            ),
511        ];
512
513        for (input, expected, schema, test_name) in test_cases {
514            println!("Testing case: {test_name}");
515
516            let attrs = &input;
517            let attr_schema = &schema["attributes"];
518
519            // Create and encode with schema
520            let mut fbb = FlatBufferBuilder::new();
521            let mut common_schema = AttributeSchema::new();
522            common_schema.add_attributes(attr_schema);
523
524            let columns = to_columns(&mut fbb, &common_schema);
525            let header = {
526                let version = fbb.create_string("1.0.0");
527                Header::create(
528                    &mut fbb,
529                    &HeaderArgs {
530                        version: Some(version),
531                        columns: Some(columns),
532                        ..Default::default()
533                    },
534                )
535            };
536            fbb.finish(header, None);
537
538            // Decode and verify
539            let finished_data = fbb.finished_data();
540            let header_buf = root_as_header(finished_data).unwrap();
541
542            let mut fbb = FlatBufferBuilder::new();
543            let feature = {
544                let (attr_buf, _) = to_fcb_attribute(&mut fbb, attrs, &common_schema);
545                let city_object = {
546                    let id = fbb.create_string("test");
547                    CityObject::create(
548                        &mut fbb,
549                        &CityObjectArgs {
550                            id: Some(id),
551                            attributes: Some(attr_buf),
552                            ..Default::default()
553                        },
554                    )
555                };
556                let objects = fbb.create_vector(&[city_object]);
557                let cf_id = fbb.create_string("test_feature");
558                CityFeature::create(
559                    &mut fbb,
560                    &CityFeatureArgs {
561                        id: Some(cf_id),
562                        objects: Some(objects),
563                        ..Default::default()
564                    },
565                )
566            };
567
568            fbb.finish(feature, None);
569
570            let finished_data = fbb.finished_data();
571            let feature_buf = root_as_city_feature(finished_data).unwrap();
572            let attributes = feature_buf.objects().unwrap().get(0).attributes().unwrap();
573
574            let decoded = decode_attributes(&header_buf.columns().unwrap(), attributes);
575
576            assert_eq!(
577                expected, decoded,
578                "decoded data should match original for {}",
579                test_name
580            );
581        }
582
583        Ok(())
584    }
585}
fcb_core/writer/attribute.rs

fcb_core/writer/
attribute.rs
