Merge pull request #15 from monoscope-tech/variant-support2

Add Variant type support and fix SLT tests

Author: tonyalaribe

Cargo.lock

@@ -20,8 +20,8 @@ log = "0.4.27"
 color-eyre = "0.6.5"
 arrow-schema = "57.1.0"
 regex = "1.11.1"
-# Updated to delta-rs with datafusion 52 Utf8View fixes (includes commits 987e535f, ffb794ba)
-deltalake = { git = "https://github.com/delta-io/delta-rs.git", rev = "ffb794ba0745394fc4b747a4ef2e11c2d4ec086a", features = [
+# Using fork with VariantType support until upstream merges the feature
+deltalake = { git = "https://github.com/tonyalaribe/delta-rs.git", rev = "ba769136c5dd9b84a7335ea67e42b67884bfcce3", features = [
   "datafusion",
   "s3",
 ] }
@@ -75,6 +75,12 @@ bincode = { version = "2.0", features = ["serde"] }
 walrus-rust = "0.2.0"
 thiserror = "2.0"
 strum = { version = "0.27", features = ["derive"] }
+datafusion-variant = { git = "https://github.com/tonyalaribe/datafusion-variant.git", rev = "8b6b270" }
+parquet-variant-compute = "57.2.0"
+parquet-variant-json = "57.2.0"
+parquet-variant = "57.2.0"
+serde_json_path = "0.7"
+base64 = "0.22"
 
 [dev-dependencies]
 sqllogictest = { git = "https://github.com/risinglightdb/sqllogictest-rs.git" }

Cargo.toml

@@ -61,7 +61,7 @@ fields:
     data_type: 'Timestamp(Microsecond, Some("UTC"))'
     nullable: true
   - name: context
-    data_type: Utf8
+    data_type: Variant
     nullable: true
   - name: context___trace_id
     data_type: Utf8
@@ -79,13 +79,13 @@ fields:
     data_type: Utf8
     nullable: true
   - name: events
-    data_type: Utf8
+    data_type: Variant
     nullable: true
   - name: links
-    data_type: Utf8
+    data_type: Variant
     nullable: true
   - name: attributes
-    data_type: Utf8
+    data_type: Variant
     nullable: true
   - name: attributes___client___address
     data_type: Utf8
@@ -235,7 +235,7 @@ fields:
     data_type: Utf8
     nullable: true
   - name: resource
-    data_type: Utf8
+    data_type: Variant
     nullable: true
   - name: resource___service___name
     data_type: Utf8

schemas/otel_logs_and_spans.yaml

@@ -1,9 +1,9 @@
 use crate::config::{self, AppConfig};
 use crate::object_store_cache::{FoyerCacheConfig, FoyerObjectStoreCache, SharedFoyerCache};
-use crate::schema_loader::{get_default_schema, get_schema};
+use crate::schema_loader::{get_default_schema, get_schema, is_variant_type};
 use crate::statistics::DeltaStatisticsExtractor;
 use anyhow::Result;
-use arrow_schema::SchemaRef;
+use arrow_schema::{Schema, SchemaRef};
 use async_trait::async_trait;
 use chrono::Utc;
 use datafusion::arrow::array::Array;
@@ -36,10 +36,10 @@ use deltalake::operations::create::CreateBuilder;
 use deltalake::{DeltaTable, DeltaTableBuilder};
 use futures::StreamExt;
 use instrumented_object_store::instrument_object_store;
-use std::sync::Mutex;
 use serde::{Deserialize, Serialize};
 use sqlx::{PgPool, postgres::PgPoolOptions};
 use std::fmt;
+use std::sync::Mutex;
 use std::sync::OnceLock;
 use std::{any::Any, collections::HashMap, sync::Arc};
 use tokio::sync::RwLock;
@@ -82,6 +82,85 @@ pub fn extract_project_id(batch: &RecordBatch) -> Option<String> {
     })
 }
 
