fix: bson type handling cleanup

crates/datafusion_ext/src/errors.rs

@@ -30,9 +30,6 @@ pub enum ExtensionError {
     #[error(transparent)]
     Arrow(#[from] datafusion::arrow::error::ArrowError),
 
-    #[error(transparent)]
-    BsonDeserialize(#[from] bson::de::Error),
-
     #[error(transparent)]
     DecimalError(#[from] decimal::DecimalError),
 
@@ -44,6 +41,9 @@ pub enum ExtensionError {
 
     #[error(transparent)]
     ListingErrBoxed(#[from] Box<dyn std::error::Error + Sync + Send>),
+
+    #[error("object store: {0}")]
+    ObjectStore(String),
 }
 
 impl ExtensionError {

crates/datasources/src/bson/builder.rs

@@ -5,9 +5,10 @@ use std::sync::Arc;
 use bitvec::{order::Lsb0, vec::BitVec};
 use bson::{RawBsonRef, RawDocument};
 use datafusion::arrow::array::{
-    Array, ArrayBuilder, ArrayRef, BinaryBuilder, BooleanBuilder, Decimal128Builder,
-    Float64Builder, Int32Builder, Int64Builder, StringBuilder, StructArray,
-    TimestampMicrosecondBuilder, TimestampMillisecondBuilder,
+    Array, ArrayBuilder, ArrayRef, BinaryBuilder, BooleanBuilder, Date32Builder, Date64Builder,
+    Decimal128Builder, Float64Builder, Int32Builder, Int64Builder, LargeBinaryBuilder,
+    LargeStringBuilder, StringBuilder, StructArray, TimestampMicrosecondBuilder,
+    TimestampMillisecondBuilder, TimestampSecondBuilder,
 };
 use datafusion::arrow::datatypes::{DataType, Field, Fields, TimeUnit};
 
@@ -182,13 +183,10 @@ fn append_value(val: RawBsonRef, typ: &DataType, col: &mut dyn ArrayBuilder) ->
     // So robust
     match (val, typ) {
         // Boolean
-        (RawBsonRef::Boolean(v), DataType::Boolean) => {
-            append_scalar!(BooleanBuilder, col, v)
-        }
+        (RawBsonRef::Boolean(v), DataType::Boolean) => append_scalar!(BooleanBuilder, col, v),
         (RawBsonRef::Boolean(v), DataType::Utf8) => {
             append_scalar!(StringBuilder, col, v.to_string())
         }
-
         // Double
         (RawBsonRef::Double(v), DataType::Int32) => append_scalar!(Int32Builder, col, v as i32),
         (RawBsonRef::Double(v), DataType::Int64) => append_scalar!(Int64Builder, col, v as i64),
@@ -214,6 +212,8 @@ fn append_value(val: RawBsonRef, typ: &DataType, col: &mut dyn ArrayBuilder) ->
         }
 
         // String
+        (RawBsonRef::String(v), DataType::Utf8) => append_scalar!(StringBuilder, col, v),
+        (RawBsonRef::String(v), DataType::LargeUtf8) => append_scalar!(LargeStringBuilder, col, v),
         (RawBsonRef::String(v), DataType::Boolean) => {
             append_scalar!(BooleanBuilder, col, v.parse().unwrap_or_default())
         }
@@ -226,31 +226,61 @@ fn append_value(val: RawBsonRef, typ: &DataType, col: &mut dyn ArrayBuilder) ->
         (RawBsonRef::String(v), DataType::Float64) => {
             append_scalar!(Float64Builder, col, v.parse().unwrap_or_default())
         }
-        (RawBsonRef::String(v), DataType::Utf8) => {
-            append_scalar!(StringBuilder, col, v)
-        }
 
         // Binary
         (RawBsonRef::Binary(v), DataType::Binary) => append_scalar!(BinaryBuilder, col, v.bytes),
+        (RawBsonRef::Binary(v), DataType::LargeBinary) => {
+            append_scalar!(LargeBinaryBuilder, col, v.bytes)
+        }
 
         // Object id
+        (RawBsonRef::ObjectId(v), DataType::Binary) => {
+            append_scalar!(BinaryBuilder, col, v.bytes())
+        }
         (RawBsonRef::ObjectId(v), DataType::Utf8) => {
             append_scalar!(StringBuilder, col, v.to_string())
         }
 
-        // Timestamp
+        // Timestamp (internal mongodb type; second specified)
+        (RawBsonRef::Timestamp(v), DataType::Timestamp(TimeUnit::Second, _)) => {
+            append_scalar!(TimestampSecondBuilder, col, v.time as i64)
+        }
+        (RawBsonRef::Timestamp(v), DataType::Timestamp(TimeUnit::Millisecond, _)) => {
+            append_scalar!(TimestampSecondBuilder, col, v.time as i64 * 1000)
+        }
         (RawBsonRef::Timestamp(v), DataType::Timestamp(TimeUnit::Microsecond, _)) => {
-            append_scalar!(TimestampMicrosecondBuilder, col, v.time as i64) // TODO: Possibly change to nanosecond.
+            append_scalar!(TimestampSecondBuilder, col, v.time as i64 * 1000 * 1000)
+        }
+        (RawBsonRef::Timestamp(v), DataType::Date64) => {
+            append_scalar!(Date64Builder, col, v.time as i64 * 1000)
+        }
+        (RawBsonRef::Timestamp(v), DataType::Date32) => {
+            append_scalar!(
+                Date32Builder,
+                col,
+                v.time
+                    .try_into()
+                    .map_err(|_| BsonError::UnhandledElementType(
+                        bson::spec::ElementType::Timestamp,
+                        DataType::Date32
+                    ))?
+            )
         }
 
-        // Datetime
+        // Datetime (actual timestamps that you'd actually use. in an application )
+        (RawBsonRef::DateTime(v), DataType::Timestamp(TimeUnit::Millisecond, _)) => {
+            append_scalar!(TimestampMillisecondBuilder, col, v.timestamp_millis())
+        }
         (RawBsonRef::DateTime(v), DataType::Timestamp(TimeUnit::Microsecond, _)) => {
             append_scalar!(
-                TimestampMicrosecondBuilder, // TODO: Possibly change to nanosecond.
+                TimestampMicrosecondBuilder,
                 col,
-                v.timestamp_millis()
+                v.timestamp_millis() * 1000
             )
         }
+        (RawBsonRef::DateTime(v), DataType::Date64) => {
+            append_scalar!(Date64Builder, col, v.timestamp_millis())
+        }
 
         // Document
         (RawBsonRef::Document(nested), DataType::Struct(_)) => {
@@ -263,13 +293,16 @@ fn append_value(val: RawBsonRef, typ: &DataType, col: &mut dyn ArrayBuilder) ->
 
         // Array
         (RawBsonRef::Array(arr), DataType::Utf8) => {
-            // TODO: Proper types.
-            let s = arr
-                .into_iter()
-                .map(|r| r.map(|v| format!("{:?}", v)).unwrap_or_default())
-                .collect::<Vec<_>>()
-                .join(", ");
-            append_scalar!(StringBuilder, col, format!("[{}]", s))
+            append_scalar!(
+                StringBuilder,
+                col,
+                serde_json::Value::try_from(
+                    bson::Array::try_from(arr)
+                        .map_err(|_| BsonError::FailedToReadRawBsonDocument)?
+                )
+                .map_err(|_| BsonError::FailedToReadRawBsonDocument)?
+                .to_string()
+            )
         }
 
         // Decimal128

crates/datasources/src/bson/errors.rs

@@ -1,16 +1,34 @@
 use datafusion::error::DataFusionError;
+use datafusion_ext::errors::ExtensionError;
+
+use crate::object_store::errors::ObjectStoreSourceError;
 
 #[derive(Debug, thiserror::Error)]
 pub enum BsonError {
     #[error("Unsupported bson type: {0}")]
-    UnsupportedBsonType(&'static str),
+    UnspportedType(&'static str),
 
     #[error("Unexpected datatype for builder {0:?}")]
     UnexpectedDataTypeForBuilder(datafusion::arrow::datatypes::DataType),
 
-    #[error("External Datafusion Error")]
+    #[error(transparent)]
+    IO(#[from] std::io::Error),
+
+    #[error(transparent)]
     Datafusion(#[from] datafusion::error::DataFusionError),
 
+    #[error(transparent)]
+    Arrow(#[from] datafusion::arrow::error::ArrowError),
+
+    #[error(transparent)]
+    Raw(#[from] bson::raw::Error),
+
+    #[error(transparent)]
+    Serialization(#[from] bson::de::Error),
+
+    #[error(transparent)]
+    ObjectStore(#[from] ObjectStoreSourceError),
+
     #[error("Unhandled element type to arrow type conversion; {0:?}, {1}")]
     UnhandledElementType(
         bson::spec::ElementType,
@@ -28,6 +46,9 @@ pub enum BsonError {
 
     #[error("Recursion limit exceeded for schema inferrence: {0}")]
     RecursionLimitExceeded(usize),
+
+    #[error("no objects found {0}")]
+    NotFound(String),
 }
 
 impl From<BsonError> for DataFusionError {
@@ -36,4 +57,10 @@ impl From<BsonError> for DataFusionError {
     }
 }
 
+impl From<BsonError> for ExtensionError {
+    fn from(e: BsonError) -> Self {
+        ExtensionError::String(e.to_string())
+    }
+}
+
 pub type Result<T, E = BsonError> = std::result::Result<T, E>;

crates/datasources/src/bson/schema.rs

@@ -2,7 +2,7 @@ use std::collections::HashMap;
 use std::iter::IntoIterator;
 
 use bson::{Bson, Document};
-use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
 
 use crate::bson::errors::{BsonError, Result};
 
@@ -11,9 +11,8 @@ use crate::bson::errors::{BsonError, Result};
 /// The MongoDB kernel rejects nesting of greater than 100.
 const RECURSION_LIMIT: usize = 100;
 
-pub fn schema_from_document(doc: &Document) -> Result<Schema> {
-    let fields = fields_from_document(0, doc.iter())?;
-    Ok(Schema::new(fields))
+pub fn schema_from_document(doc: Document) -> Result<Schema> {
+    Ok(Schema::new(fields_from_document(0, doc.iter())?))
 }
 
 fn fields_from_document<'a>(
@@ -39,61 +38,51 @@ fn fields_from_document<'a>(
 }
 
 fn bson_to_arrow_type(depth: usize, bson: &Bson) -> Result<DataType> {
-    let arrow_typ = match bson {
-        Bson::Double(_) => DataType::Float64,
-        Bson::String(val) => {
-            if val.is_empty() {
-                // TODO: We'll want to determine if this is something we should
-                // keep in. Currently when we load in the test file, we're
-                // loading a csv that might have empty values (null). Mongo
-                // reads those in and puts them as empty strings.
-                //
-                // During schema merges, we'll widen "null" types to whatever
-                // type we're merging with. This _should_ be more resilient to
-                // empty values.
-                DataType::Null
-            } else {
-                DataType::Utf8
-            }
-        }
-        Bson::Array(array_doc) => match array_doc.to_owned().pop() {
-            Some(val) => bson_to_arrow_type(0, &val)?,
-            None => DataType::Utf8,
-        },
+    Ok(match bson {
+        Bson::Array(array_doc) => DataType::new_list(
+            bson_to_arrow_type(0, &array_doc.to_owned().pop().unwrap_or(Bson::Null))?,
+            true,
+        ),
         Bson::Document(nested) => {
-            let fields = fields_from_document(depth + 1, nested.iter())?;
-            DataType::Struct(fields.into())
+            DataType::Struct(fields_from_document(depth + 1, nested.iter())?.into())
         }
+        Bson::String(_) => DataType::Utf8,
+        Bson::Double(_) => DataType::Float64,
         Bson::Boolean(_) => DataType::Boolean,
         Bson::Null => DataType::Null,
-        Bson::RegularExpression(_) => DataType::Utf8,
-        Bson::JavaScriptCode(_) => DataType::Utf8,
-        Bson::JavaScriptCodeWithScope(_) => {
-            return Err(BsonError::UnsupportedBsonType("CodeWithScope"))
-        }
         Bson::Int32(_) => DataType::Float64,
         Bson::Int64(_) => DataType::Float64,
-        Bson::Timestamp(_) => return Err(BsonError::UnsupportedBsonType("OplogTimestamp")),
         Bson::Binary(_) => DataType::Binary,
         Bson::ObjectId(_) => DataType::Utf8,
-        Bson::DateTime(_) => DataType::Timestamp(TimeUnit::Millisecond, None),
+        Bson::DateTime(_) => DataType::Date64,
         Bson::Symbol(_) => DataType::Utf8,
         Bson::Decimal128(_) => DataType::Decimal128(38, 10),
         Bson::Undefined => DataType::Null,
+        Bson::RegularExpression(_) => DataType::Utf8,
+        Bson::JavaScriptCode(_) => DataType::Utf8,
+
+        // TODO: storing these (which exist to establish a total order
+        // of types for indexing in the MongoDB server,) in documents
+        // that GlareDB would interact with is probably always an
+        // error.
         Bson::MaxKey => DataType::Utf8,
         Bson::MinKey => DataType::Utf8,
-        Bson::DbPointer(_) => return Err(BsonError::UnsupportedBsonType("DbPointer")),
-    };
-    Ok(arrow_typ)
+
+        // Deprecated or MongoDB server intrenal types
+        Bson::JavaScriptCodeWithScope(_) => return Err(BsonError::UnspportedType("CodeWithScope")),
+        Bson::Timestamp(_) => return Err(BsonError::UnspportedType("OplogTimestamp")),
+        Bson::DbPointer(_) => return Err(BsonError::UnspportedType("DbPointer")),
+    })
 }
 
 #[derive(Debug, Clone)]
 struct OrderedField(usize, Field);
 
-pub fn merge_schemas(schemas: impl IntoIterator<Item = Schema>) -> Result<Schema> {
+pub fn merge_schemas(schemas: impl IntoIterator<Item = Result<Schema>>) -> Result<Schema> {
     let mut fields: HashMap<String, OrderedField> = HashMap::new();
 
     for schema in schemas.into_iter() {
+        let schema = schema?;
         for (idx, field) in schema.fields.into_iter().enumerate() {
             match fields.get_mut(field.name()) {
                 Some(existing) => {