Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
202 changes: 170 additions & 32 deletions Cargo.lock

Large diffs are not rendered by default.

10 changes: 8 additions & 2 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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",
] }
Expand Down Expand Up @@ -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" }
Expand Down
10 changes: 5 additions & 5 deletions schemas/otel_logs_and_spans.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
105 changes: 97 additions & 8 deletions src/database.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -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(())
}

Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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);
Expand Down
Loading
Loading