+/// Convert string columns to Variant binary format where the target schema expects Variant type.
+/// This enables automatic JSON string → Variant conversion during INSERT.
+pub fn convert_variant_columns(batch: RecordBatch, target_schema: &SchemaRef) -> DFResult<RecordBatch> {
+    use datafusion::arrow::array::{ArrayRef, LargeStringArray, StringArray, StringViewArray};
+    use datafusion::arrow::datatypes::{DataType, Field};
+
+    let batch_schema = batch.schema();
+    let mut columns: Vec<ArrayRef> = batch.columns().to_vec();
+    let mut new_fields: Vec<Arc<Field>> = batch_schema.fields().iter().cloned().collect();
+
+    for (idx, target_field) in target_schema.fields().iter().enumerate() {
+        if !is_variant_type(target_field.data_type()) {
+            continue;
+        }
+        // Skip columns beyond batch length - this is normal for INSERT with fewer columns than table schema
+        // (e.g., columns with defaults or nullable columns omitted from INSERT)
+        if idx >= columns.len() {
+            continue;
+        }
+
+        let col = &columns[idx];
+        let col_type = col.data_type();
+
+        // Only convert if source is a string type and target is Variant
+        let converted: Option<ArrayRef> =
+            match col_type {
+                DataType::Utf8View => {
+                    let arr = col.as_any().downcast_ref::<StringViewArray>().ok_or_else(|| {
+                        DataFusionError::Execution(format!("Expected StringViewArray for field '{}' but downcast failed", target_field.name()))
+                    })?;
+                    Some(Arc::new(json_strings_to_variant(arr.iter())?))
+                }
+                DataType::Utf8 => {
+                    let arr = col
+                        .as_any()
+                        .downcast_ref::<StringArray>()
+                        .ok_or_else(|| DataFusionError::Execution(format!("Expected StringArray for field '{}' but downcast failed", target_field.name())))?;
+                    Some(Arc::new(json_strings_to_variant(arr.iter())?))
+                }
+                DataType::LargeUtf8 => {
+                    let arr = col.as_any().downcast_ref::<LargeStringArray>().ok_or_else(|| {
+                        DataFusionError::Execution(format!("Expected LargeStringArray for field '{}' but downcast failed", target_field.name()))
+                    })?;
+                    Some(Arc::new(json_strings_to_variant(arr.iter())?))
+                }
+                _ => None, // Already Variant or other type, skip
+            };
+
+        if let Some(variant_array) = converted {
+            columns[idx] = variant_array;
+            new_fields[idx] = target_field.clone();
+        }
+    }
+
+    let new_schema = Arc::new(Schema::new(new_fields));
+    RecordBatch::try_new(new_schema, columns).map_err(|e| DataFusionError::ArrowError(Box::new(e), None))
+}
+
+/// Convert an iterator of optional JSON strings to a Variant StructArray.
+/// Fails fast on invalid JSON to ensure data integrity.
+fn json_strings_to_variant<'a>(iter: impl Iterator<Item = Option<&'a str>>) -> DFResult<datafusion::arrow::array::StructArray> {
+    use parquet_variant_compute::VariantArrayBuilder;
+    use parquet_variant_json::JsonToVariant;
+
+    let items: Vec<_> = iter.collect();
+    let mut builder = VariantArrayBuilder::new(items.len());
+
+    for (row_idx, item) in items.into_iter().enumerate() {
+        match item {
+            Some(json_str) => builder
+                .append_json(json_str)
+                .map_err(|e| DataFusionError::Execution(format!("Invalid JSON at row {}: {} (value: '{}')", row_idx, e, json_str)))?,
+            None => builder.append_null(),
+        }
+    }
+
+    Ok(builder.build().into())
+}
+
 // Compression level for parquet files - kept for WriterProperties fallback
 const ZSTD_COMPRESSION_LEVEL: i32 = 3;
 
@@ -712,11 +791,14 @@ impl Database {
 
         self.register_pg_settings_table(ctx)?;
         self.register_set_config_udf(ctx);
-        self.register_json_functions(ctx);
 
-        // Register custom PostgreSQL-compatible functions
+        // CRITICAL: Register custom functions BEFORE JSON functions to ensure VariantAwareExprPlanner
+        // intercepts -> and ->> operators on Variant columns before JsonExprPlanner handles them as strings
         crate::functions::register_custom_functions(ctx).map_err(|e| DataFusionError::Execution(format!("Failed to register custom functions: {}", e)))?;
 
+        // JSON functions (JsonExprPlanner for -> and ->> on string columns - must come after Variant handlers)
+        self.register_json_functions(ctx);
+
         Ok(())
     }
 
@@ -1205,7 +1287,10 @@ impl Database {
 
         // Fallback to legacy batch queue if configured
         let enable_queue = self.config.core.enable_batch_queue;
-        if !skip_queue && enable_queue && let Some(ref queue) = self.batch_queue {
+        if !skip_queue
+            && enable_queue
+            && let Some(ref queue) = self.batch_queue
+        {
             span.record("use_queue", true);
             for batch in batches {
                 if let Err(e) = queue.queue(batch) {
@@ -1892,14 +1977,18 @@ impl DataSink for ProjectRoutingTable {
         let span = tracing::Span::current();
         let mut total_row_count = 0;
         let mut project_batches: HashMap<String, Vec<RecordBatch>> = HashMap::new();
+        let target_schema = self.schema();
 
-        // Collect and group batches by project_id
+        // Collect and group batches by project_id, converting variant columns
         while let Some(batch) = data.next().await.transpose()? {
             let batch_rows = batch.num_rows();
             debug!("write_all: received batch with {} rows", batch_rows);
             total_row_count += batch_rows;
             let project_id = extract_project_id(&batch).unwrap_or_else(|| self.default_project.clone());
-            project_batches.entry(project_id).or_default().push(batch);
+
+            // Convert string columns to Variant where target schema expects Variant
+            let converted_batch = convert_variant_columns(batch, &target_schema)?;
+            project_batches.entry(project_id).or_default().push(converted_batch);
         }
 
         span.record("rows.count", total_row_count);