friendlymatthew commented on code in PR #7740:
URL: https://github.com/apache/arrow-rs/pull/7740#discussion_r2164669063
##########
parquet-variant/src/builder.rs:
##########
@@ -148,189 +314,108 @@ fn make_room_for_header(buffer: &mut Vec<u8>,
start_pos: usize, header_size: usi
///
/// # Example: [`Variant::List`] of [`Variant::Object`]s
///
-/// THis example shows how to create an list of objects:
+/// This example shows how to create an list of objects:
/// ```json
/// [
-/// {
-/// "first_name": "Jiaying",
-/// "last_name": "Li"
-/// },
/// {
-/// "first_name": "Malthe",
-/// "last_name": "Karbo"
-/// }
+/// "id": 1,
+/// "type": "Cauliflower"
+/// },
+/// {
+/// "id": 2,
+/// "type": "Beets"
+/// }
/// ]
/// ```
+/// ```
+/// use parquet_variant::{Variant, VariantBuilder};
+/// let mut builder = VariantBuilder::new();
+///
+/// // Create a builder that will write elements to the list
+/// let mut list_builder = builder.new_list();
+///
+/// {
+/// let mut object_builder = list_builder.new_object();
+/// object_builder.append_value("id", 1);
+/// object_builder.append_value("type", "Cauliflower");
+/// object_builder.finish();
+/// }
+///
+/// {
+/// let mut object_builder = list_builder.new_object();
+/// object_builder.append_value("id", 2);
+/// object_builder.append_value("type", "Beets");
+/// object_builder.finish();
+/// }
+///
+/// list_builder.finish();
+/// // Finish the builder to get the metadata and value
+/// let (metadata, value) = builder.finish();
+/// // use the Variant API to verify the result
+/// let variant = Variant::try_new(&metadata, &value).unwrap();
+/// let variant_list = variant.as_list().unwrap();
+///
+///
+/// let obj1_variant = variant_list.get(0).unwrap();
+/// let obj1 = obj1_variant.as_object().unwrap();
+/// assert_eq!(
+/// obj1.field_by_name("id").unwrap(),
+/// Some(Variant::from(1))
+/// );
+/// assert_eq!(
+/// obj1.field_by_name("type").unwrap(),
+/// Some(Variant::from("Cauliflower"))
+/// );
///
-/// TODO
+/// let obj2_variant = variant_list.get(1).unwrap();
+/// let obj2 = obj2_variant.as_object().unwrap();
+///
+/// assert_eq!(
+/// obj2.field_by_name("id").unwrap(),
+/// Some(Variant::from(2))
+/// );
+/// assert_eq!(
+/// obj2.field_by_name("type").unwrap(),
+/// Some(Variant::from("Beets"))
+/// );
///
+/// ```
pub struct VariantBuilder {
- buffer: Vec<u8>,
- dict: BTreeMap<String, u32>,
- dict_keys: Vec<String>,
+ buffer: VariantBuffer,
+ field_metadata_dictionary: FieldMetadataDictionary,
}
impl VariantBuilder {
pub fn new() -> Self {
Self {
- buffer: Vec::new(),
- dict: BTreeMap::new(),
- dict_keys: Vec::new(),
- }
- }
-
- fn append_null(&mut self) {
- self.buffer
- .push(primitive_header(VariantPrimitiveType::Null));
- }
-
- fn append_bool(&mut self, value: bool) {
- let primitive_type = if value {
- VariantPrimitiveType::BooleanTrue
- } else {
- VariantPrimitiveType::BooleanFalse
- };
- self.buffer.push(primitive_header(primitive_type));
- }
-
- fn append_int8(&mut self, value: i8) {
- self.buffer
- .push(primitive_header(VariantPrimitiveType::Int8));
- self.buffer.push(value as u8);
- }
-
- fn append_int16(&mut self, value: i16) {
- self.buffer
- .push(primitive_header(VariantPrimitiveType::Int16));
- self.buffer.extend_from_slice(&value.to_le_bytes());
- }
-
- fn append_int32(&mut self, value: i32) {
- self.buffer
- .push(primitive_header(VariantPrimitiveType::Int32));
- self.buffer.extend_from_slice(&value.to_le_bytes());
- }
-
- fn append_int64(&mut self, value: i64) {
- self.buffer
- .push(primitive_header(VariantPrimitiveType::Int64));
- self.buffer.extend_from_slice(&value.to_le_bytes());
- }
-
- fn append_float(&mut self, value: f32) {
- self.buffer
- .push(primitive_header(VariantPrimitiveType::Float));
- self.buffer.extend_from_slice(&value.to_le_bytes());
- }
-
- fn append_double(&mut self, value: f64) {
- self.buffer
- .push(primitive_header(VariantPrimitiveType::Double));
- self.buffer.extend_from_slice(&value.to_le_bytes());
- }
-
- fn append_date(&mut self, value: chrono::NaiveDate) {
- self.buffer
- .push(primitive_header(VariantPrimitiveType::Date));
- let days_since_epoch =
value.signed_duration_since(UNIX_EPOCH_DATE).num_days() as i32;
- self.buffer
- .extend_from_slice(&days_since_epoch.to_le_bytes());
- }
-
- fn append_timestamp_micros(&mut self, value:
chrono::DateTime<chrono::Utc>) {
- self.buffer
- .push(primitive_header(VariantPrimitiveType::TimestampMicros));
- let micros = value.timestamp_micros();
- self.buffer.extend_from_slice(µs.to_le_bytes());
- }
-
- fn append_timestamp_ntz_micros(&mut self, value: chrono::NaiveDateTime) {
- self.buffer
- .push(primitive_header(VariantPrimitiveType::TimestampNtzMicros));
- let micros = value.and_utc().timestamp_micros();
- self.buffer.extend_from_slice(µs.to_le_bytes());
- }
-
- fn append_decimal4(&mut self, integer: i32, scale: u8) {
- self.buffer
- .push(primitive_header(VariantPrimitiveType::Decimal4));
- self.buffer.push(scale);
- self.buffer.extend_from_slice(&integer.to_le_bytes());
- }
-
- fn append_decimal8(&mut self, integer: i64, scale: u8) {
- self.buffer
- .push(primitive_header(VariantPrimitiveType::Decimal8));
- self.buffer.push(scale);
- self.buffer.extend_from_slice(&integer.to_le_bytes());
- }
-
- fn append_decimal16(&mut self, integer: i128, scale: u8) {
- self.buffer
- .push(primitive_header(VariantPrimitiveType::Decimal16));
- self.buffer.push(scale);
- self.buffer.extend_from_slice(&integer.to_le_bytes());
- }
-
- fn append_binary(&mut self, value: &[u8]) {
- self.buffer
- .push(primitive_header(VariantPrimitiveType::Binary));
- self.buffer
- .extend_from_slice(&(value.len() as u32).to_le_bytes());
- self.buffer.extend_from_slice(value);
- }
-
- fn append_short_string(&mut self, value: ShortString) {
- let inner = value.0;
- self.buffer.push(short_string_header(inner.len()));
- self.buffer.extend_from_slice(inner.as_bytes());
- }
-
- fn append_string(&mut self, value: &str) {
- self.buffer
- .push(primitive_header(VariantPrimitiveType::String));
- self.buffer
- .extend_from_slice(&(value.len() as u32).to_le_bytes());
- self.buffer.extend_from_slice(value.as_bytes());
- }
-
- /// Add key to dictionary, return its ID
- fn add_key(&mut self, key: &str) -> u32 {
- use std::collections::btree_map::Entry;
- match self.dict.entry(key.to_string()) {
- Entry::Occupied(entry) => *entry.get(),
- Entry::Vacant(entry) => {
- let id = self.dict_keys.len() as u32;
- entry.insert(id);
- self.dict_keys.push(key.to_string());
- id
- }
+ buffer: VariantBuffer::default(),
+ field_metadata_dictionary: FieldMetadataDictionary::default(),
}
}
- fn offset(&self) -> usize {
- self.buffer.len()
- }
-
/// Create an [`ListBuilder`] for creating [`Variant::List`] values.
///
/// See the examples on [`VariantBuilder`] for usage.
pub fn new_list(&mut self) -> ListBuilder {
- ListBuilder::new(self)
+ ListBuilder::new(&mut self.buffer, &mut self.field_metadata_dictionary)
}
/// Create an [`ObjectBuilder`] for creating [`Variant::Object`] values.
///
/// See the examples on [`VariantBuilder`] for usage.
pub fn new_object(&mut self) -> ObjectBuilder {
- ObjectBuilder::new(self)
+ ObjectBuilder::new(&mut self.buffer, &mut
self.field_metadata_dictionary)
+ }
+
+ pub fn append_value<'m, 'd, T: Into<Variant<'m, 'd>>>(&mut self, value: T)
{
+ self.buffer.append_value(value);
}
pub fn finish(self) -> (Vec<u8>, Vec<u8>) {
- let nkeys = self.dict_keys.len();
+ let nkeys = self.field_metadata_dictionary.num_field_names();
Review Comment:
I like that! Could we do this in a subsequent PR? I have another PR that
will follow this that does nested object and object with list building
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
