perf: improve bson iteration performance

crates/datasources/src/bson/builder.rs

@@ -90,27 +90,26 @@ impl RecordStructBuilder {
 
     pub fn append_record(&mut self, doc: &RawDocument) -> Result<()> {
         let mut cols_set: BitVec<u8, Lsb0> = BitVec::repeat(false, self.fields.len());
+        let doc_buf = doc.to_raw_document_buf();
+        let iter = doc_buf.iter_elements();
 
-        for iter_result in doc {
-            match iter_result {
-                Ok((key, val)) => {
-                    let idx = *self
-                        .field_index
-                        .get(key)
-                        .ok_or_else(|| BsonError::ColumnNotInInferredSchema(key.to_string()))?;
-
-                    if *cols_set.get(idx).unwrap() {
-                        continue;
-                    }
+        for item in iter {
+            let elem = item?;
 
-                    // Add value to columns.
-                    self.add_value_at_index(idx, Some(val))?;
+            let idx = *self
+                .field_index
+                .get(elem.key())
+                .ok_or_else(|| BsonError::ColumnNotInInferredSchema(elem.key().to_string()))?;
 
-                    // Track which columns we've added values to.
-                    cols_set.set(idx, true);
-                }
-                Err(_) => return Err(BsonError::FailedToReadRawBsonDocument),
+            if *cols_set.get(idx).unwrap() {
+                continue;
             }
+
+            // Add value to columns.
+            self.add_value_at_index(idx, Some(elem.value()?))?;
+
+            // Track which columns we've added values to.
+            cols_set.set(idx, true);
         }
 
         // Append nulls to all columns not included in the doc.

crates/datasources/src/bson/schema.rs

@@ -12,30 +12,24 @@ use crate::bson::errors::{BsonError, Result};
 const RECURSION_LIMIT: usize = 100;
 
 pub fn schema_from_document(doc: &RawDocumentBuf) -> Result<Schema> {
-    Ok(Schema::new(fields_from_document(
-        0,
-        doc.iter().map(|item| item.map_err(|e| e.into())),
-    )?))
+    Ok(Schema::new(fields_from_document(0, doc)?))
 }
 
-fn fields_from_document<'a>(
-    depth: usize,
-    doc_iter: impl Iterator<Item = Result<(&'a str, RawBsonRef<'a>)>>,
-) -> Result<Vec<Field>> {
+fn fields_from_document(depth: usize, doc: &RawDocumentBuf) -> Result<Vec<Field>> {
     if depth >= RECURSION_LIMIT {
         return Err(BsonError::RecursionLimitExceeded(RECURSION_LIMIT));
     }
 
-    // let doc_iter = doc.iter();
-    let (_, size) = doc_iter.size_hint();
+    let iter = doc.iter_elements();
+    let (_, size) = iter.size_hint();
     let mut fields = Vec::with_capacity(size.unwrap_or_default());
 
-    for item in doc_iter {
-        let (key, val) = item?;
-        let arrow_typ = bson_to_arrow_type(depth, val)?;
+    for item in iter {
+        let elem = item?;
+        let arrow_typ = bson_to_arrow_type(depth, elem.value()?)?;
 
         // Assume everything is nullable.
-        fields.push(Field::new(key, arrow_typ, true));
+        fields.push(Field::new(elem.key(), arrow_typ, true));
     }
 
     Ok(fields)
@@ -56,13 +50,9 @@ fn bson_to_arrow_type(depth: usize, bson: RawBsonRef) -> Result<DataType> {
             )?,
             true,
         ),
-        RawBsonRef::Document(nested) => DataType::Struct(
-            fields_from_document(
-                depth + 1,
-                nested.into_iter().map(|item| item.map_err(|e| e.into())),
-            )?
-            .into(),
-        ),
+        RawBsonRef::Document(nested) => {
+            DataType::Struct(fields_from_document(depth + 1, &nested.to_raw_document_buf())?.into())
+        }
         RawBsonRef::String(_) => DataType::Utf8,
         RawBsonRef::Double(_) => DataType::Float64,
         RawBsonRef::Boolean(_) => DataType::Boolean,